from __future__ import annotations

from dataclasses import dataclass

from PySide6.QtCore import QAbstractItemModel, QModelIndex, Qt, Signal
from PySide6.QtGui import QIcon

from bedit_core.models import Component, ComponentID, GraphImplementation
from bedit_core.models import Document as CoreDocument


@dataclass
class DocumentTreeNode:
    name: str
    value: object
    component_id: ComponentID | None
    parent: DocumentTreeNode | None
    children: list[DocumentTreeNode]


class DocumentTreeModel(QAbstractItemModel):
    rename_document_requested = Signal(str)
    rename_component_requested = Signal(Component, str)

    def __init__(self, editable: bool = True) -> None:
        super().__init__()
        self._editable = editable
        self._root = DocumentTreeNode("Document", None, None, None, [])
        self._component_icons: dict[ComponentID, QIcon] = {}
        self._component_nodes: dict[ComponentID, DocumentTreeNode] = {}

    def set_document(self, document: CoreDocument) -> None:
        self.set_documents([document])

    def set_documents(self, documents: list[CoreDocument]) -> None:
        self.beginResetModel()
        self._component_icons = {}
        self._component_nodes = {}
        self._root = self._build_tree(documents)
        self.endResetModel()

    def set_component_icon(self, component_id: ComponentID, icon: QIcon) -> None:
        node = self._component_nodes.get(component_id)
        if node is None or node.parent is None:
            return
        self._component_icons[component_id] = icon
        row = node.parent.children.index(node)
        index = self.createIndex(row, 1, node)
        self.dataChanged.emit(index, index, [Qt.ItemDataRole.DecorationRole])

    def rowCount(self, parent: QModelIndex | None = None) -> int:
        if parent is not None and parent.isValid() and parent.column() != 0:
            return 0
        return len(self._node(parent).children)

    def columnCount(self, _parent: QModelIndex | None = None) -> int:
        return 2

    def index(self, row: int, column: int, parent: QModelIndex | None = None) -> QModelIndex:
        if parent is not None and parent.isValid() and parent.column() != 0:
            return QModelIndex()
        parent_node = self._node(parent)
        if column not in (0,1) or row < 0 or row >= len(parent_node.children):
            return QModelIndex()
        return self.createIndex(row, column, parent_node.children[row])

    def parent(self, index: QModelIndex) -> QModelIndex:
        if not index.isValid():
            return QModelIndex()

        node = index.internalPointer()
        if not isinstance(node, DocumentTreeNode):
            return QModelIndex()

        parent_node = node.parent
        if parent_node is None or parent_node is self._root:
            return QModelIndex()

        grandparent = parent_node.parent
        if grandparent is None:
            return QModelIndex()
        row = grandparent.children.index(parent_node)
        return self.createIndex(row, 0, parent_node)

    def data(self, index: QModelIndex, role: int = Qt.ItemDataRole.DisplayRole) -> object | None:
        if not index.isValid():
            return None

        node = index.internalPointer()
        if not isinstance(node, DocumentTreeNode):
            return None

        if index.column() == 0:
            if role in (Qt.ItemDataRole.DisplayRole, Qt.ItemDataRole.EditRole):
                return node.name
        elif index.column() == 1:
            if role == Qt.ItemDataRole.DecorationRole:
                return self._component_icons.get(node.component_id)
            if role == Qt.ItemDataRole.TextAlignmentRole:
                return Qt.AlignmentFlag.AlignCenter

        return None

    def setData(self, index: QModelIndex, value: object, role: int = Qt.ItemDataRole.EditRole) -> bool:
        if role != Qt.ItemDataRole.EditRole or not index.isValid() or index.column() != 0:
            return False

        node = index.internalPointer()
        if not isinstance(node, DocumentTreeNode) or not isinstance(node.value, (CoreDocument, Component)):
            return False

        name = str(value).strip()

        if not name or name == node.name:
            return False

        if isinstance(node.value, CoreDocument):
            self.rename_document_requested.emit(name)
        elif isinstance(node.value, Component):
            self.rename_component_requested.emit(node.value, name)
        return True

    def _node(self, index: QModelIndex | None) -> DocumentTreeNode:
        if index is None or not index.isValid():
            return self._root

        node = index.internalPointer()
        return node if isinstance(node, DocumentTreeNode) else self._root

    def value(self, index: QModelIndex) -> object | None:
        if not index.isValid():
            return None
        node = index.internalPointer()
        return node.value if isinstance(node, DocumentTreeNode) else None

    def component_index(self, component_id: ComponentID) -> QModelIndex:
        node = self._component_nodes.get(component_id)
        if node is None or node.parent is None:
            return QModelIndex()
        return self.createIndex(node.parent.children.index(node), 0, node)

    def flags(self, index: QModelIndex) -> Qt.ItemFlag:
        flags = super().flags(index)

        if not index.isValid():
            return flags

        node = index.internalPointer()

        # Make the document root node editable
        if self._editable and index.column() == 0 and isinstance(node, DocumentTreeNode) and isinstance(node.value, (CoreDocument, Component)):
            flags |= Qt.ItemFlag.ItemIsEditable

        return flags

    def _build_tree(self, documents: list[CoreDocument]) -> DocumentTreeNode:
        # QT's invisible root
        root = DocumentTreeNode(
            name="",
            value=None,
            component_id=None,
            parent=None,
            children=[],
        )
        def _list_children(parent: DocumentTreeNode, components: dict[ComponentID, Component]) -> None:
            for component_id, component in sorted(components.items(), key=lambda item: (item[1].name.casefold(), item[1].name, str(item[0]))):
                component_node = DocumentTreeNode(name=component.name, value=component, component_id=component_id, parent=parent, children=[])
                parent.children.append(component_node)
                self._component_nodes[component_id] = component_node

                if isinstance(component.implementation, GraphImplementation):
                    _list_children(component_node, component.implementation.graph.components)

        for document in documents:
            document_root = DocumentTreeNode(name=document.name, value=document, component_id=None, parent=root, children=[])
            root.children.append(document_root)
            _list_children(document_root, document.root)

        return root