"""DAG graph with topological sort and execution."""

from __future__ import annotations

from collections import defaultdict, deque
from typing import Any, Self

from thryve.core.graph.node import Node
from thryve.utils import get_logger

logger = get_logger("graph")


class Graph:
    """Directed Acyclic Graph (DAG) for workflow execution.

    Nodes are arbitrary :class:`Node` subclasses; edges encode data-flow
    dependencies.  Execution order is determined via Kahn's algorithm
    (topological sort).
    """

    def __init__(self, name: str = "default") -> None:
        self.name = name
        self.nodes: dict[str, Node] = {}
        self.edges: defaultdict[str, list[str]] = defaultdict(list)  # out-edges
        self._in_degree: dict[str, int] = {}

    # ------------------------------------------------------------------
    # Construction
    # ------------------------------------------------------------------

    def add_node(self, node: Node) -> Self:
        """Add a node to the graph."""
        if node.name in self.nodes:
            raise ValueError(f"Duplicate node name: {node.name!r}")
        self.nodes[node.name] = node
        self._in_degree.setdefault(node.name, 0)
        return self

    def add_edge(self, from_node: str, to_node: str) -> Self:
        """Add a directed edge from *from_node* to *to_node*."""
        if from_node not in self.nodes:
            raise ValueError(f"Source node not found: {from_node!r}")
        if to_node not in self.nodes:
            raise ValueError(f"Target node not found: {to_node!r}")
        self.edges[from_node].append(to_node)
        self._in_degree[to_node] = self._in_degree.get(to_node, 0) + 1
        return self

    def chain(self, *nodes: Node) -> Self:
        """Convenience: add nodes in sequence and link them linearly."""
        for i, node in enumerate(nodes):
            if node.name not in self.nodes:
                self.add_node(node)
            if i > 0:
                self.add_edge(nodes[i - 1].name, node.name)
        return self

    # ------------------------------------------------------------------
    # Topological sort (Kahn)
    # ------------------------------------------------------------------

    def topological_sort(self) -> list[str]:
        """Return node names in a valid execution order.

        Raises :class:`ValueError` if the graph contains a cycle.
        """
        in_deg = dict(self._in_degree)
        queue: deque[str] = deque(n for n, d in in_deg.items() if d == 0)
        result: list[str] = []

        while queue:
            node = queue.popleft()
            result.append(node)
            for neighbour in self.edges.get(node, []):
                in_deg[neighbour] -= 1
                if in_deg[neighbour] == 0:
                    queue.append(neighbour)

        if len(result) != len(self.nodes):
            raise ValueError("Graph contains a cycle – topological sort impossible")

        return result

    # ------------------------------------------------------------------
    # Layer decomposition (for parallel execution)
    # ------------------------------------------------------------------

    def get_execution_layers(self) -> list[list[str]]:
        """Partition nodes into layers where all nodes in a layer are independent.

        Nodes within the same layer can be executed concurrently.
        """
        in_deg = dict(self._in_degree)
        layers: list[list[str]] = []
        current_layer = [n for n, d in in_deg.items() if d == 0]

        while current_layer:
            layers.append(current_layer)
            next_layer: list[str] = []
            for node in current_layer:
                for neighbour in self.edges.get(node, []):
                    in_deg[neighbour] -= 1
                    if in_deg[neighbour] == 0:
                        next_layer.append(neighbour)
            current_layer = next_layer

        return layers

    # ------------------------------------------------------------------
    # Simple sequential execute (see GraphExecutor for parallel)
    # ------------------------------------------------------------------

    async def execute(self, initial_inputs: dict[str, Any] | None = None) -> dict[str, Any]:
        """Execute all nodes in topological order, passing outputs downstream."""
        order = self.topological_sort()
        node_outputs: dict[str, Any] = {}

        for node_name in order:
            node = self.nodes[node_name]
            inputs = self._collect_inputs(node_name, node_outputs, initial_inputs or {})
            logger.debug("Executing node %r", node_name)
            output = await node.execute(inputs)
            node_outputs[node_name] = output

        return node_outputs

    def _collect_inputs(
        self,
        node_name: str,
        node_outputs: dict[str, Any],
        initial_inputs: dict[str, Any],
    ) -> dict[str, Any]:
        """Gather inputs for *node_name* from upstream outputs + initial inputs."""
        inputs = dict(initial_inputs)
        # Add outputs from all upstream nodes.
        for src, targets in self.edges.items():
            if node_name in targets:
                inputs[src] = node_outputs.get(src)
        return inputs
