"""
Navigation Tree Module

Phase 1.2 - Code Simplification (011-code-simplification)
Extracted from navigation.py

Spec Traceability:
------------------
- 002-ascent-functionality: Branching navigation tree (time-travel support)

Contains:
- NavigationTreeNode: Single node in navigation tree
- NavigationTree: Branching tree structure
"""

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional
import uuid

from intuitiveness.complexity import Dataset, ComplexityLevel
from intuitiveness.navigation.state import NavigationAction


@dataclass
class NavigationTreeNode:
    """
    A single node in the navigation tree, supporting branching paths.

    Per FR-021, each node records:
    - (a) The navigation step taken (action)
    - (b) Decision made at each step (decision_description)
    - (c) Generated output snapshot (output_snapshot)

    Attributes:
        id: Unique identifier for this node
        level: Complexity level at this node (L0-L4)
        dataset_snapshot: Full dataset at this point (for restoration)
        parent_id: Parent node ID (None for root)
        children_ids: List of child node IDs (branches)
        action: Action that created this node ("entry", "descend", "ascend", "restore")
        timestamp: When this node was created
        metadata: Additional info (enrichment used, dimensions added, etc.)
        decision_description: Human-readable description of decision (FR-021)
        output_snapshot: Summary of output at this step (FR-021)
    """
    id: str
    level: ComplexityLevel
    dataset_snapshot: Dataset
    parent_id: Optional[str]
    children_ids: List[str] = field(default_factory=list)
    action: str = "entry"
    timestamp: datetime = field(default_factory=datetime.now)
    metadata: Dict[str, Any] = field(default_factory=dict)
    decision_description: str = ""  # FR-021
    output_snapshot: Dict[str, Any] = field(default_factory=dict)  # FR-021

    @property
    def depth(self) -> int:
        """Depth in tree (for UI indentation). Root is depth 0."""
        return self.metadata.get('_depth', 0)

    def to_dict(self) -> Dict[str, Any]:
        """Serialize for JSON export (excludes dataset_snapshot)."""
        return {
            "id": self.id,
            "level": self.level.value,
            "level_name": self.level.name,
            "parent_id": self.parent_id,
            "children_ids": self.children_ids.copy(),
            "action": self.action,
            "timestamp": self.timestamp.isoformat(),
            "metadata": {k: v for k, v in self.metadata.items() if not k.startswith('_')},
            "decision_description": self.decision_description,
            "output_snapshot": self.output_snapshot
        }


class NavigationTree:
    """
    Branching tree structure tracking all navigation decisions.

    Supports time-travel navigation by preserving multiple exploration branches.
    Replaces linear NavigationHistory for sessions requiring branching.

    Usage:
        >>> tree = NavigationTree(root_dataset)
        >>> tree.branch(NavigationAction.DESCEND, new_dataset, {"step": "entities"})
        >>> tree.restore("node_abc123")  # Time-travel back
    """

    def __init__(self, root_dataset: Dataset):
        """
        Initialize navigation tree with root node.

        Args:
            root_dataset: The L4 dataset at entry point
        """
        self._nodes: Dict[str, NavigationTreeNode] = {}
        self._root_id = str(uuid.uuid4())
        self._current_id = self._root_id

        # Create root node
        root_node = NavigationTreeNode(
            id=self._root_id,
            level=root_dataset.complexity_level,
            dataset_snapshot=root_dataset,
            parent_id=None,
            action=NavigationAction.ENTRY.value,
            metadata={'_depth': 0}
        )
        self._nodes[self._root_id] = root_node

    @property
    def root_id(self) -> str:
        """Get the root node ID."""
        return self._root_id

    @property
    def current_id(self) -> str:
        """Get the current node ID."""
        return self._current_id

    @property
    def current_node(self) -> NavigationTreeNode:
        """Get the current node."""
        return self._nodes[self._current_id]

    @property
    def nodes(self) -> Dict[str, NavigationTreeNode]:
        """Get all nodes."""
        return self._nodes

    def branch(
        self,
        action: NavigationAction,
        dataset: Dataset,
        metadata: Optional[Dict[str, Any]] = None,
        decision_description: str = "",
        output_snapshot: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Create a new branch from current position.

        Per FR-021, records decision_description and output_snapshot at each node.

        Args:
            action: The navigation action taken (DESCEND, ASCEND, etc.)
            dataset: The dataset at this new position
            metadata: Additional info (enrichment used, dimensions added)
            decision_description: Human-readable description of the decision (FR-021)
            output_snapshot: Summary of output at this step (FR-021)

        Returns:
            ID of the newly created node
        """
        new_id = str(uuid.uuid4())
        parent = self._nodes[self._current_id]
        parent_depth = parent.metadata.get('_depth', 0)

        # Build metadata with depth
        node_metadata = metadata.copy() if metadata else {}
        node_metadata['_depth'] = parent_depth + 1

        # Generate output_snapshot if not provided (FR-021)
        if output_snapshot is None:
            output_snapshot = self._generate_output_snapshot(dataset)

        # Create new node
        new_node = NavigationTreeNode(
            id=new_id,
            level=dataset.complexity_level,
            dataset_snapshot=dataset,
            parent_id=self._current_id,
            action=action.value if isinstance(action, NavigationAction) else action,
            metadata=node_metadata,
            decision_description=decision_description,
            output_snapshot=output_snapshot
        )

        # Add to tree
        self._nodes[new_id] = new_node
        parent.children_ids.append(new_id)

        # Move current pointer
        self._current_id = new_id

        return new_id

    def _generate_output_snapshot(self, dataset: Dataset) -> Dict[str, Any]:
        """
        Generate an output snapshot for a dataset (FR-021).

        Args:
            dataset: The dataset to summarize

        Returns:
            Dict with output summary info
        """
        snapshot = {
            "level": dataset.complexity_level.value,
            "level_name": dataset.complexity_level.name
        }

        data = dataset.get_data()

        if dataset.complexity_level == ComplexityLevel.LEVEL_0:
            snapshot["type"] = "datum"
            snapshot["value"] = str(data)
        elif dataset.complexity_level == ComplexityLevel.LEVEL_1:
            snapshot["type"] = "vector"
            if hasattr(data, '__len__'):
                snapshot["length"] = len(data)
        elif dataset.complexity_level == ComplexityLevel.LEVEL_2:
            snapshot["type"] = "dataframe"
            if hasattr(data, 'shape'):
                snapshot["row_count"] = data.shape[0]
                snapshot["columns"] = list(data.columns) if hasattr(data, 'columns') else []
        elif dataset.complexity_level == ComplexityLevel.LEVEL_3:
            snapshot["type"] = "graph"
            if hasattr(data, 'number_of_nodes'):
                snapshot["node_count"] = data.number_of_nodes()
                snapshot["edge_count"] = data.number_of_edges()
            elif hasattr(data, 'shape'):
                snapshot["row_count"] = data.shape[0]
        elif dataset.complexity_level == ComplexityLevel.LEVEL_4:
            snapshot["type"] = "unlinkable"
            if isinstance(data, dict):
                snapshot["source_count"] = len(data)

        return snapshot

    def restore(self, node_id: str) -> Dataset:
        """
        Restore navigation state to a previous node (time-travel).

        Args:
            node_id: ID of the node to restore to

        Returns:
            The dataset at the restored node

        Raises:
            KeyError: If node_id not found in tree
        """
        if node_id not in self._nodes:
            raise KeyError(f"Node '{node_id}' not found in navigation tree")

        self._current_id = node_id
        return self._nodes[node_id].dataset_snapshot

    def get_current_branch_path(self) -> List[NavigationTreeNode]:
        """
        Get path from root to current node.

        Returns:
            List of nodes from root to current position
        """
        path = []
        current = self._current_id

        while current is not None:
            node = self._nodes[current]
            path.insert(0, node)
            current = node.parent_id

        return path

    def get_all_branches(self) -> List[List[NavigationTreeNode]]:
        """
        Get all paths from root to leaf nodes.

        Returns:
            List of paths, where each path is a list of nodes from root to leaf
        """
        branches = []

        def find_leaves(node_id: str, path: List[NavigationTreeNode]):
            node = self._nodes[node_id]
            new_path = path + [node]

            if not node.children_ids:
                branches.append(new_path)
            else:
                for child_id in node.children_ids:
                    find_leaves(child_id, new_path)

        find_leaves(self._root_id, [])
        return branches

    def export_to_json(self) -> Dict[str, Any]:
        """
        Export full tree for JSON visualization.

        Returns:
            Dict with nodes, root_id, and current_id
        """
        return {
            "nodes": [node.to_dict() for node in self._nodes.values()],
            "root_id": self._root_id,
            "current_id": self._current_id
        }

    def get_node(self, node_id: str) -> NavigationTreeNode:
        """Get a specific node by ID."""
        if node_id not in self._nodes:
            raise KeyError(f"Node '{node_id}' not found")
        return self._nodes[node_id]

    def __len__(self) -> int:
        """Number of nodes in tree."""
        return len(self._nodes)
