"""PathwayVM - The graph execution engine facade.

This is the main entry point for pathway execution. It delegates to
specialized schedulers based on execution mode:
- DAG execution: DAGScheduler
- Streaming execution: StreamingScheduler

ALL PYDANTIC — inputs/outputs are JsonObject (typed JSON), results are models.
"""

import asyncio
import logging
import time
from collections.abc import Callable

from pydantic import BaseModel

from core.common import JsonObject
from core.pathways.context import Context
from core.pathways.models.execution import PathwayEvent, PathwayStatus
from core.pathways.nodes.execution import PathwayRecord
from core.pathways.pathway import Pathway
from core.pathways.run_state import RunState
from core.pathways.vm.observability import VMEventEmitter, wrap_tools_for_observability
from core.pathways.vm.schedulers import DAGScheduler, StreamingScheduler

logger = logging.getLogger(__name__)


class VMStats(BaseModel):
    """Execution statistics from the VM."""

    total_executions: int = 0
    completed: int = 0
    failed: int = 0
    success_rate: float = 0.0


class PathwayVM:
    """The graph execution engine.

    Execute graphs by calling node.compute() directly.
    No resolvers, no type lookups - nodes know how to execute themselves.

    Usage:
        from core.pathways.context import Context
        ctx = Context(tools={...})
        vm = PathwayVM(ctx)
        record = await vm.execute(graph, {"input": "hello"})
    """

    def __init__(
        self,
        ctx: Context,
        *,
        max_parallel: int = 10,
        default_timeout: float = 300.0,
    ):
        """Create a PathwayVM.

        Args:
            ctx: Execution context with tools and services
            max_parallel: Maximum concurrent node executions
            default_timeout: Default execution timeout in seconds
        """
        self.ctx = ctx
        self.max_parallel = max_parallel
        self.default_timeout = default_timeout
        self._execution_semaphore = asyncio.Semaphore(max_parallel)
        self._node_semaphore = asyncio.Semaphore(max_parallel)
        self._executions: dict[str, PathwayRecord] = {}

        # Initialize event emitter
        self._emitter = VMEventEmitter()

        # Permanent trace collector — always on, partitions by execution_id
        from core.observability.trace import TraceCollector

        self._trace_collector = TraceCollector()
        self._emitter.add_listener(self._trace_collector)

        # Wrap tools for observability
        try:
            wrap_tools_for_observability(self.ctx, self._emitter)
        except Exception:
            logger.debug("Failed to wrap tools for observability (non-fatal)", exc_info=True)

        # Initialize schedulers
        self._dag_scheduler = DAGScheduler(ctx, max_parallel, self._emitter, self._node_semaphore)
        self._streaming_scheduler = StreamingScheduler(
            ctx,
            max_parallel,
            self._emitter,
            self._node_semaphore,
        )

    # -- Trace access (always-on) --------------------------------------------

    @property
    def last_trace(self) -> object | None:
        """The most recently completed ``ExecutionTrace``, or ``None``."""
        return self._trace_collector.last_trace

    def get_trace(self, execution_id: str) -> object | None:
        """Look up a completed ``ExecutionTrace`` by *execution_id*."""
        return self._trace_collector.get_trace(execution_id)

    def _inject_trace_metadata(self, record: PathwayRecord, execution_id: str) -> None:
        """Copy token counts from the trace into the PathwayRecord metadata.

        Called after each execution so that ``Agent.turn()`` can read
        ``record.metadata["total_tokens"]``.
        """
        trace = self._trace_collector.get_trace(execution_id)
        if trace is not None:
            record.metadata["total_tokens"] = trace.total_tokens

    # -- Listener management --------------------------------------------------

    def add_listener(self, listener: Callable[[PathwayEvent], None]) -> None:
        """Add an event listener."""
        self._emitter.add_listener(listener)

    def remove_listener(self, listener: Callable[[PathwayEvent], None]) -> None:
        """Remove a previously added event listener."""
        self._emitter.remove_listener(listener)

    async def execute(
        self,
        pathway: Pathway,
        inputs: JsonObject | None = None,
        *,
        execution_id: str | None = None,
        timeout: float | None = None,
        run_state: RunState | None = None,
        checkpoint_store: object | None = None,
    ) -> PathwayRecord:
        """Execute a pathway.

        Parameters
        ----------
        pathway
            The pathway to execute.
        inputs
            Input data for the pathway (typed JSON).
        execution_id
            Optional execution ID (auto-generated if omitted).
        timeout
            Execution timeout in seconds.
        run_state
            Per-execution mutable state.  Created fresh by the caller
            (usually the consumer's bootstrap code) for each turn.
        checkpoint_store
            Optional ``StateStore`` used to persist DAG checkpoints for
            resume-from-failure.
        """
        execution_id = execution_id or f"exec_{int(time.time() * 1000)}"
        timeout = timeout or self.default_timeout
        inputs = inputs or {}

        # Install per-execution state
        if run_state is not None:
            self.ctx.run = run_state
        else:
            from core.pathways.run_state import RunState

            self.ctx.run = RunState()

        # Update context identifiers
        self.ctx.pathway_id = pathway.id
        self.ctx.execution_id = execution_id

        # Set logging context so every log line during execution carries IDs
        from core.observability.logging import (
            reset_log_execution_id,
            reset_log_pathway_id,
            set_log_execution_id,
            set_log_pathway_id,
        )

        pw_token = set_log_pathway_id(pathway.id)
        ex_token = set_log_execution_id(execution_id)

        try:
            async with self._execution_semaphore:
                try:
                    _budget_max_tokens = int(inputs.get("budget_max_tokens", 0))
                    _budget_max_cost = float(inputs.get("budget_max_cost", 0.0))

                    result = await asyncio.wait_for(
                        self._dag_scheduler.execute(
                            pathway,
                            inputs,
                            execution_id,
                            checkpoint_store=checkpoint_store,
                            budget_max_tokens=_budget_max_tokens,
                            budget_max_cost=_budget_max_cost,
                        ),
                        timeout=timeout,
                    )
                    self._inject_trace_metadata(result, execution_id)
                    self._executions[execution_id] = result
                    return result

                except TimeoutError:
                    msg = f"Execution {execution_id} timed out after {timeout}s"
                    raise TimeoutError(msg) from None
        finally:
            reset_log_pathway_id(pw_token)
            reset_log_execution_id(ex_token)

    async def resume(
        self,
        pathway: Pathway,
        checkpoint: object,
        *,
        timeout: float | None = None,
        run_state: RunState | None = None,
        checkpoint_store: object | None = None,
    ) -> PathwayRecord:
        """Resume a DAG execution from a checkpoint.

        Skips already-completed nodes recorded in the checkpoint and
        continues execution from where it left off.

        Parameters
        ----------
        pathway
            The pathway to re-execute.
        checkpoint
            A ``DAGCheckpoint`` loaded from the state store.
        timeout
            Execution timeout in seconds.
        run_state
            Per-execution mutable state.
        checkpoint_store
            Optional store for persisting new checkpoints during resume.
        """
        from core.pathways.checkpoint import DAGCheckpoint

        if not isinstance(checkpoint, DAGCheckpoint):
            raise TypeError(f"Expected DAGCheckpoint, got {type(checkpoint).__name__}")

        execution_id = checkpoint.execution_id
        timeout = timeout or self.default_timeout
        inputs = checkpoint.inputs

        if run_state is not None:
            self.ctx.run = run_state
        else:
            self.ctx.run = RunState()

        self.ctx.pathway_id = pathway.id
        self.ctx.execution_id = execution_id

        async with self._execution_semaphore:
            try:
                result = await asyncio.wait_for(
                    self._dag_scheduler.execute(
                        pathway,
                        inputs,
                        execution_id,
                        checkpoint_store=checkpoint_store,
                        resume_checkpoint=checkpoint,
                    ),
                    timeout=timeout,
                )
                self._inject_trace_metadata(result, execution_id)
                self._executions[execution_id] = result
                return result
            except TimeoutError:
                msg = f"Resume {execution_id} timed out after {timeout}s"
                raise TimeoutError(msg) from None

    async def stream_execute(
        self,
        pathway: Pathway,
        inputs: JsonObject | None = None,
        *,
        max_events: int | None = None,
    ):
        """Execute a streaming pathway and yield per-event tick results.

        This is **true streaming execution**:
        - Each streaming node yields events over time
        - For each event, the VM executes the downstream reachable subgraph

        Yields JsonObject dicts like:
            {
              "type": "tick",
              "source_node_id": "...",
              "event": {...},
              "outputs": {node_id: node_output, ...},
            }
        """
        async for result in self._streaming_scheduler.execute(pathway, inputs or {}, max_events):
            yield result

    async def get_execution(self, execution_id: str) -> PathwayRecord | None:
        """Get execution record by ID."""
        return self._executions.get(execution_id)

    def get_stats(self) -> VMStats:
        """Get execution statistics as a Pydantic model."""
        total = len(self._executions)
        completed = sum(1 for r in self._executions.values() if r.status == PathwayStatus.COMPLETED)
        failed = sum(1 for r in self._executions.values() if r.status == PathwayStatus.FAILED)

        return VMStats(
            total_executions=total,
            completed=completed,
            failed=failed,
            success_rate=completed / total if total > 0 else 0,
        )
