"""Finite State Machine (FSM) definitions for the Xulcan Kernel.

This module defines the valid states and state transitions that govern
the agent execution lifecycle. The FSM ensures that the Kernel can only
move through legal state sequences, preventing invalid or unsafe execution paths.

Architecture:
    The FSM follows the execution flow documented in ICEBOX/KERNEL.mmd:
    
    1. Run creation and initialization
    2. State hydration (event replay)
    3. Budget governance (Bursar)
    4. Context preparation (overflow handling)
    5. LLM inference
    6. Response processing
    7. Tool execution (if needed)
    8. Error handling and recovery
    9. Completion

Design Philosophy:
    - Explicit over implicit: All transitions are declared upfront
    - Fail-safe: Invalid transitions are caught at runtime
    - Observable: State changes are logged and auditable
    - Recoverable: Error states have defined recovery paths

References:
    - ICEBOX/KERNEL.mmd: Visual state diagram
    - ICEBOX/SOTA.md: Architecture philosophy
"""

from enum import Enum
from typing import Set, Dict, FrozenSet, Any
from types import MappingProxyType


# ═══════════════════════════════════════════════════════════════════════════
# STATE DEFINITIONS
# ═══════════════════════════════════════════════════════════════════════════

class KernelState(str, Enum):
    """Valid states in the Kernel execution FSM.
    
    Each state represents a distinct phase in the agent lifecycle.
    The Kernel can only be in one state at a time, and transitions
    between states are governed by the VALID_TRANSITIONS graph.
    
    State Categories:
        - Initialization: IDLE, CREATED
        - Hydration: HYDRATING, HYDRATED
        - Governance: CHECKING_BUDGET, CHECKING_POLICY
        - Execution: PREPARING_CONTEXT, CALLING_MODEL, EXECUTING_TOOL
        - Processing: PROCESSING_RESPONSE, PARSING_TOOL_ARGS
        - Terminal: COMPLETED, FAILED
        - Recovery: HANDLING_ERROR, AWAITING_HUMAN
    
    State Semantics:
        IDLE: Initial state, no execution in progress.
        CREATED: Run has been created but not started.
        HYDRATING: Loading events from Ledger and replaying state.
        HYDRATED: State reconstruction complete, ready for governance.
        CHECKING_BUDGET: Bursar validating token/time budget.
        PREPARING_CONTEXT: Building LLM prompt from conversation history.
        COMPACTING_CONTEXT: Context overflow, applying compression strategy.
        CALLING_MODEL: LLM inference in progress.
        PROCESSING_RESPONSE: Parsing and validating LLM output.
        PARSING_TOOL_ARGS: Extracting tool call arguments from response.
        CHECKING_POLICY: Sentinel validating tool call against policies.
        AWAITING_HUMAN: Human-in-the-loop approval required.
        EXECUTING_TOOL: Tool execution in progress.
        HANDLING_ERROR: Processing error and determining recovery strategy.
        COMPLETED: Successful execution, final response generated.
        FAILED: Terminal failure, no recovery possible.
    """
    
    # Initialization
    IDLE = "idle"
    CREATED = "created"
    
    # Hydration (State Replay)
    HYDRATING = "hydrating"
    HYDRATED = "hydrated"
    
    # Governance (Budget)
    CHECKING_BUDGET = "checking_budget"
    
    # Context Management
    PREPARING_CONTEXT = "preparing_context"
    COMPACTING_CONTEXT = "compacting_context"
    
    # LLM Inference
    CALLING_MODEL = "calling_model"
    PROCESSING_RESPONSE = "processing_response"
    
    # Tool Execution Path
    PARSING_TOOL_ARGS = "parsing_tool_args"
    CHECKING_POLICY = "checking_policy"
    AWAITING_HUMAN = "awaiting_human"
    EXECUTING_TOOL = "executing_tool"
    
    # Error Recovery
    HANDLING_ERROR = "handling_error"
    
    # Terminal States
    COMPLETED = "completed"
    FAILED = "failed"


# ═══════════════════════════════════════════════════════════════════════════
# TRANSITION GRAPH
# ═══════════════════════════════════════════════════════════════════════════

# Private mutable dictionary
_VALID_TRANSITIONS_DICT: Dict[KernelState, FrozenSet[KernelState]] = {
    # From IDLE: Can only create a new run
    KernelState.IDLE: frozenset({
        KernelState.CREATED,
    }),
    
    # From CREATED: Must hydrate state
    KernelState.CREATED: frozenset({
        KernelState.HYDRATING,
        KernelState.FAILED,  # Immediate failure (e.g., invalid blueprint)
    }),
    
    # From HYDRATING: Either succeed or fail
    KernelState.HYDRATING: frozenset({
        KernelState.HYDRATED,
        KernelState.FAILED,  # Ledger errors
    }),
    
    # From HYDRATED: Check budget before proceeding
    KernelState.HYDRATED: frozenset({
        KernelState.CHECKING_BUDGET,
    }),
    
    # From CHECKING_BUDGET: Authorized or denied
    KernelState.CHECKING_BUDGET: frozenset({
        KernelState.PREPARING_CONTEXT,  # Budget OK or Exceeded but SOFT_NOTIFY
        KernelState.FAILED,             # Budget exceeded (HARD_CAP)
    }),
    
    # From PREPARING_CONTEXT: Check for overflow
    KernelState.PREPARING_CONTEXT: frozenset({
        KernelState.COMPACTING_CONTEXT,  # Context overflow detected
        KernelState.CALLING_MODEL,       # Context fits, proceed
        KernelState.HANDLING_ERROR,      # Context preparation failed
    }),
    
    # From COMPACTING_CONTEXT: Return to context prep or fail
    KernelState.COMPACTING_CONTEXT: frozenset({
        KernelState.CALLING_MODEL,  # Compaction successful
        KernelState.FAILED,         # Cannot compact further
    }),
    
    # From CALLING_MODEL: Receive response or error
    KernelState.CALLING_MODEL: frozenset({
        KernelState.PROCESSING_RESPONSE,  # Valid HTTP response
        KernelState.HANDLING_ERROR,       # Provider error
    }),
    
    # From PROCESSING_RESPONSE: Text response or tool calls
    KernelState.PROCESSING_RESPONSE: frozenset({
        KernelState.COMPLETED,         # Text-only response, done
        KernelState.PARSING_TOOL_ARGS, # Tool calls detected
        KernelState.HANDLING_ERROR,    # Invalid response format
    }),
    
    # From PARSING_TOOL_ARGS: Valid JSON or recovery attempt
    KernelState.PARSING_TOOL_ARGS: frozenset({
        KernelState.CHECKING_POLICY,  # Valid tool call
        KernelState.HANDLING_ERROR,   # JSON parsing failed (after repair attempts)
    }),
    
    # From CHECKING_POLICY: Approved, denied, or human review
    KernelState.CHECKING_POLICY: frozenset({
        KernelState.EXECUTING_TOOL,  # Policy approved
        KernelState.AWAITING_HUMAN,  # Human-in-the-loop required
        KernelState.HANDLING_ERROR,  # Policy violation
    }),
    
    # From AWAITING_HUMAN: User approves or denies
    KernelState.AWAITING_HUMAN: frozenset({
        KernelState.EXECUTING_TOOL,  # User approved
        KernelState.HANDLING_ERROR,  # User rejected
    }),
    
    # From EXECUTING_TOOL: Success or failure
    KernelState.EXECUTING_TOOL: frozenset({
        KernelState.CHECKING_BUDGET, # Tool succeeded, check budget before next step
        KernelState.HANDLING_ERROR,  # Tool execution failed
    }),
    
    # From HANDLING_ERROR: Retry or terminal failure
    KernelState.HANDLING_ERROR: frozenset({
        KernelState.CHECKING_BUDGET, # Transient error, retry (error is in RAM context now)
        KernelState.FAILED,          # Structural error, fatal
    }),
    
    # Terminal states: No outgoing transitions
    KernelState.COMPLETED: frozenset(),
    KernelState.FAILED: frozenset(),
}

for state, valid_targets in _VALID_TRANSITIONS_DICT.items():
    if state not in (KernelState.COMPLETED, KernelState.FAILED):
        _VALID_TRANSITIONS_DICT[state] = valid_targets | frozenset({KernelState.FAILED})

# Public Immutable Proxy. Prevents runtime mutation of the laws of physics.
VALID_TRANSITIONS = MappingProxyType(_VALID_TRANSITIONS_DICT)


# ═══════════════════════════════════════════════════════════════════════════
# VALIDATION UTILITIES
# ═══════════════════════════════════════════════════════════════════════════

class InvalidTransitionError(RuntimeError):
    """Raised when attempting an illegal state transition.
    
    This error indicates a bug in the Orchestrator logic, as all transitions
    should be validated before execution. If this exception is raised in
    production, it represents a critical invariant violation.
    
    Attributes:
        from_state: The current state.
        to_state: The attempted target state.
        valid_transitions: The set of legal transitions from from_state.
    """
    
    def __init__(
        self,
        from_state: KernelState,
        to_state: KernelState,
        valid_transitions: FrozenSet[KernelState],
    ):
        self.from_state = from_state
        self.to_state = to_state
        self.valid_transitions = valid_transitions
        
        # Defensive getter to prevent crashes if a non-Enum is passed
        def _get_val(state: Any) -> str:
            return getattr(state, 'value', str(state))
        
        super().__init__(
            f"Invalid transition: {_get_val(from_state)} → {_get_val(to_state)}. "
            f"Valid transitions from {_get_val(from_state)}: "
            f"{{{', '.join(_get_val(s) for s in valid_transitions)}}}"
        )


def validate_transition(
    from_state: KernelState,
    to_state: KernelState,
) -> None:
    """Validates that a state transition is legal according to the FSM.
    
    This function checks the VALID_TRANSITIONS graph to ensure that
    the proposed transition is allowed. If not, it raises an exception.
    
    Args:
        from_state: The current state.
        to_state: The target state.
    
    Raises:
        InvalidTransitionError: If the transition is not allowed.
    
    Example:
        >>> validate_transition(KernelState.IDLE, KernelState.CREATED)
        # Returns None (valid transition)
        
        >>> validate_transition(KernelState.IDLE, KernelState.COMPLETED)
        # Raises InvalidTransitionError
    """
    valid_targets = VALID_TRANSITIONS.get(from_state, frozenset())
    
    if to_state not in valid_targets:
        raise InvalidTransitionError(from_state, to_state, valid_targets)


def is_terminal_state(state: KernelState) -> bool:
    """Checks if a state is terminal (no outgoing transitions).
    
    Terminal states represent the end of execution. Once in a terminal
    state, the run cannot transition further.
    
    Args:
        state: The state to check.
    
    Returns:
        True if the state is terminal (COMPLETED or FAILED).
    
    Example:
        >>> is_terminal_state(KernelState.COMPLETED)
        True
        >>> is_terminal_state(KernelState.IDLE)
        False
    """
    return len(VALID_TRANSITIONS.get(state, frozenset())) == 0


def get_valid_transitions(state: KernelState) -> Set[KernelState]:
    """Returns the set of valid target states from the given state.
    
    Useful for debugging, logging, and visualization of the FSM.
    
    Args:
        state: The state to query.
    
    Returns:
        Set of states that can be transitioned to from the given state.
    
    Example:
        >>> get_valid_transitions(KernelState.IDLE)
        {<KernelState.CREATED: 'created'>}
    """
    return set(VALID_TRANSITIONS.get(state, frozenset()))


# ═══════════════════════════════════════════════════════════════════════════
# STATE CATEGORIZATION
# ═══════════════════════════════════════════════════════════════════════════

INITIALIZATION_STATES = frozenset({
    KernelState.IDLE,
    KernelState.CREATED,
})

HYDRATION_STATES = frozenset({
    KernelState.HYDRATING,
    KernelState.HYDRATED,
})

GOVERNANCE_STATES = frozenset({
    KernelState.CHECKING_BUDGET,
    KernelState.CHECKING_POLICY,
})

EXECUTION_STATES = frozenset({
    KernelState.PREPARING_CONTEXT,
    KernelState.COMPACTING_CONTEXT,
    KernelState.CALLING_MODEL,
    KernelState.EXECUTING_TOOL,
})

PROCESSING_STATES = frozenset({
    KernelState.PROCESSING_RESPONSE,
    KernelState.PARSING_TOOL_ARGS,
})

ERROR_STATES = frozenset({
    KernelState.HANDLING_ERROR,
    KernelState.AWAITING_HUMAN,
})

TERMINAL_STATES = frozenset({
    KernelState.COMPLETED,
    KernelState.FAILED,
})


def get_state_category(state: KernelState) -> str:
    """Returns the semantic category of a state.
    
    Useful for logging, metrics, and understanding execution flow.
    
    Args:
        state: The state to categorize.
    
    Returns:
        String name of the category.
    
    Example:
        >>> get_state_category(KernelState.CALLING_MODEL)
        'execution'
    """
    if state in INITIALIZATION_STATES:
        return "initialization"
    elif state in HYDRATION_STATES:
        return "hydration"
    elif state in GOVERNANCE_STATES:
        return "governance"
    elif state in EXECUTION_STATES:
        return "execution"
    elif state in PROCESSING_STATES:
        return "processing"
    elif state in ERROR_STATES:
        return "error_recovery"
    elif state in TERMINAL_STATES:
        return "terminal"
    else:
        return "unknown"