from __future__ import annotations

from dataclasses import dataclass

from PySide6.QtCore import QAbstractItemModel, QModelIndex, Qt

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


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


class DocumentTreeModel(QAbstractItemModel):
    def __init__(self) -> None:
        super().__init__()
        self._document: CoreDocument | None = None
        self._root = DocumentTreeNode("Document", None, None, [])

    def set_document(self, document: CoreDocument) -> None:
        self.beginResetModel()
        self._document = document
        self._root = self._build_tree(document)
        self.endResetModel()

    def rowCount(self, parent: QModelIndex | None = None) -> int:
        return len(self._node(parent).children)

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

    def index(
        self,
        row: int,
        column: int,
        parent: QModelIndex | None = None,
    ) -> QModelIndex:
        parent_node = self._node(parent)
        if (
            column != 0
            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 (
            role == Qt.ItemDataRole.DisplayRole
            and isinstance(node, DocumentTreeNode)
        ):
            return node.name
        return None

    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 _build_tree(self, document: CoreDocument) -> DocumentTreeNode:
        # QT's invisible root
        root = DocumentTreeNode(
            name="",
            value=None,
            parent=None,
            children=[],
        )
        # Add itself as a child so the document root is visible in the tree
        document_root = DocumentTreeNode(
            name = document.name,
            value = document,
            parent = root,
            children = []
        )
        root.children.append(document_root)

        def _list_children(root: DocumentTreeNode, components: dict[ComponentID, Component]) -> None:
            for component_id, component in components.items():
                component_node = DocumentTreeNode(
                    name = component.name,
                    value = (component_id, component),
                    parent = root,
                    children = []
                )
                root.children.append(component_node)

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

        return root