"""DAG Scheduler — parallel execution of directed acyclic graphs.

Operates on NodeInputs / NodeOutputs Pydantic models end-to-end.
"""

from __future__ import annotations

import asyncio
import contextlib
import logging
import time
import traceback
from datetime import UTC, datetime
from typing import TYPE_CHECKING

from core.common import JsonObject
from core.pathways.models.io import NodeInputs, NodeOutputs
from core.pathways.nodes.execution import (
    PathwayEvent,
    PathwayRecord,
    PathwayStatus,
)
from core.pathways.pathway import PSEUDO_NODE_IDS, Pathway
from core.pathways.vm.graph import topological_sort
from core.pathways.vm.sanitizer import sanitize_for_logging
from core.pathways.vm.utils import gather_inputs, should_execute

if TYPE_CHECKING:
    from core.pathways.context import Context
    from core.pathways.vm.observability import VMEventEmitter

logger = logging.getLogger(__name__)


class DAGScheduler:
    """Scheduler for DAG (directed acyclic graph) execution."""

    def __init__(
        self,
        ctx: Context,
        max_parallel: int,
        emitter: VMEventEmitter,
        node_semaphore: asyncio.Semaphore,
    ) -> None:
        self.ctx = ctx
        self.max_parallel = max_parallel
        self.emitter = emitter
        self._node_semaphore = node_semaphore

    async def execute(
        self,
        pathway: Pathway,
        inputs: JsonObject,
        execution_id: str,
    ) -> PathwayRecord:
        """Execute a pathway using DAG scheduling.

        All node interactions use NodeInputs / NodeOutputs Pydantic models.
        """
        record = PathwayRecord(
            id=execution_id,
            pathway_id=pathway.id,
            inputs=inputs,
            status=PathwayStatus.RUNNING,
        )
        record.metrics.start_time = datetime.now(UTC)
        run_span_id = f"run:{execution_id}"

        self.emitter.emit_event(
            PathwayEvent(
                event_type="execution_started",
                pathway_id=pathway.id,
                data={
                    "execution_id": execution_id,
                    "span_id": run_span_id,
                    "parent_span_id": None,
                    "inputs": sanitize_for_logging(inputs),
                },
            )
        )

        try:
            topo_order = topological_sort(pathway)
            order_index = {nid: i for i, nid in enumerate(topo_order)}
            results: dict[str, NodeOutputs] = {}

            # Build dependency graph
            deps: dict[str, set[str]] = {nid: set() for nid in pathway.nodes}
            dependents: dict[str, set[str]] = {nid: set() for nid in pathway.nodes}
            for conn in pathway.connections:
                from_is_pseudo = (
                    conn.from_node in PSEUDO_NODE_IDS and conn.from_node not in pathway.nodes
                )
                to_is_pseudo = conn.to_node in PSEUDO_NODE_IDS and conn.to_node not in pathway.nodes
                if from_is_pseudo or to_is_pseudo:
                    continue
                if conn.from_node in deps and conn.to_node in deps:
                    deps[conn.to_node].add(conn.from_node)
                    dependents[conn.from_node].add(conn.to_node)

            completed: set[str] = set()
            in_flight: dict[str, asyncio.Task[tuple[str, NodeOutputs, bool, float, str | None, str | None, NodeInputs]]] = {}
            ready: list[str] = [nid for nid in topo_order if not deps.get(nid)]
            stop_on_error = bool(inputs.get("_stop_on_error"))
            lock = asyncio.Lock()

            async def _run_node(
                node_id: str,
            ) -> tuple[str, NodeOutputs, bool, float, str | None, str | None, NodeInputs]:
                """Run a single node under the global node semaphore.

                Respects per-node ``timeout_ms`` and ``retry_count`` / ``retry_delay_ms``
                fields on NodeBase.  Retries use exponential back-off capped at the
                node-level timeout.
                """
                node = pathway.nodes.get(node_id)
                empty_inputs = NodeInputs()
                if node is None:
                    return (node_id, NodeOutputs(), True, 0.0, None, None, empty_inputs)

                async with lock:
                    node_inputs = gather_inputs(pathway, node_id, results, inputs)

                node_span_id = f"node:{execution_id}:{node_id}"
                node_type = getattr(node, "type", node.__class__.__name__)
                node_timeout_s: float | None = (
                    node.timeout_ms / 1000.0 if node.timeout_ms else None
                )
                max_attempts = max(1, node.retry_count + 1)

                async with self._node_semaphore:
                    self.emitter.emit_event(
                        PathwayEvent(
                            event_type="node_started",
                            node_id=node_id,
                            pathway_id=pathway.id,
                            data={
                                "execution_id": execution_id,
                                "span_id": node_span_id,
                                "parent_span_id": run_span_id,
                                "node_type": node_type,
                                "inputs": sanitize_for_logging(node_inputs.to_dict()),
                            },
                        )
                    )

                    last_exc: Exception | None = None
                    started_at = time.time()

                    for attempt in range(max_attempts):
                        try:
                            with contextlib.suppress(Exception):
                                self.ctx.run._current_node_id = node_id  # noqa: SLF001

                            if node_timeout_s is not None:
                                out = await asyncio.wait_for(
                                    node.compute(node_inputs, self.ctx),
                                    timeout=node_timeout_s,
                                )
                            else:
                                out = await node.compute(node_inputs, self.ctx)

                            duration_ms = (time.time() - started_at) * 1000
                            return (node_id, out, True, duration_ms, None, None, node_inputs)

                        except TimeoutError:
                            duration_ms = (time.time() - started_at) * 1000
                            tb = traceback.format_exc()[-2000:]  # truncate
                            self.emitter.emit_event(
                                PathwayEvent(
                                    event_type="node_timeout",
                                    node_id=node_id,
                                    pathway_id=pathway.id,
                                    error=f"Node timed out after {node.timeout_ms}ms",
                                    data={
                                        "execution_id": execution_id,
                                        "timeout_ms": node.timeout_ms,
                                        "duration_ms": duration_ms,
                                    },
                                )
                            )
                            return (
                                node_id,
                                NodeOutputs(
                                    error=f"Timeout after {node.timeout_ms}ms",
                                    error_type="TimeoutError",
                                    error_traceback=tb,
                                    error_node_id=node_id,
                                ),
                                False,
                                duration_ms,
                                f"Timeout after {node.timeout_ms}ms",
                                "TimeoutError",
                                node_inputs,
                            )

                        except Exception as e:  # noqa: BLE001
                            last_exc = e
                            if attempt < max_attempts - 1:
                                delay_s = (node.retry_delay_ms / 1000.0) * (2 ** attempt)
                                self.emitter.emit_event(
                                    PathwayEvent(
                                        event_type="node_retry",
                                        node_id=node_id,
                                        pathway_id=pathway.id,
                                        error=str(e),
                                        data={
                                            "execution_id": execution_id,
                                            "attempt": attempt + 1,
                                            "max_attempts": max_attempts,
                                            "retry_delay_ms": delay_s * 1000,
                                        },
                                    )
                                )
                                await asyncio.sleep(delay_s)

                        finally:
                            with contextlib.suppress(Exception):
                                if self.ctx.run._current_node_id == node_id:  # noqa: SLF001
                                    self.ctx.run._current_node_id = None  # noqa: SLF001

                    # All retries exhausted
                    duration_ms = (time.time() - started_at) * 1000
                    err = last_exc or RuntimeError("Unknown error")
                    tb = traceback.format_exc()[-2000:]  # truncate
                    return (
                        node_id,
                        NodeOutputs(
                            error=str(err),
                            error_type=type(err).__name__,
                            error_traceback=tb,
                            error_node_id=node_id,
                        ),
                        False,
                        duration_ms,
                        str(err),
                        type(err).__name__,
                        node_inputs,
                    )

            def _mark_ready_if_unblocked(nid: str) -> None:
                if nid in completed or nid in in_flight:
                    return
                if deps.get(nid, set()).issubset(completed):
                    ready.append(nid)

            # Scheduler loop: execute independent nodes concurrently
            while ready or in_flight:
                ready.sort(key=lambda nid: order_index.get(nid, 10**9))
                while ready and len(in_flight) < self.max_parallel:
                    node_id = ready.pop(0)

                    async with lock:
                        should = should_execute(pathway, node_id, results)
                    if not should:
                        record.metrics.nodes_skipped += 1
                        completed.add(node_id)
                        for dep in dependents.get(node_id, set()):
                            _mark_ready_if_unblocked(dep)
                        continue

                    in_flight[node_id] = asyncio.create_task(_run_node(node_id))

                if not in_flight:
                    break

                done, _pending = await asyncio.wait(
                    in_flight.values(), return_when=asyncio.FIRST_COMPLETED
                )
                for task in done:
                    node_id, out, ok, duration_ms, error_msg, error_type, node_inputs = (
                        task.result()
                    )
                    in_flight.pop(node_id, None)

                    async with lock:
                        results[node_id] = out
                        record.metrics.nodes_executed += 1
                        if ok:
                            record.metrics.nodes_succeeded += 1
                        else:
                            record.metrics.nodes_failed += 1

                    node = pathway.nodes.get(node_id)
                    node_type = (
                        getattr(node, "type", node.__class__.__name__) if node else "unknown"
                    )
                    node_span_id = f"node:{execution_id}:{node_id}"

                    if ok:
                        self.emitter.emit_event(
                            PathwayEvent(
                                event_type="node_completed",
                                node_id=node_id,
                                pathway_id=pathway.id,
                                data={
                                    "execution_id": execution_id,
                                    "span_id": node_span_id,
                                    "parent_span_id": run_span_id,
                                    "success": True,
                                    "duration_ms": duration_ms,
                                    "node_type": node_type,
                                    "inputs": sanitize_for_logging(node_inputs.to_dict()),
                                    "outputs": sanitize_for_logging(out.model_dump()),
                                },
                            )
                        )
                    else:
                        logger.error(
                            "Node execution failed: node_id=%s, type=%s, error=%s, execution_id=%s, duration_ms=%.2f",
                            node_id,
                            error_type,
                            error_msg,
                            execution_id,
                            duration_ms,
                        )
                        self.emitter.emit_event(
                            PathwayEvent(
                                event_type="node_failed",
                                node_id=node_id,
                                pathway_id=pathway.id,
                                error=error_msg or "",
                                data={
                                    "execution_id": execution_id,
                                    "span_id": node_span_id,
                                    "parent_span_id": run_span_id,
                                    "duration_ms": duration_ms,
                                    "error_type": error_type,
                                    "node_type": node_type,
                                    "inputs": sanitize_for_logging(node_inputs.to_dict()),
                                    "outputs": sanitize_for_logging(out.model_dump()),
                                },
                            )
                        )

                        if stop_on_error:
                            record.status = PathwayStatus.FAILED
                            record.error = f"{node_id} ({error_type}): {error_msg}"
                            for t in in_flight.values():
                                t.cancel()
                            await asyncio.gather(*in_flight.values(), return_exceptions=True)
                            in_flight.clear()
                            ready.clear()
                            break

                    completed.add(node_id)
                    for dep in dependents.get(node_id, set()):
                        _mark_ready_if_unblocked(dep)

                if record.status == PathwayStatus.FAILED and stop_on_error:
                    break

            # Execute loops (ReAct pattern)
            if record.status != PathwayStatus.FAILED:
                for loop in pathway.loops:
                    iteration = 0
                    while loop.should_continue(results, iteration):
                        self.emitter.emit(
                            "loop_iteration",
                            pathway_id=pathway.id,
                            execution_id=execution_id,
                            iteration=iteration,
                            loop_nodes=loop.body_nodes,
                        )

                        for node_id in loop.body_nodes:
                            node = pathway.nodes.get(node_id)
                            if node is None:
                                continue

                            loop_inputs_data = gather_inputs(
                                pathway, node_id, results, {**inputs, "_iteration": iteration}
                            )

                            started_at = time.time()
                            try:
                                async with self._node_semaphore:
                                    self.ctx.run._current_node_id = node_id  # noqa: SLF001
                                    result = await node.compute(loop_inputs_data, self.ctx)
                                duration_ms = (time.time() - started_at) * 1000

                                results[node_id] = result
                                record.metrics.nodes_executed += 1
                                record.metrics.nodes_succeeded += 1

                                self.emitter.emit(
                                    "node_completed",
                                    node_id=node_id,
                                    execution_id=execution_id,
                                    success=True,
                                    duration_ms=duration_ms,
                                    iteration=iteration,
                                )
                            except Exception as e:
                                duration_ms = (time.time() - started_at) * 1000
                                results[node_id] = NodeOutputs(error=str(e))
                                record.metrics.nodes_executed += 1
                                record.metrics.nodes_failed += 1

                                self.emitter.emit(
                                    "node_failed",
                                    node_id=node_id,
                                    error=str(e),
                                    execution_id=execution_id,
                                    duration_ms=duration_ms,
                                    iteration=iteration,
                                )
                            finally:
                                if self.ctx.run._current_node_id == node_id:  # noqa: SLF001
                                    self.ctx.run._current_node_id = None  # noqa: SLF001
                                break

                        iteration += 1

                    record.metadata["_loop_iterations"] = iteration

            # Extract outputs: prefer OutputNode result, fallback to all results.
            # Serialise to dicts here — this is the API boundary where typed
            # Python objects must become JSON-compatible.
            final_outputs: JsonObject = {
                k: v.to_dict() for k, v in results.items()
            }

            # Look for an OutputNode and promote its mapped values to top-level
            output_node_id = None
            for node_id, node in pathway.nodes.items():
                if getattr(node, "type", None) == "output":
                    output_node_id = node_id
                    break

            if output_node_id and output_node_id in results:
                output_data = results[output_node_id].to_dict()
                for key, value in output_data.items():
                    if key not in final_outputs:
                        final_outputs[key] = value

            record.outputs = final_outputs

            record.metadata["_diagnostics"] = {
                "node_count": len(pathway.nodes),
                "nodes_executed": record.metrics.nodes_executed,
                "nodes_succeeded": record.metrics.nodes_succeeded,
                "nodes_failed": record.metrics.nodes_failed,
                "nodes_skipped": record.metrics.nodes_skipped,
                "output_node_id": output_node_id,
                "output_keys": list(final_outputs.keys()) if final_outputs else [],
                "parse_errors": pathway.metadata.get("_parse_errors", []),
            }

            if len(pathway.nodes) == 0:
                logger.warning(
                    "Pathway '%s' executed with 0 nodes - likely a loading/parsing issue. "
                    "Check pathway YAML structure and node type definitions.",
                    pathway.id,
                )
                record.metadata["_diagnostics"]["warning"] = "Empty pathway executed (0 nodes)"

            if record.status != PathwayStatus.FAILED:
                record.status = PathwayStatus.COMPLETED

        except Exception as e:
            record.status = PathwayStatus.FAILED
            record.error = str(e)
            self.emitter.emit(
                "execution_error",
                pathway_id=pathway.id,
                error=str(e),
                execution_id=execution_id,
            )

        finally:
            record.metrics.end_time = datetime.now(UTC)
            if record.metrics.start_time:
                duration = (record.metrics.end_time - record.metrics.start_time).total_seconds()
                record.metrics.duration_ms = duration * 1000

            self.emitter.emit(
                "execution_completed",
                pathway_id=pathway.id,
                execution_id=execution_id,
                span_id=run_span_id,
                status=record.status.value,
                success=record.status == PathwayStatus.COMPLETED,
                duration_ms=record.metrics.duration_ms,
                outputs=sanitize_for_logging(record.outputs),
                error=record.error,
                metrics={
                    "nodes_executed": record.metrics.nodes_executed,
                    "nodes_succeeded": record.metrics.nodes_succeeded,
                    "nodes_failed": record.metrics.nodes_failed,
                    "nodes_skipped": record.metrics.nodes_skipped,
                },
            )

        return record
