"""Type definitions for the chat module.

Contains enums, protocols, and dataclasses used throughout the chat system.
"""

from dataclasses import dataclass
from datetime import datetime
from enum import Enum, auto
from typing import Protocol, Union
from uuid import UUID

from rich.console import RenderableType
from alloy_runtime_types.dtos.sessions import MessageResponse, SessionSummary


class ChatPhase(Enum):
    """Current phase of chat interaction.

    Used to control UI state and prevent conflicting operations.
    """

    IDLE = auto()
    SENDING = auto()
    LOADING_SESSION = auto()
    LOADING_SESSIONS = auto()
    CREATING_SESSION = auto()


class RenderMode(Enum):
    """Rendering mode for message display."""

    MARKDOWN = auto()
    PLAIN = auto()


@dataclass(frozen=True)
class SessionContext:
    """Immutable context for the active chat session.

    Supports two modes:
    - Agent mode: agent_id and agent_name are set, model comes from agent config
    - Model mode: provider_key and model_name are set, with optional system_instruction
    """

    session_id: UUID
    session_name: str | None
    agent_id: UUID | None
    agent_name: str | None
    message_count: int
    # Model info for direct model mode (no agent)
    provider_key: str | None = None
    model_name: str | None = None
    system_instruction: str | None = None

    @property
    def is_model_mode(self) -> bool:
        """Check if this session uses direct model mode."""
        return self.provider_key is not None and self.model_name is not None

    @property
    def is_agent_mode(self) -> bool:
        """Check if this session uses agent mode."""
        return self.agent_id is not None

    @property
    def display_name(self) -> str:
        """Get the display name for the assistant (agent name or model identifier)."""
        if self.agent_name:
            return self.agent_name
        if self.provider_key and self.model_name:
            return f"{self.provider_key}:{self.model_name}"
        if self.model_name:
            return self.model_name
        return "Assistant"


@dataclass(frozen=True)
class ChatState:
    """Immutable snapshot of chat state.

    All fields are immutable (frozen dataclass with tuple for collections).
    State transitions create new instances rather than mutating.
    """

    phase: ChatPhase = ChatPhase.IDLE
    session: SessionContext | None = None
    messages: tuple[MessageResponse, ...] = ()
    error: str | None = None

    # Optimistic UI: user message shown before server confirms
    pending_user_message: str | None = None

    # Session sidebar state
    sessions: tuple[SessionSummary, ...] = ()
    session_search_query: str = ""

    # UI state
    sidebar_visible: bool = False  # Default: hidden for cleaner initial view
    thinking_collapsed: bool = False  # Whether thinking blocks are collapsed

    # Rendering
    render_mode: RenderMode = RenderMode.MARKDOWN


class MessageRenderer(Protocol):
    """Protocol for message rendering strategies.

    Implement this to create custom renderers (markdown, plain text, etc.).
    """

    def render(
        self,
        messages: tuple[MessageResponse, ...],
        pending_user_message: str | None = None,
        assistant_name: str = "Assistant",
        is_generating: bool = False,
        thinking_collapsed: bool = False,
    ) -> Union[str, RenderableType]:
        """Render messages to a displayable format.

        Args:
            messages: Tuple of message responses to render.
            pending_user_message: User message shown optimistically before server confirms.
            assistant_name: Name to display for assistant (agent name or model identifier).
            is_generating: Whether currently waiting for a final response.
            thinking_collapsed: Whether thinking blocks should be collapsed.

        Returns:
            Formatted content suitable for display in a Static widget.
            Can be a string (for plain text) or a Rich renderable (for markdown).
        """
        ...


def format_time(dt: datetime) -> str:
    """Format datetime for message display."""
    return dt.strftime("%H:%M")


def format_session_time(dt: datetime) -> str:
    """Format datetime for session sidebar display."""
    now = datetime.now(dt.tzinfo)
    if dt.date() == now.date():
        return dt.strftime("%H:%M")
    return dt.strftime("%m/%d")


def extract_message_text(msg: MessageResponse) -> str:
    """Extract text content from a message's content parts (excludes thinking content)."""
    text_parts: list[str] = []
    for part in msg.content_parts:
        # Skip thinking content - it's extracted separately
        if part.content_type_id == "thinking":
            continue
        if part.content_text:
            text_parts.append(part.content_text)
    return "".join(text_parts)


def extract_message_text_for_display(
    msg: MessageResponse,
) -> tuple[str, bool]:
    """Extract text content from a message for display.

    Args:
        msg: The message to extract text from

    Returns:
        Tuple of (text, was_any_truncated)
        - text: The extracted text for display
        - was_any_truncated: Always False in the current TUI contract
    """
    text_parts: list[str] = []

    for part in msg.content_parts:
        if part.content_type_id == "thinking":
            continue
        if part.content_text:
            text_parts.append(part.content_text)

    return "".join(text_parts), False


def extract_thinking_content(msg: MessageResponse) -> str | None:
    """Extract thinking/reasoning content from a message's content parts.

    Returns:
        The thinking content if present, None otherwise.
    """
    for part in msg.content_parts:
        if part.content_type_id == "thinking" and part.content_text:
            return part.content_text
    return None
