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

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

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

from core.common import JsonObject
from core.pathways.context import Context
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.observability import VMEventEmitter
from core.pathways.vm.sanitizer import sanitize_for_logging
from core.pathways.vm.utils import gather_inputs, should_execute

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

                            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:
                            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:
                                    self.ctx.run._current_node_id = None

                    # 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 (iterative refinement)
            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
                                    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:
                                    self.ctx.run._current_node_id = None
                                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
