from __future__ import annotations

import string
import sys
from enum import IntEnum
from random import choices
from typing import List, Optional, Tuple, Union, cast

from PySide2.QtCore import (
    QAbstractItemModel,
    QModelIndex,
    QObject,
    QSortFilterProxyModel,
    Qt,
    Signal,
)
from PySide2.QtWidgets import QApplication, QTreeView


class Node:
    """Node inside the tree model."""

    def __init__(self, name: Optional[str]) -> None:
        """Create a new Node."""
        self._name = name
        self._children: List[Node] = []
        self._parent: Optional[Node] = None
        self._check_state: Qt.CheckState = Qt.CheckState.Unchecked

    @property
    def name(self) -> Optional[str]:
        """Return the name of the Node."""
        return self._name

    @property
    def child_count(self) -> int:
        """Return the number of children for the Node."""
        return len(self._children)

    def child(self, row: int) -> Node:
        """Return the child for the given row."""
        return self._children[row]

    @property
    def children(self) -> List[Node]:
        """Return the list of children."""
        return self._children

    @property
    def parent(self) -> Optional[Node]:
        """Return the Node's parent element."""
        return self._parent

    @parent.setter
    def parent(self, parent: Node) -> None:
        """Set the Node's parent element."""
        self._parent = parent

    @property
    def row(self) -> int:
        """Return the row number of the Node."""
        if not self.parent:
            return -1
        return self.parent.children.index(self)

    def add_child(self, child: Node) -> None:
        """Add the given child to the Node's children."""
        child.parent = self
        self._children.append(child)

    def remove_child(self, child: Node) -> None:
        """Remove the given child to the Node's children."""
        self._children.remove(child)
        child.parent = None

    @property
    def check_state(self) -> Qt.CheckState:
        """Return the check state of the Node."""
        return self._check_state

    @check_state.setter
    def check_state(self, check_state: Qt.CheckState) -> None:
        """Set the check state of the Node."""
        self._check_state = check_state


class Model(QAbstractItemModel):
    """Model for data in a QTreeView."""

    add_filter = Signal(tuple, name="add_filter")
    remove_filter = Signal(tuple, name="remove_filter")

    class Header(IntEnum):
        """Header definitions."""

        NAME = 0

    def __init__(self, parent: Optional[QObject] = None) -> None:
        """Create a new Model to show data in a QTreeView."""
        super().__init__(parent)
        self._header_labels = {
            Model.Header.NAME: self.tr("Name"),
        }
        self._root = Node(None)
        self._path_cache: List[Tuple[str, ...]] = []
        self._blocked = False

    def flags(self, index: QModelIndex) -> Qt.ItemFlags:
        """Return the flags for the given index and model."""
        if index.isValid():
            return super().flags(index) | Qt.ItemFlag.ItemIsUserCheckable
        return super().flags(index)

    def columnCount(  # pylint: disable=invalid-name, no-self-use
        self, _: QModelIndex = QModelIndex()
    ) -> int:
        """Return the column count."""
        return len(self._header_labels)

    def rowCount(  # pylint: disable=invalid-name
        self, parent: QModelIndex = QModelIndex()
    ) -> int:
        """Return the row count."""
        if parent.isValid():
            return cast(Node, parent.internalPointer()).child_count
        count = self._root.child_count
        return count

    def headerData(  # pylint: disable=invalid-name
        self,
        section: int,
        orientation: Qt.Orientation,
        role: int = Qt.ItemDataRole.DisplayRole,
    ) -> Optional[str]:
        """Return the header data for the given section."""
        if (
            orientation == Qt.Orientation.Horizontal
            and role == Qt.ItemDataRole.DisplayRole
        ):
            return self._header_labels[Model.Header(section)]
        return None

    def on_add_path(self, path: Tuple[str, ...]) -> None:
        """Add a path to the model if it doesn't exist yet."""
        if path not in self._path_cache:
            self._add_path_to_node(path, QModelIndex())
            self._path_cache.append(path)

    def _add_path_to_node(
        self, path: Tuple[str, ...], index: QModelIndex
    ) -> None:
        """Add the path to the given node if it doesn't exist yet."""
        if index.isValid():
            node: Node = index.internalPointer()
        else:
            node = self._root
        for child in node.children:
            if child.name == path[0]:
                self._add_path_to_node(
                    path[1:], self.index(child.row, 0, index)
                )
                return
        new_node = Node(path[0])
        self.beginInsertRows(index, node.child_count, node.child_count)
        node.add_child(new_node)
        self.endInsertRows()
        if path[1:]:
            self._add_path_to_node(
                path[1:], self.index(new_node.row, 0, index)
            )

    def parent(self, child: QModelIndex) -> QModelIndex:  # type: ignore
        """Return the parent index for a given index."""
        if not child.isValid():
            return QModelIndex()

        child_item = child.internalPointer()
        parent_item = child_item.parent

        if parent_item is self._root:
            return QModelIndex()

        return self.createIndex(parent_item.row, 0, parent_item)

    def hasChildren(  # pylint: disable=invalid-name
        self, index: QModelIndex = QModelIndex()
    ) -> bool:
        """Evaluate if children exist for the given index."""
        if index.isValid():
            return bool(index.internalPointer().child_count)
        return bool(self._root.child_count)

    def index(
        self, row: int, col: int, parent: QModelIndex = QModelIndex()
    ) -> QModelIndex:
        """Return an index for the given row and column and parent."""
        # a not existent index should never be requested
        assert self.hasIndex(row, col, parent)

        if not parent.isValid():
            parent_item = self._root
        else:
            parent_item = parent.internalPointer()

        return self.createIndex(row, col, parent_item.child(row))

    def data(  # pylint: disable=no-self-use
        self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole
    ) -> Union[str, Qt.CheckState, None]:
        """Return data for the given index and role."""
        if role == Qt.ItemDataRole.DisplayRole:
            item = cast(Node, index.internalPointer())
            return item.name
        if role == Qt.ItemDataRole.CheckStateRole:
            item = cast(Node, index.internalPointer())
            return item.check_state
        return None

    def setData(  # pylint: disable=invalid-name
            self,
            index: QModelIndex,
            value: Qt.CheckState,
            role: int = Qt.ItemDataRole.EditRole,
    ) -> bool:
        """Set the data for the given index and role."""
        if role == Qt.ItemDataRole.CheckStateRole:
            self.layoutAboutToBeChanged.emit()
            item = cast(Node, index.internalPointer())
            item.check_state = Qt.CheckState(value)
            self.dataChanged.emit(
                index, index, [Qt.ItemDataRole.CheckStateRole]
            )
            self.layoutChanged.emit()
            return True
        return False


if __name__ == "__main__":
    APP = QApplication(sys.argv)
    model = Model()
    rand_words = [
        "".join(choices(string.ascii_uppercase + string.digits, k=6))
        for _ in range(20)
    ]
    for char in "ABCDEFG":
        for number in range(100000, 100010):
            for word in rand_words:
                model.on_add_path((char, str(number), word))
    filter_model = QSortFilterProxyModel()
    filter_model.setSourceModel(model)
    view = QTreeView()
    view.setModel(filter_model)
    view.show()
    sys.exit(APP.exec_())