"""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
- Stateful execution: StatefulScheduler

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.session import PathwaySession
from core.pathways.vm.observability import VMEventEmitter, wrap_tools_for_observability
from core.pathways.vm.schedulers import DAGScheduler, StatefulScheduler, 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()

        # 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,
        )
        self._stateful_scheduler = StatefulScheduler(
            ctx,
            max_parallel,
            self._emitter,
            self._node_semaphore,
        )

    def add_listener(self, listener: Callable[[PathwayEvent], None]) -> None:
        """Add an event listener."""
        self._emitter.add_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

        async with self._execution_semaphore:
            try:
                from core.pathways.pathway import PathwayMode

                if pathway.mode == PathwayMode.STATEFUL:
                    # Stateful pathways need a session for multi-turn state.
                    # Auto-create one so callers don't have to wire it manually.
                    session = PathwaySession(pathway_id=pathway.id)
                    turn_result = await asyncio.wait_for(
                        self._stateful_scheduler.execute_turn(pathway, session, inputs, timeout),
                        timeout=timeout,
                    )
                    # Convert stateful result into a PathwayRecord so the
                    # return type stays consistent for Agent.turn().
                    outputs = turn_result if isinstance(turn_result, dict) else {}
                    has_error = bool(outputs.get("error"))
                    record = PathwayRecord(
                        pathway_id=pathway.id,
                        execution_id=execution_id,
                        status=(PathwayStatus.FAILED if has_error else PathwayStatus.COMPLETED),
                        outputs=outputs,
                    )
                    record.metadata["_stateful_session_id"] = session.id
                    record.metadata["_stateful_status"] = session.status.value
                    self._executions[execution_id] = record
                    return record

                # Extract budget from inputs if passed by Agent.turn()
                _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._executions[execution_id] = result
                return result

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

    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._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 execute_stateful(
        self,
        pathway: Pathway,
        session: PathwaySession,
        inputs: JsonObject | None = None,
        *,
        timeout: float | None = None,
    ) -> JsonObject:
        """Execute one turn of a stateful pathway.

        Unlike execute() which runs the full DAG, this:
        - Runs only the current stage
        - Evaluates transitions to determine next stage
        - Updates session state
        - Supports cycles and multi-turn pathways

        Args:
            pathway: The stateful pathway to execute
            session: Session tracking state across turns
            inputs: Input data for this turn (typed JSON)
            timeout: Execution timeout in seconds

        Returns
        -------
            JsonObject with stage result, next_stage, and session status
        """
        return await self._stateful_scheduler.execute_turn(pathway, session, inputs or {}, timeout)

    async def run_stateful_to_completion(
        self,
        pathway: Pathway,
        session: PathwaySession,
        initial_inputs: JsonObject | None = None,
        *,
        max_turns: int = 50,
        turn_inputs_fn: Callable[..., JsonObject] | None = None,
    ) -> JsonObject:
        """Run a stateful pathway to completion.

        Keeps executing turns until:
        - Pathway completes
        - Session is paused (awaiting input/approval)
        - Max turns reached
        - Error occurs

        Args:
            pathway: The stateful pathway
            session: Session tracking state
            initial_inputs: Inputs for first turn (typed JSON)
            max_turns: Safety limit on turns
            turn_inputs_fn: Optional function to provide inputs for each turn

        Returns
        -------
            Final result as JsonObject with session state
        """
        return await self._stateful_scheduler.run_to_completion(
            pathway,
            session,
            initial_inputs or {},
            max_turns,
            turn_inputs_fn,
        )

    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,
        )
