"""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.approval import ApprovalRequiredError
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,
        *,
        checkpoint_store: object | None = None,
        resume_checkpoint: object | None = None,
        budget_max_tokens: int = 0,
        budget_max_cost: float = 0.0,
    ) -> 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()
            skipped: set[str] = set()
            in_flight: dict[
                str,
                asyncio.Task[
                    tuple[str, NodeOutputs, bool, float, str | None, str | None, NodeInputs]
                ],
            ] = {}

            # Resume from checkpoint: pre-populate completed nodes
            if resume_checkpoint is not None:
                from core.pathways.checkpoint import DAGCheckpoint

                if isinstance(resume_checkpoint, DAGCheckpoint):
                    for nid, out_data in resume_checkpoint.completed_nodes.items():
                        results[nid] = NodeOutputs.model_validate(out_data)
                        completed.add(nid)
                    record.metrics.nodes_executed = len(completed)
                    record.metrics.nodes_succeeded = len(completed)

            ready: list[str] = [
                nid
                for nid in topo_order
                if nid not in completed and deps.get(nid, set()).issubset(completed)
            ]
            stop_on_error = bool(inputs.get("_stop_on_error"))
            budget_exceeded = False
            cumulative_tokens = 0
            cumulative_cost = 0.0
            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 ApprovalRequiredError:
                            # Let ApprovalRequiredError propagate — handled by the scheduler loop
                            raise

                        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, skipped)
                    if not should:
                        record.metrics.nodes_skipped += 1
                        completed.add(node_id)
                        skipped.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:
                    # Check for ApprovalRequiredError before accessing result
                    exc = task.exception() if not task.cancelled() else None
                    if isinstance(exc, ApprovalRequiredError):
                        approval_node_id = None
                        for nid, t in list(in_flight.items()):
                            if t is task:
                                approval_node_id = nid
                                break
                        in_flight.pop(approval_node_id, None)

                        # Save checkpoint and set awaiting_approval
                        if checkpoint_store is not None:
                            from core.pathways.checkpoint import DAGCheckpoint

                            ckpt = DAGCheckpoint(
                                execution_id=execution_id,
                                pathway_id=pathway.id,
                                completed_nodes={
                                    nid: res.to_dict() for nid, res in results.items()
                                },
                                pending_approval_node=approval_node_id,
                                inputs=inputs,
                            )
                            try:
                                await checkpoint_store.save_checkpoint(ckpt)
                            except Exception:
                                logger.debug("Failed to save approval checkpoint", exc_info=True)

                        # Cancel remaining in-flight tasks
                        for t in in_flight.values():
                            t.cancel()
                        await asyncio.gather(*in_flight.values(), return_exceptions=True)
                        in_flight.clear()
                        ready.clear()

                        record.status = PathwayStatus.AWAITING_APPROVAL
                        record.metadata["approval_node"] = approval_node_id
                        record.metadata["approval_message"] = exc.approval_message
                        record.metadata["approval_roles"] = exc.roles
                        record.outputs = {k: v.to_dict() for k, v in results.items()}
                        return record

                    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

                    # Accumulate token/cost from LLM node outputs for budget tracking
                    out_data = out.to_dict()
                    llm_output = out_data.get("output")
                    if isinstance(llm_output, dict):
                        usage = llm_output.get("usage")
                        if isinstance(usage, dict):
                            cumulative_tokens += int(usage.get("tokens_used", 0))
                        cumulative_cost += float(llm_output.get("cost_usd", 0.0))
                        # Also check metadata (iterative tool calling stores totals there)
                        meta = llm_output.get("metadata")
                        if isinstance(meta, dict):
                            cumulative_tokens += int(meta.get("total_tokens", 0))
                            cumulative_cost += float(meta.get("total_cost_usd", 0.0))

                    # Pathway-level budget enforcement
                    if (
                        budget_max_tokens > 0
                        and cumulative_tokens >= budget_max_tokens
                        and not budget_exceeded
                    ):
                        budget_exceeded = True
                        logger.warning(
                            "[DAG] Token budget exceeded: %d/%d — aborting remaining nodes",
                            cumulative_tokens,
                            budget_max_tokens,
                        )
                        for t in in_flight.values():
                            t.cancel()
                        await asyncio.gather(*in_flight.values(), return_exceptions=True)
                        in_flight.clear()
                        ready.clear()
                        record.metadata["budget_exceeded"] = True
                        record.metadata["budget_reason"] = (
                            f"Token budget: {cumulative_tokens}/{budget_max_tokens}"
                        )
                    elif (
                        budget_max_cost > 0
                        and cumulative_cost >= budget_max_cost
                        and not budget_exceeded
                    ):
                        budget_exceeded = True
                        logger.warning(
                            "[DAG] Cost budget exceeded: $%.4f/$%.4f — aborting remaining nodes",
                            cumulative_cost,
                            budget_max_cost,
                        )
                        for t in in_flight.values():
                            t.cancel()
                        await asyncio.gather(*in_flight.values(), return_exceptions=True)
                        in_flight.clear()
                        ready.clear()
                        record.metadata["budget_exceeded"] = True
                        record.metadata["budget_reason"] = (
                            f"Cost budget: ${cumulative_cost:.4f}/${budget_max_cost:.4f}"
                        )

                    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)

                    # Save checkpoint after each node (for workflow resume)
                    if checkpoint_store is not None:
                        from core.pathways.checkpoint import DAGCheckpoint

                        ckpt = DAGCheckpoint(
                            execution_id=execution_id,
                            pathway_id=pathway.id,
                            completed_nodes={nid: res.to_dict() for nid, res in results.items()},
                            failed_node=node_id if not ok else None,
                            inputs=inputs,
                        )
                        try:
                            await checkpoint_store.save_checkpoint(ckpt)
                        except Exception:
                            logger.debug(
                                "Failed to save checkpoint for %s (non-fatal)",
                                execution_id,
                                exc_info=True,
                            )

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

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

            # NOTE: Loop iteration is now handled by LoopNode (a proper graph node)
            # rather than the legacy pathway.loops list.

            # 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
            elif output_node_id is None:
                # No OutputNode — auto-promote the last LLM node's response
                # so Agent.turn() always finds a top-level "response" key.
                for nid in reversed(list(results)):
                    node_out = results[nid].to_dict()
                    llm_out = node_out.get("output")
                    if isinstance(llm_out, dict) and "response" in llm_out:
                        resp = llm_out["response"]
                        if isinstance(resp, str) and resp.strip():
                            final_outputs["response"] = resp
                            break

            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
