"""
Checkpoint manager for human-in-the-loop interaction.

Provides async pause/resume functionality for checkpoints with
persistence via FileSessionMemory for resumability across sessions.
"""

import uuid
from typing import Callable, Optional, TYPE_CHECKING
import anyio

from gsd_rlm.checkpoints.types import Checkpoint, CheckpointType

if TYPE_CHECKING:
    from gsd_rlm.session.memory import FileSessionMemory, SessionState


class CheckpointManager:
    """
    Manages checkpoint lifecycle with async pause/resume.

    The CheckpointManager handles:
    - Creating checkpoints via factory methods
    - Pausing execution to wait for user input
    - Persisting checkpoint state for resumability
    - Resuming from pending checkpoints

    Attributes:
        session_memory: FileSessionMemory for persistence.
        on_checkpoint: Optional callback when checkpoint is reached.
        pending_checkpoints: Dict of active checkpoints by session_id.

    Example:
        >>> from gsd_rlm.session.memory import FileSessionMemory
        >>> from pathlib import Path
        >>>
        >>> memory = FileSessionMemory(Path("/tmp/sessions"))
        >>> manager = CheckpointManager(memory)
        >>>
        >>> # Create and pause at checkpoint
        >>> checkpoint = manager.create_human_verify(
        ...     what_built="Login form",
        ...     how_to_verify="Test login",
        ...     session_id="sess-001"
        ... )
        >>> response = await manager.pause(checkpoint)
    """

    def __init__(
        self,
        session_memory: "FileSessionMemory",
        on_checkpoint: Optional[Callable[[Checkpoint], str]] = None,
    ):
        """
        Initialize CheckpointManager.

        Args:
            session_memory: FileSessionMemory instance for persistence.
            on_checkpoint: Optional callback when checkpoint is reached.
                           Receives Checkpoint, returns user response string.
                           If None, pause() returns "approved" immediately.
        """
        self.session_memory = session_memory
        self.on_checkpoint = on_checkpoint
        self.pending_checkpoints: dict[str, Checkpoint] = {}

    def _generate_checkpoint_id(self) -> str:
        """Generate unique checkpoint ID."""
        return f"cp-{uuid.uuid4().hex[:12]}"

    def _store_pending_checkpoint(self, checkpoint: Checkpoint) -> None:
        """
        Store checkpoint in session memory for persistence.

        Args:
            checkpoint: Checkpoint to persist.
        """
        # Store in memory dict for quick access
        self.pending_checkpoints[checkpoint.session_id] = checkpoint

        # Persist to session metadata
        session = self.session_memory.load(checkpoint.session_id)
        if session:
            # Store checkpoint in session metadata
            session.metadata["pending_checkpoint"] = checkpoint.to_dict()
            self.session_memory.save(session)

    def _clear_pending_checkpoint(self, session_id: str) -> None:
        """
        Clear pending checkpoint after response.

        Args:
            session_id: Session to clear checkpoint for.
        """
        # Remove from memory dict
        if session_id in self.pending_checkpoints:
            del self.pending_checkpoints[session_id]

        # Clear from session metadata
        session = self.session_memory.load(session_id)
        if session and "pending_checkpoint" in session.metadata:
            del session.metadata["pending_checkpoint"]
            self.session_memory.save(session)

    async def pause(self, checkpoint: Checkpoint) -> str:
        """
        Pause execution and wait for user response.

        Stores checkpoint in session memory for resumability,
        calls on_checkpoint callback (or returns "approved" for testing),
        then clears the pending checkpoint.

        Args:
            checkpoint: Checkpoint to pause at.

        Returns:
            User's response string (typically "approved" or feedback).

        Example:
            >>> checkpoint = manager.create_human_verify(
            ...     "Built feature", "Test it", "sess-001"
            ... )
            >>> response = await manager.pause(checkpoint)
            >>> if response == "approved":
            ...     print("User approved!")
        """
        # Store checkpoint for persistence/resumability
        self._store_pending_checkpoint(checkpoint)

        try:
            # Call callback or return default
            if self.on_checkpoint:
                response = await self._call_callback(checkpoint)
            else:
                # Default behavior for testing - auto-approve
                response = "approved"

            return response
        finally:
            # Clear pending checkpoint after response
            self._clear_pending_checkpoint(checkpoint.session_id)

    async def _call_callback(self, checkpoint: Checkpoint) -> str:
        """
        Call on_checkpoint callback, handling sync/async.

        Args:
            checkpoint: Checkpoint to pass to callback.

        Returns:
            Response string from callback.
        """
        if self.on_checkpoint is None:
            return "approved"

        result = self.on_checkpoint(checkpoint)

        # Handle both sync and async callbacks
        if hasattr(result, "__await__"):
            return await result
        return result

    def resume_pending(self, session_id: str) -> Optional[Checkpoint]:
        """
        Get pending checkpoint for session resumption.

        Args:
            session_id: Session identifier to check.

        Returns:
            Pending Checkpoint if exists, None otherwise.

        Example:
            >>> checkpoint = manager.resume_pending("sess-001")
            >>> if checkpoint:
            ...     print(f"Found pending checkpoint: {checkpoint.checkpoint_id}")
        """
        # Check memory dict first
        if session_id in self.pending_checkpoints:
            return self.pending_checkpoints[session_id]

        # Check session metadata
        session = self.session_memory.load(session_id)
        if session and "pending_checkpoint" in session.metadata:
            data = session.metadata["pending_checkpoint"]
            checkpoint = Checkpoint.from_dict(data)
            # Restore to memory dict
            self.pending_checkpoints[session_id] = checkpoint
            return checkpoint

        return None

    def create_human_verify(
        self,
        what_built: str,
        how_to_verify: str,
        session_id: str,
        *,
        gate: str = "blocking",
        resume_signal: str = "Type 'approved' or describe issues",
        timeout_seconds: Optional[int] = None,
    ) -> Checkpoint:
        """
        Factory for HUMAN_VERIFY checkpoints.

        Args:
            what_built: Description of what was automated.
            how_to_verify: Steps for user to verify.
            session_id: Session identifier.
            gate: "blocking" or "non-blocking" (default: "blocking").
            resume_signal: Prompt for user response.
            timeout_seconds: Optional timeout.

        Returns:
            Configured HUMAN_VERIFY checkpoint.

        Example:
            >>> checkpoint = manager.create_human_verify(
            ...     what_built="Login form with validation",
            ...     how_to_verify="Visit /login and test credentials",
            ...     session_id="sess-001"
            ... )
        """
        return Checkpoint(
            type=CheckpointType.HUMAN_VERIFY,
            gate=gate,
            checkpoint_id=self._generate_checkpoint_id(),
            session_id=session_id,
            what_built=what_built,
            how_to_verify=how_to_verify,
            resume_signal=resume_signal,
            timeout_seconds=timeout_seconds,
        )

    def create_decision(
        self,
        decision: str,
        context: str,
        options: list[dict[str, str]],
        session_id: str,
        *,
        gate: str = "blocking",
        resume_signal: str = "Select an option (1, 2, 3...)",
        timeout_seconds: Optional[int] = None,
    ) -> Checkpoint:
        """
        Factory for DECISION checkpoints.

        Args:
            decision: What's being decided.
            context: Why this decision matters.
            options: List of option dicts with name/pros/cons.
            session_id: Session identifier.
            gate: "blocking" or "non-blocking" (default: "blocking").
            resume_signal: Prompt for user response.
            timeout_seconds: Optional timeout.

        Returns:
            Configured DECISION checkpoint.

        Example:
            >>> checkpoint = manager.create_decision(
            ...     decision="Choose database",
            ...     context="Need ACID compliance",
            ...     options=[
            ...         {"name": "PostgreSQL", "pros": "ACID", "cons": "Schema"},
            ...         {"name": "MongoDB", "pros": "Flexible", "cons": "No ACID"},
            ...     ],
            ...     session_id="sess-001"
            ... )
        """
        return Checkpoint(
            type=CheckpointType.DECISION,
            gate=gate,
            checkpoint_id=self._generate_checkpoint_id(),
            session_id=session_id,
            decision=decision,
            context=context,
            options=options,
            resume_signal=resume_signal,
            timeout_seconds=timeout_seconds,
        )

    def create_human_action(
        self,
        action: str,
        instructions: str,
        verification: str,
        session_id: str,
        *,
        gate: str = "blocking",
        resume_signal: str = "Type 'done' when complete",
        timeout_seconds: Optional[int] = None,
    ) -> Checkpoint:
        """
        Factory for HUMAN_ACTION checkpoints.

        Args:
            action: The action user must perform.
            instructions: How to perform the action.
            verification: How to verify completion.
            session_id: Session identifier.
            gate: "blocking" or "non-blocking" (default: "blocking").
            resume_signal: Prompt for user response.
            timeout_seconds: Optional timeout.

        Returns:
            Configured HUMAN_ACTION checkpoint.

        Example:
            >>> checkpoint = manager.create_human_action(
            ...     action="Click verification link in email",
            ...     instructions="Check inbox for verification email",
            ...     verification="Run 'auth status' to confirm",
            ...     session_id="sess-001"
            ... )
        """
        return Checkpoint(
            type=CheckpointType.HUMAN_ACTION,
            gate=gate,
            checkpoint_id=self._generate_checkpoint_id(),
            session_id=session_id,
            action=action,
            instructions=instructions,
            verification=verification,
            resume_signal=resume_signal,
            timeout_seconds=timeout_seconds,
        )

    async def pause_with_timeout(
        self,
        checkpoint: Checkpoint,
        timeout: float,
    ) -> str:
        """
        Pause with a timeout, returning "timeout" if exceeded.

        Args:
            checkpoint: Checkpoint to pause at.
            timeout: Timeout in seconds.

        Returns:
            User response or "timeout" if timeout exceeded.

        Example:
            >>> checkpoint = manager.create_human_verify(
            ...     "Feature", "Test it", "sess-001"
            ... )
            >>> response = await manager.pause_with_timeout(checkpoint, 60.0)
            >>> if response == "timeout":
            ...     print("User did not respond in time")
        """
        try:
            with anyio.fail_after(timeout):
                return await self.pause(checkpoint)
        except TimeoutError:
            # Clear pending checkpoint on timeout
            self._clear_pending_checkpoint(checkpoint.session_id)
            return "timeout"
