"""
Tree model for displaying the EMDS Workbench hierarchy.
"""

from __future__ import annotations

from typing import Any, Dict, Optional

from .qt_compat import QIcon, QTreeWidgetItem, Qt
from .models import Analysis, LayerKind, Result, Study, Task, Workspace, WorkspaceManager


class TreeNodeData:
    """Data associated with each tree node."""

    def __init__(
        self,
        node_type: str,
        item_id: str,
        item: Any,
        metadata: Optional[Dict[str, Any]] = None,
    ):
        self.node_type = node_type
        self.item_id = item_id
        self.item = item
        self.metadata = metadata or {}


class EMDSTreeModel:
    """Manages the Advanced tree structure for EMDS Workbench."""

    THEME_ICONS = {
        "workspace": "folder",
        "study": "view-statistics",
        "analysis": "folder-saved-search",
        "task": "system-run",
        "result": "view-list-details",
        # Legacy node-type aliases accepted during transition.
        "project": "folder",
        "assessment": "view-statistics",
        "provenance": "folder-saved-search",
        "workflow": "system-run",
        "results": "view-list-details",
    }

    NODE_LABELS = {
        "workspace": "Workspace",
        "study": "Study",
        "analysis": "Analysis",
        "task": "Task",
        "result": "Result",
        "project": "Workspace",
        "assessment": "Study",
        "provenance": "Analysis",
        "workflow": "Task",
        "results": "Result",
    }

    LEGACY_NODE_TYPES = {
        "project": "workspace",
        "assessment": "study",
        "provenance": "analysis",
        "workflow": "task",
        "results": "result",
    }

    @classmethod
    def normalize_node_type(cls, node_type: Optional[str]) -> Optional[str]:
        if node_type is None:
            return None
        return cls.LEGACY_NODE_TYPES.get(node_type, node_type)

    @staticmethod
    def _build_context_metadata(
        *,
        workspace_id: Optional[str] = None,
        study_id: Optional[str] = None,
        analysis_id: Optional[str] = None,
        task_id: Optional[str] = None,
        result_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Attach both v1 and legacy identity keys while old call sites are retired."""
        merged = dict(metadata or {})
        merged.update(
            {
                "workspace_id": workspace_id or "",
                "study_id": study_id or "",
                "analysis_id": analysis_id or "",
                "task_id": task_id or "",
                "result_id": result_id or "",
                "project_id": workspace_id or "",
                "assessment_id": study_id or "",
                "provenance_id": analysis_id or "",
                "workflow_id": task_id or "",
                "results_id": result_id or "",
            }
        )
        return merged

    @classmethod
    def create_item(
        cls,
        text: str,
        node_type: str,
        item_id: str,
        item: Any,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> QTreeWidgetItem:
        """Create a tree widget item with associated data."""
        normalized_type = cls.normalize_node_type(node_type) or node_type
        tree_item = QTreeWidgetItem()
        tree_item.setText(0, text)
        tree_item.setToolTip(0, cls.NODE_LABELS.get(normalized_type, normalized_type.title()))
        icon = QIcon.fromTheme(cls.THEME_ICONS.get(normalized_type, ""))
        if not icon.isNull():
            tree_item.setIcon(0, icon)

        tree_item.setData(0, Qt.ItemDataRole.UserRole, TreeNodeData(normalized_type, item_id, item, metadata=metadata))
        tree_item.setFlags(tree_item.flags() | Qt.ItemFlag.ItemIsSelectable)
        return tree_item

    @staticmethod
    def format_result_label(result: Result) -> str:
        visible_text = "visible" if result.visible else "hidden"
        return f"{result.layer_name} ({result.feature_count} features, {visible_text})"

    @staticmethod
    def format_results_label(result: Result) -> str:
        return EMDSTreeModel.format_result_label(result)

    # Raster-support Phase 4b: study-area kind badges.
    _STUDY_KIND_BADGES = {
        LayerKind.VECTOR: "[V]",
        LayerKind.RASTER: "[R]",
    }

    @staticmethod
    def _format_study_label(study: Study) -> str:
        """Render a study-area tree label with a kind badge prefix.

        Reads the canonical Study.kind populated by Phase 1's
        apply_kind_metadata (defaults to VECTOR for legacy studies
        that pre-date the kind field).
        """
        kind = getattr(study, "kind", LayerKind.VECTOR)
        badge = EMDSTreeModel._STUDY_KIND_BADGES.get(kind, "[V]")
        return f"{badge} {study.name}"

    @staticmethod
    def populate_tree(tree_widget, workspace_manager: WorkspaceManager) -> None:
        """Populate tree with all workspaces and their hierarchies."""
        tree_widget.clear()

        for workspace in workspace_manager.workspaces:
            workspace_item = EMDSTreeModel.create_item(
                workspace.name,
                "workspace",
                workspace.id,
                workspace,
                metadata=EMDSTreeModel._build_context_metadata(
                    workspace_id=workspace.id,
                    metadata=getattr(workspace, "metadata", {}) or {},
                ),
            )
            tree_widget.addTopLevelItem(workspace_item)

            for study in workspace.studies:
                # Raster-support Phase 4b: kind badge prefix so users
                # can tell at a glance whether a study is vector or
                # raster. Reads the typed Study.kind populated by
                # apply_kind_metadata; defaults to VECTOR for studies
                # that pre-date the kind field.
                study_label = EMDSTreeModel._format_study_label(study)
                study_item = EMDSTreeModel.create_item(
                    study_label,
                    "study",
                    study.id,
                    study,
                    metadata=EMDSTreeModel._build_context_metadata(
                        workspace_id=workspace.id,
                        study_id=study.id,
                        metadata=getattr(study, "metadata", {}) or {},
                    ),
                )
                workspace_item.addChild(study_item)

                for analysis in study.analyses:
                    analysis_item = EMDSTreeModel.create_item(
                        f"Analysis: {analysis.name}",
                        "analysis",
                        analysis.id,
                        analysis,
                        metadata=EMDSTreeModel._build_context_metadata(
                            workspace_id=workspace.id,
                            study_id=study.id,
                            analysis_id=analysis.id,
                            metadata=getattr(analysis, "metadata", {}) or {},
                        ),
                    )
                    study_item.addChild(analysis_item)

                    for task in analysis.tasks:
                        EMDSTreeModel._add_task(
                            analysis_item,
                            task,
                            workspace_id=workspace.id,
                            study_id=study.id,
                            analysis_id=analysis.id,
                        )

    @staticmethod
    def _add_task(
        parent_item: QTreeWidgetItem,
        task: Task,
        *,
        workspace_id: str,
        study_id: str,
        analysis_id: str,
    ) -> QTreeWidgetItem:
        """Recursively add task and child task nodes."""
        parameters = getattr(task, "parameters", {}) or {}
        task_prefix = "Workflow Task" if (
            task.engine_type.value == "Workflow"
            or str(parameters.get("task_type") or parameters.get("task_kind") or "").strip().lower()
            in {"workflow", "workflow_definition"}
        ) else "Task"
        task_item = EMDSTreeModel.create_item(
            f"{task_prefix}: {task.name}",
            "task",
            task.id,
            task,
            metadata=EMDSTreeModel._build_context_metadata(
                workspace_id=workspace_id,
                study_id=study_id,
                analysis_id=analysis_id,
                task_id=task.id,
                metadata={
                    **(getattr(task, "metadata", {}) or {}),
                    **(getattr(task, "parameters", {}) or {}),
                },
            ),
        )
        parent_item.addChild(task_item)

        if task.result:
            result = task.result
            result_item = EMDSTreeModel.create_item(
                EMDSTreeModel.format_result_label(result),
                "result",
                result.id,
                result,
                metadata=EMDSTreeModel._build_context_metadata(
                    workspace_id=workspace_id,
                    study_id=study_id,
                    analysis_id=analysis_id,
                    task_id=task.id,
                    result_id=result.id,
                    metadata=getattr(result, "metadata", {}) or {},
                ),
            )
            task_item.addChild(result_item)

        for child_task in task.children:
            EMDSTreeModel._add_task(
                task_item,
                child_task,
                workspace_id=workspace_id,
                study_id=study_id,
                analysis_id=analysis_id,
            )

        return task_item

    @staticmethod
    def _add_workflow_step(
        parent_item: QTreeWidgetItem,
        workflow_step: Task,
        *,
        project_id: str,
        assessment_id: str,
        provenance_id: str,
    ) -> QTreeWidgetItem:
        return EMDSTreeModel._add_task(
            parent_item,
            workflow_step,
            workspace_id=project_id,
            study_id=assessment_id,
            analysis_id=provenance_id,
        )

    @staticmethod
    def collapse_all_except_path(tree_widget, active_item: Optional[QTreeWidgetItem]) -> None:
        """Collapse all items except those in the active item's parent path."""
        for index in range(tree_widget.topLevelItemCount()):
            tree_widget.collapseItem(tree_widget.topLevelItem(index))

        if active_item:
            items_to_expand = []
            current = active_item
            while current is not None:
                items_to_expand.append(current)
                current = current.parent()
            for item in reversed(items_to_expand):
                tree_widget.expandItem(item)

    @staticmethod
    def get_node_path(tree_widget, item: QTreeWidgetItem) -> Dict[str, Optional[str]]:
        """Return the identity path for a tree item."""
        path = {
            "workspace_id": None,
            "study_id": None,
            "analysis_id": None,
            "task_id": None,
            "result_id": None,
            "project_id": None,
            "assessment_id": None,
            "provenance_id": None,
            "workflow_id": None,
            "results_id": None,
        }

        current = item
        while current is not None:
            data = current.data(0, Qt.ItemDataRole.UserRole)
            if isinstance(data, TreeNodeData):
                node_type = EMDSTreeModel.normalize_node_type(data.node_type)
                if node_type == "workspace":
                    path["workspace_id"] = data.item_id
                    path["project_id"] = data.item_id
                elif node_type == "study":
                    path["study_id"] = data.item_id
                    path["assessment_id"] = data.item_id
                elif node_type == "analysis":
                    path["analysis_id"] = data.item_id
                    path["provenance_id"] = data.item_id
                elif node_type == "task":
                    # Walk is bottom-up (deepest item first), so only the FIRST
                    # task seen is the clicked node. Without this guard a nested
                    # child task would be overwritten by its parent task as we
                    # ascend, resolving the wrong task (e.g. tea21 -> gdf3).
                    if path["task_id"] is None:
                        path["task_id"] = data.item_id
                        path["workflow_id"] = data.item_id
                elif node_type == "result":
                    path["result_id"] = data.item_id
                    path["results_id"] = data.item_id
            current = current.parent()

        return path
