"""Command for loading a chat session and its messages."""

from dataclasses import dataclass
from uuid import UUID

from alloy_runtime_sdk.api_client.client import ApiClient
from alloy_runtime_types.dtos.sessions import MessageResponse
from alloy_runtime_sdk.logging.config import get_logger

from cli.tui.chat.commands.base import CommandResult
from cli.tui.chat.store import ChatStore
from cli.tui.chat.types import ChatPhase

logger = get_logger(__name__)


@dataclass
class LoadSessionResult:
    """Result data for load session operation.

    Attributes:
        session_id: ID of the loaded session.
        session_name: Name of the session.
        agent_id: ID of the agent (extracted from messages).
        agent_name: Name of the agent (if available).
        messages: List of messages in the session.
        total_messages: Total message count.
    """

    session_id: UUID
    session_name: str | None
    agent_id: UUID | None
    agent_name: str | None
    messages: list[MessageResponse]
    total_messages: int


@dataclass
class LoadSessionCommand:
    """Command to load a chat session and its messages.

    This command:
    1. Transitions store to LOADING_SESSION phase
    2. Fetches session data from the API
    3. Extracts agent info from messages
    4. Updates store with session context and messages
    5. Transitions back to IDLE

    Attributes:
        client: Server client for API calls.
        store: Chat store for state management.
        session_id: ID of the session to load.
        message_limit: Maximum number of messages to fetch.
    """

    client: ApiClient
    store: ChatStore
    session_id: UUID
    message_limit: int = 100

    async def execute(self) -> CommandResult[LoadSessionResult]:
        """Execute the load session operation.

        Returns:
            CommandResult with session data on success.
        """
        logger.info("load_session_starting", session_id=str(self.session_id))
        self.store.set_phase(ChatPhase.LOADING_SESSION)

        try:
            session = await self.client.get_session(
                self.session_id, limit=self.message_limit
            )

            # Extract agent info from messages
            agent_id = None
            agent_name = None
            for msg in reversed(session.messages):
                if msg.role == "assistant" and msg.agent_execution:
                    exec_meta = msg.agent_execution
                    if exec_meta.agent_id:
                        agent_id = exec_meta.agent_id
                    break

            # Update store with session data
            self.store.set_session(
                session_id=self.session_id,
                session_name=session.name,
                agent_id=agent_id,
                agent_name=agent_name,
                message_count=session.total_messages,
            )
            self.store.set_messages(list(session.messages))
            self.store.set_phase(ChatPhase.IDLE)

            logger.info(
                "load_session_completed",
                session_id=str(self.session_id),
                session_name=session.name,
                message_count=len(session.messages),
            )
            return CommandResult[LoadSessionResult](
                success=True,
                data=LoadSessionResult(
                    session_id=self.session_id,
                    session_name=session.name,
                    agent_id=agent_id,
                    agent_name=agent_name,
                    messages=list(session.messages),
                    total_messages=session.total_messages,
                ),
            )

        except Exception as e:
            logger.warning(
                "load_session_failed", session_id=str(self.session_id), error=str(e)
            )
            self.store.set_phase(ChatPhase.IDLE)
            return CommandResult[LoadSessionResult](success=False, error=str(e))
