"""PathwaySession - Multi-turn state management for stateful pathways.

A PathwaySession tracks:
- Current stage in the pathway
- Accumulated context from previous turns
- Checkpoints for resume capability
- History of transitions

This enables:
- Multi-turn pathways that span multiple user interactions
- Resume from checkpoint after interruption
- Human-in-the-loop approval flows
- Long-running orchestration patterns

Usage:
    # Start a new session
    session = PathwaySession(
        pathway_id="onboarding-flow",
        initial_stage="intake",
    )

    # Execute a turn
    result = await vm.execute_stateful(pathway, session, user_input)

    # Session tracks state across turns
    print(session.current_stage)  # "review"
    print(session.context)        # accumulated data

    # Resume later
    session = await pathway_service.load_session(session_id)
    result = await vm.execute_stateful(pathway, session, next_input)
"""

import logging
from datetime import UTC, datetime
from enum import StrEnum
from typing import Any
from uuid import uuid4

from pydantic import BaseModel, Field


class SessionStatus(StrEnum):
    """Status of a pathway session."""

    ACTIVE = "active"  # Session is in progress
    PAUSED = "paused"  # Waiting for input/approval
    COMPLETED = "completed"  # Pathway finished successfully
    FAILED = "failed"  # Pathway failed
    EXPIRED = "expired"  # Session timed out


class SessionTransition(BaseModel):
    """Record of a stage transition."""

    from_stage: str
    to_stage: str
    reason: str | None = None
    timestamp: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())
    turn_number: int = 0


class SessionCheckpoint(BaseModel):
    """A saved checkpoint within a session.

    Supports role-based approval routing where specific roles
    (e.g., "manager", "technician") are required to approve checkpoints.
    """

    checkpoint_id: str
    node_id: str
    created_at: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())
    data: dict[str, Any] = Field(default_factory=dict)
    require_approval: bool = False
    assignee_role: str | None = None  # Role required to approve (e.g., "manager")
    approved: bool = False
    approved_at: str | None = None
    approved_by: str | None = None
    approver_role: str | None = None  # Role of the person who approved


class PathwaySession(BaseModel):
    """Session state for a stateful pathway execution.

    Tracks all state needed to pause and resume a multi-turn pathway:
    - Current stage in the state machine
    - Accumulated context from all turns
    - Checkpoints for recovery
    - Full transition history

    Attributes
    ----------
        id: Unique session ID
        pathway_id: ID of the pathway being executed
        status: Current session status
        current_stage: ID of the current stage node
        turn_count: Number of turns executed
        context: Accumulated context (results from all stages)
        checkpoints: Saved checkpoints
        transitions: History of stage transitions
        metadata: Additional session metadata
    """

    model_config = {"extra": "allow"}

    id: str = Field(default_factory=lambda: f"session_{uuid4().hex[:12]}")
    pathway_id: str
    status: SessionStatus = SessionStatus.ACTIVE

    # Current position in the pathway
    current_stage: str | None = None
    turn_count: int = 0

    # Track how many times each node has been visited (for retry limits)
    node_visit_counts: dict[str, int] = Field(default_factory=dict)

    # Accumulated context - results from all executed nodes/stages
    context: dict[str, Any] = Field(default_factory=dict)

    # Checkpoints for resume capability
    checkpoints: list[SessionCheckpoint] = Field(default_factory=list)

    # Transition history
    transitions: list[SessionTransition] = Field(default_factory=list)

    # Timestamps
    created_at: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())
    updated_at: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())
    completed_at: str | None = None

    # Session metadata
    metadata: dict[str, Any] = Field(default_factory=dict)

    # Error tracking
    last_error: str | None = None

    def transition_to(self, new_stage: str, reason: str | None = None) -> None:
        """Record a transition to a new stage.

        Args:
            new_stage: ID of the stage to transition to
            reason: Why we're transitioning (e.g., the condition that matched)
        """
        if self.current_stage:
            self.transitions.append(
                SessionTransition(
                    from_stage=self.current_stage,
                    to_stage=new_stage,
                    reason=reason,
                    turn_number=self.turn_count,
                ),
            )

        self.current_stage = new_stage
        self.updated_at = datetime.now(UTC).isoformat()

    def add_checkpoint(self, checkpoint: SessionCheckpoint) -> None:
        """Add a checkpoint to the session."""
        self.checkpoints.append(checkpoint)
        self.updated_at = datetime.now(UTC).isoformat()

    def get_checkpoint(self, checkpoint_id: str) -> SessionCheckpoint | None:
        """Get a checkpoint by ID."""
        for cp in self.checkpoints:
            if cp.checkpoint_id == checkpoint_id:
                return cp
        return None

    def approve_checkpoint(
        self,
        checkpoint_id: str,
        approved_by: str | None = None,
        approver_role: str | None = None,
    ) -> tuple[bool, str | None]:
        """Approve a checkpoint that requires approval.

        Supports role-based approval where the approver's role must match
        the checkpoint's assignee_role if specified.

        Args:
            checkpoint_id: ID of the checkpoint to approve
            approved_by: Who approved it (user ID or name)
            approver_role: Role of the approver (e.g., "manager", "technician")

        Returns
        -------
            Tuple of (success, error_message). error_message is None on success.
        """
        checkpoint = self.get_checkpoint(checkpoint_id)
        if not checkpoint:
            return False, f"Checkpoint not found: {checkpoint_id}"

        if not checkpoint.require_approval:
            return False, "Checkpoint does not require approval"

        if checkpoint.approved:
            return False, "Checkpoint already approved"

        # Role-based validation: if checkpoint specifies an assignee_role,
        # the approver must have that role
        if checkpoint.assignee_role and approver_role != checkpoint.assignee_role:
            return (
                False,
                f"Approval requires role '{checkpoint.assignee_role}', but approver has role '{approver_role}'",
            )

        checkpoint.approved = True
        checkpoint.approved_at = datetime.now(UTC).isoformat()
        checkpoint.approved_by = approved_by
        checkpoint.approver_role = approver_role
        self.updated_at = datetime.now(UTC).isoformat()

        # If we were paused waiting for approval, resume
        if self.status == SessionStatus.PAUSED:
            self.status = SessionStatus.ACTIVE

        return True, None

    def get_checkpoints_for_role(self, role: str) -> list[SessionCheckpoint]:
        """Get all checkpoints that can be approved by a specific role.

        Args:
            role: The role to filter by (e.g., "manager", "technician")

        Returns
        -------
            List of checkpoints awaiting approval that this role can approve
        """
        return [
            cp
            for cp in self.checkpoints
            if cp.require_approval
            and not cp.approved
            and (cp.assignee_role is None or cp.assignee_role == role)
        ]

    def update_context(self, updates: dict[str, Any]) -> None:
        """Merge updates into the session context."""
        self.context.update(updates)
        self.updated_at = datetime.now(UTC).isoformat()

    def complete(self) -> None:
        """Mark the session as completed."""
        self.status = SessionStatus.COMPLETED
        self.completed_at = datetime.now(UTC).isoformat()
        self.updated_at = self.completed_at

    def fail(self, error: str) -> None:
        """Mark the session as failed."""
        self.status = SessionStatus.FAILED
        self.last_error = error
        self.updated_at = datetime.now(UTC).isoformat()

    def pause(self) -> None:
        """Pause the session (waiting for input/approval)."""
        self.status = SessionStatus.PAUSED
        self.updated_at = datetime.now(UTC).isoformat()

    def resume(self) -> None:
        """Resume a paused session."""
        if self.status == SessionStatus.PAUSED:
            self.status = SessionStatus.ACTIVE
            self.updated_at = datetime.now(UTC).isoformat()

    def increment_turn(self) -> int:
        """Increment and return the turn count."""
        self.turn_count += 1
        self.updated_at = datetime.now(UTC).isoformat()
        return self.turn_count

    def increment_node_visit(self, node_id: str) -> int:
        """Increment and return the visit count for a node.

        Useful for retry limits: state.node_visit_counts.get("writer", 0) < 3
        """
        self.node_visit_counts[node_id] = self.node_visit_counts.get(node_id, 0) + 1
        self.updated_at = datetime.now(UTC).isoformat()
        return self.node_visit_counts[node_id]

    @property
    def is_active(self) -> bool:
        """Check if session is active (can execute)."""
        return self.status == SessionStatus.ACTIVE

    @property
    def is_complete(self) -> bool:
        """Check if session has ended (completed or failed)."""
        return self.status in (SessionStatus.COMPLETED, SessionStatus.FAILED, SessionStatus.EXPIRED)

    @property
    def awaiting_approval(self) -> list[SessionCheckpoint]:
        """Get checkpoints awaiting approval."""
        return [cp for cp in self.checkpoints if cp.require_approval and not cp.approved]


# =============================================================================
# SESSION MANAGER (in-memory)
# =============================================================================

_logger = logging.getLogger(__name__)


class PathwaySessionManager:
    """Manages pathway sessions in memory."""

    def __init__(self) -> None:
        self._sessions: dict[str, PathwaySession] = {}

    async def save(self, session: PathwaySession) -> str:
        """Save a pathway session for later resume. Returns session ID."""
        self._sessions[session.id] = session
        return session.id

    async def load(self, session_id: str) -> PathwaySession | None:
        """Load a pathway session by ID."""
        return self._sessions.get(session_id)

    async def delete(self, session_id: str) -> bool:
        """Delete a pathway session. Returns True always."""
        self._sessions.pop(session_id, None)
        return True

    def list_sessions(
        self,
        *,
        pathway_id: str | None = None,
        status: str | None = None,
    ) -> list[PathwaySession]:
        """List all sessions, optionally filtered."""
        results = []
        for session in self._sessions.values():
            if pathway_id and session.pathway_id != pathway_id:
                continue
            if status and session.status.value != status:
                continue
            results.append(session)
        return results


__all__ = [
    "PathwaySession",
    "PathwaySessionManager",
    "SessionCheckpoint",
    "SessionStatus",
    "SessionTransition",
]
