from typing import Annotated, Literal, Optional, Union
from datetime import datetime

from pydantic import BaseModel, Field


class TextContent(BaseModel):
    type: Literal["text"] = "text"
    text: str
    suggestions: list[str] | None = None


class SqlContent(BaseModel):
    type: Literal["sql"] = "sql"
    sql: str


MessageContent = Annotated[
    Union[TextContent, SqlContent],
    Field(discriminator="type"),
]


class GenerateSqlRequest(BaseModel):
    user_db_uuid: str = Field(..., description="UUID of the user database to query")
    user_input: str = Field(..., min_length=1, description="User input")
    conversation_uuid: Optional[str] = Field(
        None, description="UUID of existing conversation (None to create new)"
    )


class SqlConversationResponse(BaseModel):
    uuid: str
    user_db_uuid: str
    title: str | None
    created_at: datetime
    updated_at: datetime


class SqlMessageResponse(BaseModel):
    uuid: str
    role: str
    content: MessageContent
    created_at: datetime


class SqlQueryResponse(BaseModel):
    uuid: str
    user_db_uuid: str
    name: str | None
    sql: str
    created_at: datetime
    updated_at: datetime


class UpdateSqlQueryRequest(BaseModel):
    user_db_uuid: Optional[str] = Field(None, description="UUID of the user database")
    name: Optional[str] = Field(None, min_length=1, max_length=500)
    sql: Optional[str] = Field(None, min_length=1)


class CreateSqlQueryRequest(BaseModel):
    message_uuid: str = Field(
        ..., description="UUID of the SQL message to attach this query to"
    )
    name: Optional[str] = Field(None, max_length=500)
    sql: str = Field(..., min_length=1)
