"""
Graph Utilities for AGoT's DAG (Directed Acyclic Graph)

Manages the nested graph structure:
  G_h = (V_h, E_h, F_h)
    V_h  — nodes (thoughts)
    E_h  — edges (dependencies)
    F_h  — features (scores, metadata)

Supports:
- Node addition with depth tracking
- Edge validation (no cycles)
- Graph traversal
- Sub-graph nesting (recursive AGoT calls)
- Graph serialization for prompts
"""

from __future__ import annotations
from typing import Any, Optional


class AGoTGraph:
    """A DAG of reasoning thoughts with nested sub-graphs."""

    def __init__(self, parent: Optional[AGoTGraph] = None, depth: int = 0):
        self.nodes: list[dict] = []       # V_h: list of {id, content, depth, score, ...}
        self.edges: list[dict] = []        # E_h: list of {from, to, label}
        self.features: dict = {}           # F_h: metadata
        self.parent: Optional[AGoTGraph] = parent
        self.depth: int = depth
        self.sub_graphs: dict[int, AGoTGraph] = {}  # node_id -> AGoTGraph (recursive)
        self._next_id: int = 0
        self.final_answer: Optional[str] = None

    # ── Node Operations ─────────────────────────────────────────

    def add_node(self, content: str, **kwargs) -> int:
        """Add a thought node. Returns its ID."""
        node_id = self._next_id
        self._next_id += 1
        node = {
            "id": node_id,
            "content": content,
            "depth": self.depth,
            "score": 0.0,
            "is_complex": False,
            "is_final": False,
            **kwargs,
        }
        self.nodes.append(node)
        return node_id

    def get_node(self, node_id: int) -> Optional[dict]:
        for n in self.nodes:
            if n["id"] == node_id:
                return n
        return None

    def update_node(self, node_id: int, **updates):
        node = self.get_node(node_id)
        if node:
            node.update(updates)

    # ── Edge Operations ─────────────────────────────────────────

    def add_edge(self, from_id: int, to_id: int, label: str = "") -> bool:
        """Add a directed edge. Returns False if it would create a cycle."""
        # Validate both nodes exist
        if not self.get_node(from_id) or not self.get_node(to_id):
            return False
        # Validate no duplicate
        for e in self.edges:
            if e["from"] == from_id and e["to"] == to_id:
                return True  # already exists
        # Cycle check: would to_id → ... → from_id?
        if self._would_cycle(to_id, from_id):
            return False
        self.edges.append({"from": from_id, "to": to_id, "label": label})
        return True

    def _would_cycle(self, start: int, target: int, visited: set = None) -> bool:
        """DFS: can we reach `target` starting from `start`?"""
        if visited is None:
            visited = set()
        if start == target:
            return True
        if start in visited:
            return False
        visited.add(start)
        for e in self.edges:
            if e["from"] == start:
                if self._would_cycle(e["to"], target, visited):
                    return True
        return False

    # ── Sub-graph Operations ───────────────────────────────────

    def get_or_create_subgraph(self, node_id: int) -> AGoTGraph:
        """Get or create a nested sub-graph for a complex node."""
        if node_id not in self.sub_graphs:
            self.sub_graphs[node_id] = AGoTGraph(parent=self, depth=self.depth + 1)
        return self.sub_graphs[node_id]

    def get_all_nodes_recursive(self) -> list[dict]:
        """Get all nodes including those in sub-graphs."""
        all_nodes = list(self.nodes)
        for sg in self.sub_graphs.values():
            all_nodes.extend(sg.get_all_nodes_recursive())
        return all_nodes

    def get_all_edges_recursive(self) -> list[dict]:
        all_edges = list(self.edges)
        for sg in self.sub_graphs.values():
            all_edges.extend(sg.get_all_edges_recursive())
        return all_edges

    # ── Query ──────────────────────────────────────────────────

    @property
    def node_count(self) -> int:
        return len(self.nodes)

    @property
    def edge_count(self) -> int:
        return len(self.edges)

    def get_leaves(self) -> list[dict]:
        """Nodes with no outgoing edges."""
        outgoing = {e["from"] for e in self.edges}
        return [n for n in self.nodes if n["id"] not in outgoing]

    def get_root_nodes(self) -> list[dict]:
        """Nodes with no incoming edges."""
        incoming = {e["to"] for e in self.edges}
        return [n for n in self.nodes if n["id"] not in incoming]

    # ── Serialization ─────────────────────────────────────────

    def to_dict(self) -> dict:
        return {
            "nodes": self.nodes,
            "edges": self.edges,
            "features": self.features,
            "depth": self.depth,
            "final_answer": self.final_answer,
            "sub_graphs": {str(k): v.to_dict() for k, v in self.sub_graphs.items()},
        }

    @classmethod
    def from_dict(cls, data: dict, parent: Optional[AGoTGraph] = None) -> AGoTGraph:
        g = cls(parent=parent, depth=data.get("depth", 0))
        g.nodes = data.get("nodes", [])
        g.edges = data.get("edges", [])
        g.features = data.get("features", {})
        g.final_answer = data.get("final_answer")
        g._next_id = max((n["id"] for n in g.nodes), default=0) + 1
        for k, v in data.get("sub_graphs", {}).items():
            g.sub_graphs[int(k)] = cls.from_dict(v, parent=g)
        return g

    def __repr__(self) -> str:
        return f"AGoTGraph(depth={self.depth}, nodes={self.node_count}, edges={self.edge_count})"
