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

from __future__ import annotations

import time
from typing import TYPE_CHECKING

from core.common import JsonObject
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

if TYPE_CHECKING:
    import asyncio
    from collections.abc import Callable

    from core.pathways.context import Context
    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
        visit_count = 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:
                        if 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,
        }
