"""
State management for orchestration patterns.

This module provides state tracking and checkpoint functionality for
pattern execution, enabling pause/resume and state recovery.
"""

from typing import Any, Dict, Optional
from pydantic import BaseModel, Field
from datetime import datetime
from enum import Enum


class StateStatus(str, Enum):
    """
    Enumeration of execution state statuses.

    Tracks the lifecycle of pattern execution:
    - INITIALIZED: State created, not yet started
    - RUNNING: Actively executing
    - PAUSED: Execution temporarily suspended
    - COMPLETED: Successfully finished
    - FAILED: Execution failed with error
    """

    INITIALIZED = "initialized"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"


class State(BaseModel):
    """
    State model for tracking pattern execution state.

    This Pydantic model maintains execution status, data, checkpoints,
    and timestamps for pattern lifecycle management.

    :param status: StateStatus - Current execution status.
    :param data: Dict[str, Any] - Arbitrary state data storage.
    :param checkpoint: Optional[Dict[str, Any]] - Saved checkpoint for recovery.
    :param created_at: datetime - Timestamp when state was created.
    :param updated_at: datetime - Timestamp of last update.
    """

    status: StateStatus = Field(StateStatus.INITIALIZED, description="Current state status")
    data: Dict[str, Any] = Field(default_factory=dict)
    checkpoint: Optional[Dict[str, Any]] = None
    created_at: datetime = Field(default_factory=datetime.utcnow)
    updated_at: datetime = Field(default_factory=datetime.utcnow)

    def update_status(self, status: StateStatus) -> None:
        """
        Update the execution status and timestamp.

        :param status: StateStatus - New status to set.
        """
        self.status = status
        self.updated_at = datetime.utcnow()

    def set(self, key: str, value: Any) -> None:
        """
        Set a data value and update timestamp.

        :param key: str - Data key.
        :param value: Any - Value to store.
        """
        self.data[key] = value
        self.updated_at = datetime.utcnow()

    def get(self, key: str, default: Any = None) -> Any:
        """
        Get a data value with optional default.

        :param key: str - Data key.
        :param default: Any - Default value if key not found.
        :return: Any - Stored value or default.
        """
        return self.data.get(key, default)

    def create_checkpoint(self) -> None:
        """
        Create a checkpoint snapshot of current state.

        Captures status, data, and timestamp for later restoration.
        """
        self.checkpoint = {
            "status": self.status.value,
            "data": self.data.copy(),
            "timestamp": datetime.utcnow().isoformat(),
        }

    def restore_checkpoint(self) -> bool:
        """
        Restore state from the most recent checkpoint.

        :return: bool - True if restored successfully, False if no checkpoint exists.
        """
        if self.checkpoint is None:
            return False
        self.status = StateStatus(self.checkpoint["status"])
        self.data = self.checkpoint["data"].copy()
        self.updated_at = datetime.utcnow()
        return True

    class Config:
        json_encoders = {datetime: lambda v: v.isoformat()}


class StateManager:
    """
    Manager for multiple pattern execution states.

    Provides centralized management of multiple state instances,
    enabling orchestration of concurrent pattern executions.
    """

    def __init__(self):
        """Initialize the state manager with empty state registry."""
        self._states: Dict[str, State] = {}

    def create_state(self, state_id: str, initial_data: Optional[Dict[str, Any]] = None) -> State:
        """
        Create and register a new state.

        :param state_id: str - Unique identifier for the state.
        :param initial_data: Optional[Dict[str, Any]] - Initial data for the state.
        :return: State - The newly created state instance.
        """
        state = State(data=initial_data or {})
        self._states[state_id] = state
        return state

    def get_state(self, state_id: str) -> Optional[State]:
        """
        Retrieve a state by ID.

        :param state_id: str - State identifier.
        :return: Optional[State] - The state if found, None otherwise.
        """
        return self._states.get(state_id)

    def update_state(
        self, state_id: str, status: StateStatus, data: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        Update an existing state's status and data.

        :param state_id: str - State identifier.
        :param status: StateStatus - New status to set.
        :param data: Optional[Dict[str, Any]] - Data to merge into state.
        :return: bool - True if updated, False if state not found.
        """
        state = self._states.get(state_id)
        if state is None:
            return False
        state.update_status(status)
        if data:
            state.data.update(data)
        return True

    def delete_state(self, state_id: str) -> None:
        """
        Delete a state by ID.

        :param state_id: str - State identifier to delete.
        """
        if state_id in self._states:
            del self._states[state_id]

    def list_states(self) -> list:
        """
        List all registered state IDs.

        :return: list - List of state identifiers.
        """
        return list(self._states.keys())

    def clear_all(self) -> None:
        self._states.clear()
