"""Pydantic schemas for the AI Chat Builder."""
from typing import Any, Literal, Optional
from uuid import UUID

from pydantic import BaseModel, Field


# ── Requests ──


class AIChatStartRequest(BaseModel):
    """Start a new chat builder session."""

    flow_type: Literal["troubleshooting", "procedural"] = Field(
        ..., description="Type of flow to build"
    )
    tree_id: Optional[str] = Field(
        default=None,
        description="ID of existing tree for editor-embedded sessions",
    )


VALID_ACTION_TYPES = Literal[
    "generate_full",
    "generate_branch",
    "modify_node",
    "add_steps",
    "quick_action",
    "open_chat",
    "variable_inference",
]


class AIChatMessageRequest(BaseModel):
    """Send a user message in a chat session."""

    content: str = Field(..., min_length=1, max_length=5000)
    action_type: Optional[VALID_ACTION_TYPES] = Field(
        default="open_chat",
        description="Type of AI action to perform",
    )
    focal_node_id: Optional[str] = Field(
        default=None,
        description="ID of the node/step being acted on",
    )
    flow_context: Optional[dict[str, Any]] = Field(
        default=None,
        description="Live flow structure from the editor (tree structure, steps, intake form)",
    )


class AIChatImportRequest(BaseModel):
    """Import generated tree with optional metadata overrides."""

    name: Optional[str] = Field(None, min_length=1, max_length=255)
    description: Optional[str] = Field(None, max_length=2000)
    category_id: Optional[UUID] = None
    tags: list[str] = Field(default_factory=list)


# ── Responses ──


class AIChatStartResponse(BaseModel):
    """Response after creating a chat session."""

    session_id: UUID
    greeting: str
    current_phase: str


class AIChatMessageResponse(BaseModel):
    """Response after sending a message."""

    content: str
    current_phase: str
    working_tree: Optional[dict[str, Any]] = None
    tree_metadata: Optional[dict[str, Any]] = None


class AIChatSessionResponse(BaseModel):
    """Full session state for resume."""

    session_id: UUID
    status: str
    current_phase: str
    flow_type: str
    conversation_history: list[dict[str, Any]]
    working_tree: Optional[dict[str, Any]] = None
    tree_metadata: Optional[dict[str, Any]] = None
    message_count: int
    generated_tree: Optional[dict[str, Any]] = None


class AIChatGenerateResponse(BaseModel):
    """Response with the final generated tree."""

    tree_structure: dict[str, Any]
    tree_metadata: dict[str, Any]
    status: str


class AIChatImportResponse(BaseModel):
    """Response after importing tree to editor."""

    tree_id: UUID
    tree_type: str