tree_model

class BulkTreeModel(QAbstractItemModel):
    """Qt tree model populated in one shot from a single background fetch.

    Unlike :class:`LazyTreeModel`, nothing is fetched per node on
    expand: the whole hierarchy is fetched in a single ``fetch_all``
    call - run once on a background thread via the shared
    :class:`AsyncTaskQueue` - and the entire tree is built from that
    one response. Use this instead of :class:`LazyTreeModel` when a
    fast bulk endpoint can return the whole hierarchy at once, so a
    tree doesn't need to pay for one round trip per expanded level.

    Args:
        fetch_all: Callable, taking no arguments, that returns the
            whole hierarchy as a ``{parent_id: children}`` mapping
            (``None`` for the root's own children).
        no_async: When ``True``, the fetch runs synchronously on the
            calling thread instead of via the :class:`AsyncTaskQueue`.
            Useful in tests to avoid worker-thread/paint-event races.
        parent: Optional parent QObject.

    Example::

        def fetch_all():
            return {
                None: [TreeNode("root", "Root", has_children=True)],
                "root": [TreeNode("child", "Child")],
            }

        model = BulkTreeModel(fetch_all=fetch_all)
    """

    #: Role exposing TreeNode.filter_text (falls back to the label) for
    #: search proxies to match against, independent of DisplayRole.
    FILTER_ROLE = Qt.ItemDataRole.UserRole + 1

    loading_changed = Signal(bool)  # True while the fetch is in-flight
    fetch_error = Signal(str)  # error message when the fetch fails

    def __init__(
        self,
        fetch_all: Callable[[], dict[str | None, list[TreeNode]]],
        no_async: bool = False,
        parent: QObject | None = None,
    ) -> None:
        super().__init__(parent)
        self._fetch_all = fetch_all
        self._no_async: bool = no_async
        self._reset_counter: int = 0
        self._context_id: str = f"btm_{id(self)}_v0"
        self._is_loading: bool = False
        self._root = _InternalNode(tree_node=None, parent=None)
        # _all_nodes keeps every node reachable so Python's GC does not
        # collect objects that are held only via QModelIndex.internalPointer().
        self._all_nodes: set[_InternalNode] = {self._root}

    def _index_for_node(self, node: _InternalNode) -> QModelIndex:
        """Return the QModelIndex that identifies node (column 0)."""
        if node.is_root:
            return QModelIndex()
        parent = node.parent
        row = parent.children.index(node)  # type: ignore[union-attr]
        return self.createIndex(row, 0, node)

    def _fetch_all_async(self) -> None:
        """Kick off the single bulk fetch on the shared task queue."""
        ctx = self._context_id
        self._is_loading = True
        self.loading_changed.emit(True)

        if self._no_async:
            try:
                result = self._fetch_all()
            except Exception:
                log.exception("fetch_all raised")
                result = None
            self._on_data_ready(ctx, result)
            return

        task = AsyncTask(
            name="fetch_all_tree",
            function=self._fetch_all,
            callback=lambda result: self._on_data_ready(ctx, result),
            priority=1,
            context_id=ctx,
            cancellable=True,
        )
        get_task_queue().enqueue(task)

    def _on_data_ready(
        self,
        context_id: str,
        result: dict[str | None, list[TreeNode]] | None,
    ) -> None:
        """Handle the bulk fetch result on the main thread.

        Called via Qt.QueuedConnection from AsyncTaskQueue so it always
        executes on the Qt main thread, making it safe to call Qt model
        mutation methods.

        Args:
            context_id: The context_id active when the task was
                enqueued. Used to discard results from a pre-reset fetch.
            result: ``{parent_id: children}`` mapping for the whole
                hierarchy, or ``None`` on error.
        """
        if context_id != self._context_id:
            return
        self._is_loading = False

        if result is None:
            self.fetch_error.emit("Failed to fetch tree data")
            self.loading_changed.emit(False)
            return

        self.beginResetModel()
        new_root = _InternalNode(tree_node=None, parent=None)
        all_nodes: set[_InternalNode] = {new_root}
        # Iterative (rather than recursive) so hierarchy depth can
        # never hit Python's recursion limit.
        stack = [new_root]
        while stack:
            node = stack.pop()
            parent_id = None if node.is_root else node.tree_node.id  # type: ignore[union-attr]
            children = result.get(parent_id) or []
            new_nodes = [
                _InternalNode(tree_node=tn, parent=node) for tn in children
            ]
            node.children = new_nodes
            all_nodes.update(new_nodes)
            stack.extend(new_nodes)
        self._root = new_root
        self._all_nodes = all_nodes
        self.endResetModel()

        self.loading_changed.emit(False)

    def _node_from_index(
        self, index: QModelIndex | QPersistentModelIndex
    ) -> _InternalNode:
        """Return the internal node for a model index."""
        if not index.isValid():
            return self._root
        return index.internalPointer()  # type: ignore[return-value]

    def index(
        self,
        row: int,
        column: int,
        parent: QModelIndex | QPersistentModelIndex = QModelIndex(),
    ) -> QModelIndex:
        parent_node = self._node_from_index(parent)
        if row < 0 or row >= len(parent_node.children):
            return QModelIndex()
        child_node = parent_node.children[row]
        return self.createIndex(row, column, child_node)

    def parent(self, index: QModelIndex) -> QModelIndex:  # type: ignore[override]
        if not index.isValid():
            return QModelIndex()

        node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
        parent_node = node.parent
        if parent_node is None or parent_node.is_root:
            return QModelIndex()

        # grandparent is always non-None for a non-root parent_node
        grandparent = parent_node.parent
        row = grandparent.children.index(parent_node)  # type: ignore[union-attr]
        return self.createIndex(row, 0, parent_node)

    def rowCount(
        self, parent: QModelIndex | QPersistentModelIndex = QModelIndex()
    ) -> int:
        node = self._node_from_index(parent)
        return len(node.children)

    def columnCount(
        self, parent: QModelIndex | QPersistentModelIndex = QModelIndex()
    ) -> int:
        return 1

    def flags(
        self,
        index: QModelIndex | QPersistentModelIndex,
    ) -> Qt.ItemFlag:
        flags = super().flags(index)
        if index.isValid():
            node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
            if node.tree_node and not node.tree_node.selectable:
                flags &= ~Qt.ItemFlag.ItemIsSelectable
        return flags

    def data(
        self,
        index: QModelIndex | QPersistentModelIndex,
        role: int = Qt.ItemDataRole.DisplayRole,
    ) -> object:
        if not index.isValid():
            return None
        node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
        if role == Qt.ItemDataRole.DisplayRole:
            return node.tree_node.label if node.tree_node else None
        elif role == Qt.ItemDataRole.DecorationRole:
            if node.tree_node and node.tree_node.icon:
                return get_icon(
                    node.tree_node.icon,
                    node.tree_node.icon_color,
                    fill=node.tree_node.icon_fill,
                )
        elif role == Qt.ItemDataRole.UserRole:
            if node.tree_node and node.tree_node.data:
                return node.tree_node.data
        elif role == self.FILTER_ROLE:
            if node.tree_node:
                return node.tree_node.filter_text or node.tree_node.label
        return None

    def hasChildren(
        self, parent: QModelIndex | QPersistentModelIndex = QModelIndex()
    ) -> bool:
        return bool(self._node_from_index(parent).children)

    def get_node_id(self, index: QModelIndex) -> str | None:
        """Return the TreeNode.id for the given model index."""
        if not index.isValid():
            return None
        node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
        if node.tree_node is None:
            return None
        return node.tree_node.id

    def get_index_by_id(self, node_id: str) -> QModelIndex:
        """Return the index for a node ID, if present in the tree."""
        stack = list(self._root.children)
        while stack:
            node = stack.pop()
            if node.tree_node and node.tree_node.id == node_id:
                return self._index_for_node(node)
            stack.extend(node.children)
        return QModelIndex()

    @property
    def is_loading(self) -> bool:
        """Return True while the bulk fetch is in-flight."""
        return self._is_loading

    def reset(self) -> None:
        """Clear the model and re-run the bulk fetch.

        Also the recommended recovery path after a failed fetch (i.e.
        when :attr:`fetch_error` was emitted).
        """
        old_ctx = self._context_id
        if not self._no_async:
            get_task_queue().clear_context_tasks(old_ctx)
        self._reset_counter += 1
        self._context_id = f"btm_{id(self)}_v{self._reset_counter}"
        self.beginResetModel()
        self._root = _InternalNode(tree_node=None, parent=None)
        self._all_nodes = {self._root}
        self.endResetModel()
        self._fetch_all_async()

class LazyTreeModel(QAbstractItemModel):
    """Qt model that lazily loads children via a callback.

    Children are fetched on demand using Qt's canFetchMore/fetchMore
    protocol. The root level is loaded asynchronously on construction.

    Args:
        fetch_children: Callable that takes a parent node ID (``None``
            for root) and returns a list of :class:`TreeNode` instances.
        no_async: When ``True``, children are fetched synchronously on
            the main thread instead of via the :class:`AsyncTaskQueue`.
            Useful in tests to avoid worker-thread/paint-event races.
        parent: Optional parent QObject.

    .. warning::
        **Async mode and ``expandAll()``**: Because root children are
        fetched asynchronously, calling ``view.expandAll()`` immediately
        after constructing the model will silently expand nothing —
        ``hasChildren()`` returns ``False`` until the first fetch
        completes.  Connect to :attr:`loading_changed` and call
        ``expandAll()`` once it emits ``False``::

            model.loading_changed.connect(
                lambda loading: view.expandAll() if not loading else None
            )

    Example::

        def fetch(parent_id):
            if parent_id is None:
                return [TreeNode("root", "Root", has_children=True)]
            return []

        model = LazyTreeModel(fetch_children=fetch)
    """

    #: Role exposing TreeNode.filter_text (falls back to the label) for
    #: search proxies to match against, independent of DisplayRole.
    FILTER_ROLE = Qt.ItemDataRole.UserRole + 1

    loading_changed = Signal(bool)  # True while any fetch is in-flight
    fetch_error = Signal(str)  # error message when a fetch fails
    pending_count_changed = Signal(int)  # number of in-flight fetch tasks

    def __init__(
        self,
        fetch_children: Callable[[str | None], list[TreeNode]],
        no_async: bool = False,
        parent: QObject | None = None,
    ) -> None:
        super().__init__(parent)
        self._fetch_children = fetch_children
        self._no_async: bool = no_async
        self._reset_counter: int = 0
        self._context_id: str = f"ltm_{id(self)}_v0"
        self._pending_tasks: int = 0
        self._root = _InternalNode(tree_node=None, parent=None)
        # _all_nodes keeps every node reachable so Python's GC does not
        # collect objects that are held only via QModelIndex.internalPointer().
        self._all_nodes: set[_InternalNode] = {self._root}
        self._fetch_children_async(self._root)

    def _update_loading_state(self) -> None:
        """Emit loading-progress signals based on current pending task
        count."""
        self.pending_count_changed.emit(self._pending_tasks)
        self.loading_changed.emit(self._pending_tasks > 0)

    def _index_for_node(self, node: _InternalNode) -> QModelIndex:
        """Return the QModelIndex that identifies node (column 0).

        Args:
            node: A non-root internal node.

        Returns:
            Valid QModelIndex for non-root nodes, invalid for root.
        """
        if node.is_root:
            return QModelIndex()
        parent = node.parent
        row = parent.children.index(node)  # type: ignore[union-attr]
        return self.createIndex(row, 0, node)

    def _on_children_ready(
        self,
        node: _InternalNode,
        context_id: str,
        results: list[TreeNode] | None,
    ) -> None:
        """Handle fetch results on the main thread.

        Called via Qt.QueuedConnection from AsyncTaskQueue so it always
        executes on the Qt main thread, making it safe to call Qt model
        mutation methods.

        Args:
            node: The node whose children were fetched.
            context_id: The context_id active when the task was enqueued.
                Used to discard stale results.
            results: TreeNode list returned by fetch_children, or None on
                error.
        """
        # Always release the fetching guard so the node is never stuck
        # in-flight, even for stale callbacks.
        node.is_fetching = False

        # Discard stale results (model was reset while this task was running).
        if context_id != self._context_id or node not in self._all_nodes:
            log.debug(
                "Discarding stale fetch result for node %r (ctx %s vs %s)",
                node.tree_node.id if node.tree_node else "root",
                context_id,
                self._context_id,
            )
            return

        self._pending_tasks = max(0, self._pending_tasks - 1)

        if results is None:
            node_id = node.tree_node.id if node.tree_node else "root"
            self.fetch_error.emit(
                f"Failed to fetch children for node {node_id!r}"
            )
            # Do NOT set children_loaded=True on error so that the
            # caller can retry by calling reset().  The node stays in
            # a "not loaded" state; is_fetching was already cleared
            # above so a subsequent fetchMore() or reset() can re-try.
            self._update_loading_state()
            return

        new_nodes = [
            _InternalNode(tree_node=tn, parent=node) for tn in results
        ]
        if new_nodes:
            parent_index = self._index_for_node(node)
            self.beginInsertRows(parent_index, 0, len(new_nodes) - 1)
            node.children = new_nodes
            for child in new_nodes:
                self._all_nodes.add(child)
            node.children_loaded = True
            self.endInsertRows()
        else:
            node.children_loaded = True

        self._update_loading_state()

    def _fetch_children_async(self, node: _InternalNode) -> None:
        """Enqueue an async task to fetch children for *node*.

        Uses a per-node re-entrancy guard (``node.is_fetching``) to
        prevent duplicate tasks.  Results are delivered back on the main
        thread via :meth:`_on_children_ready`.

        Args:
            node: The node whose children to fetch.
        """
        if node.is_fetching or node.children_loaded:
            return
        node.is_fetching = True

        parent_id = None if node.is_root else node.tree_node.id  # type: ignore[union-attr]
        ctx = self._context_id

        self._pending_tasks += 1
        self._update_loading_state()

        if self._no_async:
            try:
                result: list[TreeNode] | None = self._fetch_children(parent_id)
            except Exception:
                log.exception(
                    "fetch_children raised for parent_id=%r", parent_id
                )
                result = None
            self._on_children_ready(node, ctx, result)
            return

        task = AsyncTask(
            name=f"fetch_children_{parent_id or 'root'}",
            function=lambda: self._fetch_children(parent_id),
            callback=lambda result: self._on_children_ready(node, ctx, result),
            priority=1,
            context_id=ctx,
            cancellable=True,
        )
        get_task_queue().enqueue(task)

    def _node_from_index(
        self, index: QModelIndex | QPersistentModelIndex
    ) -> _InternalNode:
        """Return the internal node for a model index.

        Args:
            index: A valid or invalid QModelIndex.

        Returns:
            The corresponding _InternalNode; returns root for invalid
            indexes.
        """
        if not index.isValid():
            return self._root
        return index.internalPointer()  # type: ignore[return-value]

    # keep old name as alias for backward compatibility
    def _node_for_index(self, index: QModelIndex) -> _InternalNode:
        """Alias for _node_from_index."""
        return self._node_from_index(index)

    def index(
        self,
        row: int,
        column: int,
        parent: QModelIndex | QPersistentModelIndex = QModelIndex(),
    ) -> QModelIndex:
        """Create a model index for the given row/column under parent.

        Args:
            row: Row number under the parent.
            column: Column number (always 0).
            parent: Parent index.

        Returns:
            A valid QModelIndex, or an invalid one if out of range.
        """
        parent_node = self._node_from_index(parent)
        if row < 0 or row >= len(parent_node.children):
            return QModelIndex()
        child_node = parent_node.children[row]
        return self.createIndex(row, column, child_node)

    def parent(self, index: QModelIndex) -> QModelIndex:  # type: ignore[override]
        """Return the parent index of the given index.

        Args:
            index: A model index whose parent is requested.

        Returns:
            Parent QModelIndex, or an invalid index if parent is root.
        """
        if not index.isValid():
            return QModelIndex()

        node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
        parent_node = node.parent
        if parent_node is None or parent_node.is_root:
            return QModelIndex()

        # grandparent is always non-None for a non-root parent_node
        grandparent = parent_node.parent
        row = grandparent.children.index(parent_node)  # type: ignore[union-attr]
        return self.createIndex(row, 0, parent_node)

    def rowCount(
        self, parent: QModelIndex | QPersistentModelIndex = QModelIndex()
    ) -> int:
        """Return the number of rows under the given parent.

        Args:
            parent: The parent index.

        Returns:
            Number of loaded children. Returns 0 for nodes whose
            children have not been fetched yet; Qt uses hasChildren()
            to decide whether to show a disclosure triangle.
        """
        node = self._node_from_index(parent)
        return len(node.children)

    def columnCount(
        self, parent: QModelIndex | QPersistentModelIndex = QModelIndex()
    ) -> int:
        """Return the number of columns (always 1).

        Args:
            parent: Unused parent index.

        Returns:
            Always 1.
        """
        return 1

    def flags(
        self,
        index: QModelIndex | QPersistentModelIndex,
    ) -> Qt.ItemFlag:
        flags = super().flags(index)
        if index.isValid():
            node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
            if node.tree_node and not node.tree_node.selectable:
                flags &= ~Qt.ItemFlag.ItemIsSelectable
        return flags

    def data(
        self,
        index: QModelIndex | QPersistentModelIndex,
        role: int = Qt.ItemDataRole.DisplayRole,
    ) -> object:
        """Return data for the given index and role.

        Args:
            index: The model index to query.
            role: The data role.

        Returns:
            The node label for DisplayRole, None otherwise.
        """
        if not index.isValid():
            return None
        node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
        if role == Qt.ItemDataRole.DisplayRole:
            return node.tree_node.label if node.tree_node else None
        elif role == Qt.ItemDataRole.DecorationRole:
            if node.tree_node and node.tree_node.icon:
                return get_icon(
                    node.tree_node.icon,
                    node.tree_node.icon_color,
                    fill=node.tree_node.icon_fill,
                )
        elif role == Qt.ItemDataRole.UserRole:
            if node.tree_node and node.tree_node.data:
                return node.tree_node.data
        elif role == self.FILTER_ROLE:
            if node.tree_node:
                return node.tree_node.filter_text or node.tree_node.label
        return None

    def hasChildren(
        self, parent: QModelIndex | QPersistentModelIndex = QModelIndex()
    ) -> bool:
        """Return whether the node at parent has or can have children.

        This is used by Qt to decide whether to draw the disclosure
        triangle, even before children are loaded.

        Args:
            parent: The parent index.

        Returns:
            True if the node has loaded children or has_children is
            True on the TreeNode.
        """
        node = self._node_from_index(parent)
        if node.is_root:
            return bool(node.children)
        # node.tree_node is non-None here (is_root is False)
        if node.children_loaded:
            return bool(node.children)
        return node.tree_node.has_children  # type: ignore[union-attr]

    def canFetchMore(
        self, parent: QModelIndex | QPersistentModelIndex
    ) -> bool:
        """Return True if more children can be fetched for parent.

        Args:
            parent: The parent index to check.

        Returns:
            True if the node has potential children but they have not
            been loaded yet.
        """
        node = self._node_from_index(parent)
        if node.is_root:
            return False
        # node.tree_node is non-None here (is_root is False)
        return (  # type: ignore[union-attr]
            node.tree_node.has_children and not node.children_loaded
        )

    def fetchMore(self, parent: QModelIndex | QPersistentModelIndex) -> None:
        """Fetch and insert children for the node at parent.

        Args:
            parent: The parent index whose children should be loaded.
        """
        self._fetch_children_async(self._node_from_index(parent))

    def get_node_id(self, index: QModelIndex) -> str | None:
        """Return the TreeNode.id for the given model index.

        Args:
            index: A valid QModelIndex.

        Returns:
            The node's string ID, or None if the index is invalid.
        """
        if not index.isValid():
            return None
        node: _InternalNode = index.internalPointer()  # type: ignore[assignment]
        if node.tree_node is None:
            return None
        return node.tree_node.id

    def get_index_by_id(self, node_id: str) -> QModelIndex:
        """Return the loaded index for a node ID, if present."""
        stack = list(self._root.children)
        while stack:
            node = stack.pop()
            if node.tree_node and node.tree_node.id == node_id:
                return self._index_for_node(node)
            stack.extend(node.children)
        return QModelIndex()

    @property
    def is_loading(self) -> bool:
        """Return True while at least one fetch task is in-flight."""
        return self._pending_tasks > 0

    def reset(self) -> None:
        """Reset the model to its initial state and re-fetch root children.

        This is the recommended recovery path after a failed root fetch
        (i.e. when :attr:`fetch_error` was emitted for the root node).
        """
        old_ctx = self._context_id
        if not self._no_async:
            get_task_queue().clear_context_tasks(old_ctx)
        self._reset_counter += 1
        self._context_id = f"ltm_{id(self)}_v{self._reset_counter}"
        self._pending_tasks = 0
        self.beginResetModel()
        self._root = _InternalNode(tree_node=None, parent=None)
        self._all_nodes = {self._root}
        self.endResetModel()
        # Emit loading signals *after* endResetModel so that any slot
        # connected to loading_changed observes a consistent model state.
        self._update_loading_state()
        self._fetch_children_async(self._root)

@dataclass
class TreeNode:
    """Represents a single node in a tree model.

    Attributes:
        id: Unique identifier for this node.
        label: Display label shown in the view.
        has_children: Whether this node can have children.
        icon: Optional icon name or path.
        selectable: Whether this node can be selected in the view.
        data: Arbitrary extra data associated with this node.
        filter_text: Casefolded text to match search queries against
            (e.g. a folder's path plus its label-based path). Falls
            back to ``label`` when empty; see ``FILTER_ROLE`` on
            :class:`LazyTreeModel`/:class:`BulkTreeModel`.
    """

    id: str
    label: str
    has_children: bool = False
    icon: str = ""
    icon_color: str = "#f4f5f5"
    icon_fill: bool = False
    selectable: bool = True
    data: dict = field(default_factory=dict)
    filter_text: str = ""