"""Stateful Scheduler - Multi-turn execution with persistent state.

This scheduler handles stateful pathways that execute one stage at a time,
maintaining session state across multiple turns.
"""

import asyncio
import time
from collections.abc import Callable

from core.common import JsonObject
from core.pathways.context import Context
from core.pathways.models.io import NodeInputs
from core.pathways.nodes.checkpoint import CheckpointNode
from core.pathways.nodes.stage import StageNode
from core.pathways.pathway import Pathway, PathwayMode
from core.pathways.session import PathwaySession, SessionCheckpoint
from core.pathways.vm.observability import VMEventEmitter


class StatefulScheduler:
    """Scheduler for stateful multi-turn execution."""

    def __init__(
        self,
        ctx: Context,
        max_parallel: int,
        emitter: VMEventEmitter,
        node_semaphore: asyncio.Semaphore,
    ) -> None:
        """Initialize stateful scheduler.

        Args:
            ctx: Execution context
            max_parallel: Maximum concurrent node executions
            emitter: Event emitter for observability
            node_semaphore: Semaphore for node execution concurrency
        """
        self.ctx = ctx
        self.max_parallel = max_parallel
        self.emitter = emitter
        self._node_semaphore = node_semaphore
        self.default_timeout = 300.0  # Default timeout for stateful execution

    async def execute_turn(
        self,
        pathway: Pathway,
        session: PathwaySession,
        inputs: JsonObject,
        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
            timeout: Execution timeout in seconds

        Returns
        -------
            Dict with stage result, next_stage, and session status
        """
        # Validate pathway mode
        if pathway.mode != PathwayMode.STATEFUL:
            msg = f"execute_stateful requires mode=stateful, got mode={pathway.mode}"
            raise ValueError(msg)

        # Check session is active
        if not session.is_active:
            return {
                "success": False,
                "error": f"Session is not active (status={session.status})",
                "session_id": session.id,
            }

        timeout = timeout or self.default_timeout

        # Increment turn count
        turn_number = session.increment_turn()

        # Get current stage
        current_stage_id = session.current_stage
        if not current_stage_id:
            # Find initial stage (first StageNode or first node)
            for node_id, node in pathway.nodes.items():
                if isinstance(node, StageNode):
                    current_stage_id = node_id
                    session.transition_to(node_id, "initial")
                    break

            if not current_stage_id:
                # No StageNodes - use first node
                current_stage_id = next(iter(pathway.nodes.keys()), None)
                if current_stage_id:
                    session.transition_to(current_stage_id, "initial")

        if not current_stage_id or current_stage_id not in pathway.nodes:
            session.fail(f"Stage not found: {current_stage_id}")
            return {
                "success": False,
                "error": f"Stage not found: {current_stage_id}",
                "session_id": session.id,
            }

        # Get the current stage node
        current_node = pathway.nodes[current_stage_id]

        # Track node visit count (for retry limits)
        # Increments on node ENTRY (before execution), so < 3 means max 3 total runs
        # Example: 1 initial run + 2 retries = 3 total visits
        session.increment_node_visit(current_stage_id)

        # Merge session context with inputs as NodeInputs
        execution_inputs = NodeInputs(
            **session.context,
            **inputs,
            _turn=turn_number,
            node_visit_counts=session.node_visit_counts,
        )

        # Emit turn started event
        self.emitter.emit(
            "stateful_turn_started",
            pathway_id=pathway.id,
            session_id=session.id,
            stage_id=current_stage_id,
            turn_number=turn_number,
        )

        try:
            # Execute the current node
            async with self._node_semaphore:
                self.ctx.run._current_node_id = current_stage_id
                self.ctx.run._session = session

                started_at = time.time()
                result = await current_node.compute(execution_inputs, self.ctx)
                duration_ms = (time.time() - started_at) * 1000

            # Update session context with results (serialise at session boundary)
            result_data = result.to_dict()
            session.update_context({current_stage_id: result_data})

            # Handle checkpoint nodes
            if isinstance(current_node, CheckpointNode):
                checkpoint = result.get("checkpoint")
                if checkpoint:
                    session.add_checkpoint(SessionCheckpoint(**checkpoint))

                if result.get("pause_execution"):
                    session.pause()
                    output_data = result.get("output", {})
                    return {
                        "success": True,
                        "paused": True,
                        "checkpoint_id": output_data.get("checkpoint_id")
                        if isinstance(output_data, dict)
                        else None,
                        "require_approval": output_data.get("require_approval")
                        if isinstance(output_data, dict)
                        else None,
                        "session_id": session.id,
                        "stage_id": current_stage_id,
                        "turn_number": turn_number,
                    }

            # Determine next stage from node result or conditional connections
            output_field = result.get("output", {})
            next_stage = result.get("selected_route") or (
                output_field.get("next_stage") if isinstance(output_field, dict) else None
            )
            transition_reason = (
                output_field.get("transition_reason") if isinstance(output_field, dict) else None
            )

            # If no explicit next stage, check conditional connections
            if not next_stage:
                eval_context = {
                    **execution_inputs.to_dict(),
                    **result_data,
                    current_stage_id: result_data,
                }
                for conn in pathway.connections:
                    if conn.from_node == current_stage_id and conn.should_follow(eval_context):
                        next_stage = conn.to_node
                        transition_reason = conn.when or "default_connection"
                        break

            # Update session with transition
            if next_stage and next_stage != current_stage_id:
                session.transition_to(next_stage, transition_reason)

            # Check if we've reached a terminal state or should await input
            is_complete = False
            awaiting_input = False

            if next_stage in {"complete", "end"}:
                # Explicit terminal state
                session.complete()
                is_complete = True
            elif next_stage is None:
                # No explicit transition - check if this stage has outgoing connections
                # If it does, this is a conversational stage awaiting user input
                has_outgoing = any(
                    conn.from_node == current_stage_id for conn in pathway.connections
                )
                if has_outgoing:
                    # Stage has potential transitions but none matched - await input
                    session.pause()
                    awaiting_input = True
                else:
                    # No outgoing connections - this is a terminal stage
                    session.complete()
                    is_complete = True

            # Emit turn completed event
            self.emitter.emit(
                "stateful_turn_completed",
                pathway_id=pathway.id,
                session_id=session.id,
                stage_id=current_stage_id,
                next_stage=next_stage,
                turn_number=turn_number,
                duration_ms=duration_ms,
                is_complete=is_complete,
                awaiting_input=awaiting_input,
            )

            return {
                "success": True,
                "stage_id": current_stage_id,
                "stage_result": result.get("output", result_data),
                "next_stage": next_stage,
                "transition_reason": transition_reason,
                "is_complete": is_complete,
                "awaiting_input": awaiting_input,
                "session_id": session.id,
                "session_status": session.status.value,
                "turn_number": turn_number,
                "duration_ms": duration_ms,
            }

        except Exception as e:
            session.fail(str(e))

            self.emitter.emit(
                "stateful_turn_failed",
                pathway_id=pathway.id,
                session_id=session.id,
                stage_id=current_stage_id,
                turn_number=turn_number,
                error=str(e),
            )

            return {
                "success": False,
                "error": str(e),
                "stage_id": current_stage_id,
                "session_id": session.id,
                "session_status": session.status.value,
                "turn_number": turn_number,
            }

        finally:
            if self.ctx.run._current_node_id == current_stage_id:
                self.ctx.run._current_node_id = None
            self.ctx.run._session = None

    async def run_to_completion(
        self,
        pathway: Pathway,
        session: PathwaySession,
        initial_inputs: JsonObject,
        max_turns: int = 50,
        turn_inputs_fn: Callable[[PathwaySession, int], 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
            max_turns: Safety limit on turns
            turn_inputs_fn: Optional function to provide inputs for each turn

        Returns
        -------
            Final result with session state
        """
        inputs = initial_inputs or {}

        for turn in range(max_turns):
            # Get inputs for this turn
            if turn > 0 and turn_inputs_fn:
                inputs = turn_inputs_fn(session, turn)

            # Execute one turn
            result = await self.execute_turn(pathway, session, inputs)

            # Clear inputs after first turn unless provided by callback
            if turn == 0 and not turn_inputs_fn:
                inputs = {}

            # Check termination conditions
            if not result.get("success"):
                return result

            if result.get("is_complete"):
                return result

            if result.get("paused"):
                return result

        # Max turns reached
        return {
            "success": False,
            "error": f"Max turns ({max_turns}) reached without completion",
            "session_id": session.id,
            "session_status": session.status.value,
            "turn_number": session.turn_count,
        }
