from __future__ import annotations

from collections import defaultdict, deque
from typing import Any

from nexus.core.events.bus import EventBus
from nexus.core.events.types import GraphMutationEvent
from nexus.core.graph.loop_detective import LoopDetective, LoopStrategy
from nexus.core.graph.node import AgentNode, NodeConfig, NodeState
from nexus.resilience.errors import GraphMutationError


class DynamicGraph:
    """Mutable DAG for runtime orchestration."""

    def __init__(self, bus: EventBus, loop_detective: LoopDetective) -> None:
        self.bus = bus
        self.loop_detective = loop_detective
        self._nodes: dict[str, AgentNode] = {}
        self._adjacency: dict[str, list[str]] = defaultdict(list)

    def add_node(self, config: NodeConfig) -> AgentNode:
        """Add node to graph and emit mutation event."""
        if config.node_id in self._nodes:
            raise GraphMutationError(f"Duplicate node_id: {config.node_id}", node_id=config.node_id)
        node = AgentNode(config)
        self._nodes[config.node_id] = node
        self._adjacency.setdefault(config.node_id, [])
        self.bus.publish(
            GraphMutationEvent(source_id="dynamic_graph", mutation_type="add_node", payload={"node_id": config.node_id})
        )
        return node

    def remove_node(self, node_id: str) -> None:
        """Remove node and all related edges."""
        if node_id not in self._nodes:
            return
        self._nodes.pop(node_id)
        self._adjacency.pop(node_id, None)
        for src, dsts in self._adjacency.items():
            self._adjacency[src] = [d for d in dsts if d != node_id]

    def add_edge(self, from_id: str, to_id: str) -> None:
        """Add directed edge and reject mutations that introduce cycles."""
        if from_id not in self._nodes or to_id not in self._nodes:
            raise GraphMutationError(f"Cannot add edge {from_id}->{to_id}; missing node")
        self._adjacency[from_id].append(to_id)
        cycle = self.loop_detective.detect_cycles(dict(self._adjacency))
        if cycle.has_cycle:
            self._adjacency[from_id].remove(to_id)
            if self.loop_detective.should_break(from_id, 0, LoopStrategy.HARD_STOP):
                raise GraphMutationError(f"Edge creates cycle: {' -> '.join(cycle.cycle_path)}")

    def get_ready_nodes(self, completed: set[str]) -> list[AgentNode]:
        """Return pending nodes whose dependencies are all completed."""
        return [
            node
            for node in self._nodes.values()
            if node.state == NodeState.PENDING and node.is_ready(completed)
        ]

    def topological_order(self) -> list[str]:
        """Compute Kahn topological ordering."""
        indeg = {node_id: 0 for node_id in self._nodes}
        for src, dsts in self._adjacency.items():
            for dst in dsts:
                indeg[dst] += 1
        q = deque([n for n, d in indeg.items() if d == 0])
        out: list[str] = []
        while q:
            cur = q.popleft()
            out.append(cur)
            for nxt in self._adjacency.get(cur, []):
                indeg[nxt] -= 1
                if indeg[nxt] == 0:
                    q.append(nxt)
        if len(out) != len(self._nodes):
            raise GraphMutationError("Graph contains a cycle and cannot be sorted")
        return out

    async def apply_mutation(self, event: GraphMutationEvent) -> None:
        """Apply graph mutation event to graph state."""
        payload = event.payload
        if event.mutation_type == "add_node":
            self.add_node(NodeConfig.model_validate(payload["config"]))
        elif event.mutation_type == "remove_node":
            self.remove_node(str(payload["node_id"]))
        elif event.mutation_type == "add_edge":
            self.add_edge(str(payload["from_id"]), str(payload["to_id"]))
        elif event.mutation_type == "fork":
            return
        else:
            raise GraphMutationError(f"Unknown mutation type: {event.mutation_type}")

    def to_dict(self) -> dict[str, Any]:
        """Serialize graph state for checkpointing."""
        return {
            "nodes": {node_id: node.to_dict() for node_id, node in self._nodes.items()},
            "adjacency": dict(self._adjacency),
        }

    @classmethod
    def from_dict(cls, data: dict, bus: EventBus, detective: LoopDetective) -> DynamicGraph:
        """Recreate graph from serialized state."""
        graph = cls(bus=bus, loop_detective=detective)
        for node_id, payload in data.get("nodes", {}).items():
            cfg = NodeConfig.model_validate(payload["config"])
            cfg.task_fn = None
            node = AgentNode(cfg)
            graph._nodes[node_id] = node
        graph._adjacency = defaultdict(list, data.get("adjacency", {}))
        return graph

    def node(self, node_id: str) -> AgentNode:
        """Fetch node by id."""
        return self._nodes[node_id]

    def all_nodes(self) -> list[AgentNode]:
        """Return all graph nodes."""
        return list(self._nodes.values())

    def is_complete(self) -> bool:
        """Graph is complete when every node is complete or skipped."""
        return all(node.state in {NodeState.COMPLETE, NodeState.SKIPPED} for node in self._nodes.values())

    def terminal_node(self) -> AgentNode | None:
        """Return explicit terminal-tagged node or sink node."""
        for node in self._nodes.values():
            if "terminal" in node.config.tags:
                return node
        outgoing = {src for src, dsts in self._adjacency.items() if dsts}
        sinks = [n for n in self._nodes if n not in outgoing]
        if not sinks:
            return None
        return self._nodes[sinks[0]]
