"""
GraphOrchestrator — static graph execution mode (like LangGraph).

Design notes:
  - Users define nodes (agents or functions) and edges (data flow)
  - Supports conditional edges for dynamic branching within a static topology
  - Typed shared state flows through the graph
  - Automatically detects parallelizable paths
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any, Generic, TypeVar

from pydantic import BaseModel

from nexus.core.state import SharedState

logger = logging.getLogger("nexus.orchestration.graph")

T = TypeVar("T", bound=BaseModel)

# A node function takes the shared state and returns updated fields
NodeFunc = Callable[[dict[str, Any]], Awaitable[dict[str, Any]]]
# A condition function takes the state and returns the name of the next node
ConditionFunc = Callable[[dict[str, Any]], str]

# Special sentinels
START = "__start__"
END = "__end__"


@dataclass
class Node:
    """A node in the execution graph."""

    name: str
    func: NodeFunc | None = None
    agent_id: str | None = None  # If set, the node delegates to this agent
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class Edge:
    """An edge in the execution graph (static or conditional)."""

    source: str
    target: str
    condition: ConditionFunc | None = None  # If set, this is a conditional edge
    label: str = ""


class GraphOrchestrator(Generic[T]):
    """
    Static graph-based orchestration — define nodes and edges, then execute.

    Inspired by LangGraph's StateGraph but with simpler semantics.

    Usage:
        class MyState(BaseModel):
            plan: str = ""
            code: str = ""
            review: str = ""

        graph = GraphOrchestrator[MyState]()

        # Add nodes
        graph.add_node("planner", planner_func)
        graph.add_node("coder", coder_func)
        graph.add_node("reviewer", reviewer_func)

        # Add edges
        graph.add_edge(START, "planner")
        graph.add_edge("planner", "coder")
        graph.add_edge("coder", "reviewer")
        graph.add_conditional_edge("reviewer", should_retry, {
            "retry": "planner",
            "pass": END,
        })

        # Run
        state = SharedState(MyState())
        final_state = await graph.run(state)
    """

    def __init__(self) -> None:
        self._nodes: dict[str, Node] = {}
        self._edges: list[Edge] = []
        self._conditional_edges: dict[str, tuple[ConditionFunc, dict[str, str]]] = {}

    def add_node(
        self,
        name: str,
        func: NodeFunc | None = None,
        agent_id: str | None = None,
    ) -> None:
        """Add a node to the graph."""
        if name in (START, END):
            raise ValueError(f"Cannot use reserved name '{name}' for a node")
        self._nodes[name] = Node(name=name, func=func, agent_id=agent_id)

    def add_edge(self, source: str, target: str, label: str = "") -> None:
        """Add a static edge between two nodes."""
        self._edges.append(Edge(source=source, target=target, label=label))

    def add_conditional_edge(
        self,
        source: str,
        condition: ConditionFunc,
        route_map: dict[str, str],
    ) -> None:
        """
        Add a conditional edge.

        The condition function receives the current state and returns a key.
        The route_map maps that key to the next node name.
        """
        self._conditional_edges[source] = (condition, route_map)

    def _get_next_nodes(self, current: str, state_dict: dict[str, Any]) -> list[str]:
        """Determine the next node(s) after the current one."""
        # Check conditional edges first
        if current in self._conditional_edges:
            condition, route_map = self._conditional_edges[current]
            result = condition(state_dict)
            target = route_map.get(result)
            if target:
                return [target]
            logger.warning(
                f"Conditional edge from '{current}' returned '{result}' "
                f"which has no mapping. Available: {list(route_map.keys())}"
            )
            return []

        # Static edges
        targets = [e.target for e in self._edges if e.source == current]
        return targets

    async def run(
        self,
        state: SharedState[T],
        max_steps: int = 100,
    ) -> SharedState[T]:
        """
        Execute the graph starting from START, following edges until END.

        Returns the final state after all nodes have been processed.
        """
        current_nodes = self._get_next_nodes(START, state.data.model_dump())
        if not current_nodes:
            raise ValueError("No edges from START — the graph has no entry point")

        steps = 0
        visited: list[str] = []

        while current_nodes and steps < max_steps:
            steps += 1
            next_round: list[str] = []

            # Execute current nodes (potentially in parallel)
            if len(current_nodes) == 1:
                node_name = current_nodes[0]
                if node_name == END:
                    logger.info(f"Graph reached END after {steps} steps")
                    break
                await self._execute_node(node_name, state)
                visited.append(node_name)
                next_round = self._get_next_nodes(node_name, state.data.model_dump())
            else:
                # Parallel execution
                tasks = []
                for node_name in current_nodes:
                    if node_name == END:
                        continue
                    tasks.append(self._execute_node(node_name, state))
                    visited.append(node_name)

                await asyncio.gather(*tasks)

                # Collect all next nodes
                for node_name in current_nodes:
                    if node_name != END:
                        next_round.extend(
                            self._get_next_nodes(node_name, state.data.model_dump())
                        )

                if END in current_nodes:
                    break

            current_nodes = list(set(next_round))  # Deduplicate

        if steps >= max_steps:
            logger.warning(f"Graph execution hit max_steps ({max_steps})")

        logger.info(f"Graph execution complete. Path: {' → '.join(visited)}")
        return state

    async def _execute_node(self, name: str, state: SharedState[T]) -> None:
        """Execute a single node, updating the shared state."""
        node = self._nodes.get(name)
        if not node:
            raise KeyError(f"Node '{name}' not found in graph")

        if not node.func:
            logger.warning(f"Node '{name}' has no function, skipping")
            return

        logger.debug(f"Executing node '{name}'")

        try:
            state_dict = state.data.model_dump()
            updates = await node.func(state_dict)
            if updates and isinstance(updates, dict):
                state.update(description=f"Node '{name}' update", **updates)
        except Exception as exc:
            logger.error(f"Node '{name}' failed: {exc}")
            raise

    # ---- Visualization ----

    def to_mermaid(self) -> str:
        """Generate a Mermaid diagram of the graph."""
        lines = ["graph TD"]

        for edge in self._edges:
            src = edge.source if edge.source != START else "START([Start])"
            tgt = edge.target if edge.target != END else "END([End])"
            if edge.label:
                lines.append(f"    {src} -->|{edge.label}| {tgt}")
            else:
                lines.append(f"    {src} --> {tgt}")

        for source, (_, route_map) in self._conditional_edges.items():
            for label, target in route_map.items():
                tgt = target if target != END else "END([End])"
                lines.append(f"    {source} -->|{label}| {tgt}")

        return "\n".join(lines)
