"""Graph executor with layer-parallel execution support."""

from __future__ import annotations

import asyncio
from typing import Any

from thryve.core.graph.graph import Graph
from thryve.core.graph.models import NodeResult, NodeStatus
from thryve.utils import get_logger

logger = get_logger("graph.executor")


class GraphExecutor:
    """Execute a :class:`Graph` with concurrent execution of independent nodes.

    Nodes within the same layer (i.e., no mutual dependencies) run in
    parallel via ``asyncio.gather``.
    """

    def __init__(self, max_workers: int = 4) -> None:
        self._max_workers = max_workers

    async def execute(
        self,
        graph: Graph,
        initial_inputs: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Execute *graph* layer by layer.

        Returns a dict mapping node names to their outputs.
        """
        layers = graph.get_execution_layers()
        node_outputs: dict[str, Any] = {}
        init = initial_inputs or {}

        for layer in layers:
            # Limit concurrency via semaphore.
            sem = asyncio.Semaphore(self._max_workers)
            tasks = [
                self._run_node(graph, name, node_outputs, init, sem)
                for name in layer
            ]
            results: list[NodeResult] = await asyncio.gather(*tasks)

            for res in results:
                if res.status == NodeStatus.COMPLETED:
                    node_outputs[res.node_name] = res.output
                elif res.status == NodeStatus.FAILED:
                    logger.error(
                        "Node %r failed: %s", res.node_name, res.error,
                    )
                    # Propagate failure – downstream nodes won't run because
                    # their inputs are missing, which is acceptable.

        return node_outputs

    async def execute_with_recovery(
        self,
        graph: Graph,
        initial_inputs: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Execute with best-effort recovery: failed non-critical nodes are skipped."""
        layers = graph.get_execution_layers()
        node_outputs: dict[str, Any] = {}
        init = initial_inputs or {}
        failed_nodes: list[str] = []

        for layer in layers:
            sem = asyncio.Semaphore(self._max_workers)
            tasks = [
                self._run_node(graph, name, node_outputs, init, sem)
                for name in layer
            ]
            results = await asyncio.gather(*tasks)

            for res in results:
                if res.status == NodeStatus.COMPLETED:
                    node_outputs[res.node_name] = res.output
                else:
                    failed_nodes.append(res.node_name)

        if failed_nodes:
            logger.warning("Nodes that failed: %s", failed_nodes)

        return node_outputs

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    async def _run_node(
        self,
        graph: Graph,
        node_name: str,
        node_outputs: dict[str, Any],
        initial_inputs: dict[str, Any],
        sem: asyncio.Semaphore,
    ) -> NodeResult:
        async with sem:
            node = graph.nodes[node_name]
            inputs = graph._collect_inputs(node_name, node_outputs, initial_inputs)
            try:
                output = await node.execute(inputs)
                return NodeResult(
                    node_name=node_name,
                    output=output,
                    status=NodeStatus.COMPLETED,
                )
            except Exception as exc:
                logger.exception("Node %r raised an error", node_name)
                return NodeResult(
                    node_name=node_name,
                    status=NodeStatus.FAILED,
                    error=f"{type(exc).__name__}: {exc}",
                )
