"""State management for agent lifecycle."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from typing import TYPE_CHECKING, Any

import structlog

from curadise_agent.state.storage import StateStorage

if TYPE_CHECKING:
    from curadise_agent.config.schema import StateConfig

log = structlog.get_logger(__name__)


class AgentState(Enum):
    """Agent lifecycle states."""

    INITIALIZING = "initializing"
    REGISTERING = "registering"
    RUNNING = "running"
    DEGRADED = "degraded"
    SHUTTING_DOWN = "shutting_down"
    STOPPED = "stopped"


@dataclass
class StateManager:
    """
    Manages agent state and lifecycle.

    Coordinates between runtime state and persistent storage,
    handling state transitions and persistence of critical data.
    """

    storage: StateStorage
    _state: AgentState = field(default=AgentState.INITIALIZING, init=False)
    _agent_id: str | None = field(default=None, init=False)
    _token: str | None = field(default=None, init=False)
    _state_changed_at: datetime = field(default_factory=lambda: datetime.now(UTC), init=False)
    _metadata: dict[str, Any] = field(default_factory=dict, init=False)

    @property
    def state(self) -> AgentState:
        """Get current agent state."""
        return self._state

    @property
    def agent_id(self) -> str | None:
        """Get agent ID."""
        return self._agent_id

    @property
    def token(self) -> str | None:
        """Get authentication token."""
        return self._token

    @property
    def is_registered(self) -> bool:
        """Check if agent is registered."""
        return self._agent_id is not None

    @property
    def is_authenticated(self) -> bool:
        """Check if agent has a valid token."""
        return self._token is not None

    @property
    def is_running(self) -> bool:
        """Check if agent is in running state."""
        return self._state == AgentState.RUNNING

    @property
    def state_duration_seconds(self) -> float:
        """Get duration in current state."""
        return (datetime.now(UTC) - self._state_changed_at).total_seconds()

    def transition_to(self, new_state: AgentState) -> None:
        """
        Transition to a new state.

        Args:
            new_state: Target state

        Raises:
            ValueError: If transition is invalid
        """
        valid_transitions = {
            AgentState.INITIALIZING: {
                AgentState.REGISTERING,
                AgentState.RUNNING,
                AgentState.SHUTTING_DOWN,
            },
            AgentState.REGISTERING: {AgentState.RUNNING, AgentState.SHUTTING_DOWN},
            AgentState.RUNNING: {AgentState.DEGRADED, AgentState.SHUTTING_DOWN},
            AgentState.DEGRADED: {AgentState.RUNNING, AgentState.SHUTTING_DOWN},
            AgentState.SHUTTING_DOWN: {AgentState.STOPPED},
            AgentState.STOPPED: set(),
        }

        if new_state not in valid_transitions.get(self._state, set()):
            log.warning(
                "invalid_state_transition",
                current_state=self._state.value,
                target_state=new_state.value,
            )
            # Allow transition anyway for flexibility

        old_state = self._state
        self._state = new_state
        self._state_changed_at = datetime.now(UTC)

        log.info(
            "state_transition",
            from_state=old_state.value,
            to_state=new_state.value,
        )

    async def load(self) -> None:
        """Load persisted state from storage."""
        # Load agent ID
        agent_id = self.storage.get_agent_id()
        if agent_id:
            self._agent_id = agent_id
            log.info("agent_id_loaded", agent_id=agent_id)

        # Load token
        token = self.storage.get_token()
        if token:
            self._token = token
            log.info("token_loaded")

        # Load metadata
        metadata = self.storage.get("metadata")
        if metadata:
            self._metadata = metadata
            log.debug("metadata_loaded", keys=list(metadata.keys()))

    async def save(self) -> None:
        """Persist current state to storage."""
        if self._agent_id:
            self.storage.save_agent_id(self._agent_id)

        if self._token:
            self.storage.save_token(self._token)

        if self._metadata:
            self.storage.set("metadata", self._metadata)

    def set_agent_id(self, agent_id: str, persist: bool = True) -> None:
        """
        Set the agent ID.

        Args:
            agent_id: Agent ID to set
            persist: Whether to persist to storage
        """
        self._agent_id = agent_id
        if persist:
            self.storage.save_agent_id(agent_id)
        log.info("agent_id_set", agent_id=agent_id)

    def set_token(self, token: str, persist: bool = True) -> None:
        """
        Set the authentication token.

        Args:
            token: Token to set
            persist: Whether to persist to storage
        """
        self._token = token
        if persist:
            self.storage.save_token(token)
        log.info("token_set")

    def clear_credentials(self, persist: bool = True) -> None:
        """
        Clear agent credentials.

        Args:
            persist: Whether to remove from storage
        """
        self._agent_id = None
        self._token = None

        if persist:
            self.storage.delete_agent_id()
            self.storage.delete_token()

        log.info("credentials_cleared")

    def set_metadata(self, key: str, value: Any) -> None:
        """Set a metadata value."""
        self._metadata[key] = value

    def get_metadata(self, key: str, default: Any = None) -> Any:
        """Get a metadata value."""
        return self._metadata.get(key, default)

    def get_status(self) -> dict[str, Any]:
        """
        Get comprehensive status information.

        Returns:
            Status dictionary
        """
        return {
            "state": self._state.value,
            "state_duration_seconds": self.state_duration_seconds,
            "is_registered": self.is_registered,
            "is_authenticated": self.is_authenticated,
            "agent_id": self._agent_id,
            "metadata": self._metadata,
        }


def create_state_manager(config: StateConfig) -> StateManager:
    """
    Factory function to create a configured state manager.

    Args:
        config: State configuration

    Returns:
        Configured StateManager
    """
    storage = StateStorage(directory=config.directory)
    return StateManager(storage=storage)
