"""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.
"""

from __future__ import annotations

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

from pydantic import BaseModel

from core.common import JsonObject
from core.pathways.models.execution import PathwayEvent, PathwayStatus
from core.pathways.nodes.execution import PathwayRecord
from core.pathways.vm.observability import VMEventEmitter, wrap_tools_for_observability
from core.pathways.vm.schedulers import DAGScheduler, StatefulScheduler, StreamingScheduler

if TYPE_CHECKING:
    from core.pathways.context import Context
    from core.pathways.pathway import Pathway
    from core.pathways.run_state import RunState
    from core.pathways.session import PathwaySession

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,
    ) -> 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 ``Runtime``) for each turn.
        """
        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  # noqa: PLC0415

            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:
                result = await asyncio.wait_for(
                    self._dag_scheduler.execute(pathway, inputs, execution_id),
                    timeout=timeout,
                )
                self._executions[execution_id] = result
                return result

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

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