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 ChartContent(BaseModel):
    type: Literal["chart"] = "chart"
    vega_chart: dict


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


# --- Conversation ---


class ChartConversationResponse(BaseModel):
    uuid: str
    sql_query_uuid: str
    title: str | None
    created_at: datetime
    updated_at: datetime


# --- Message ---


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


# --- Chart ---


class CreateChartRequest(BaseModel):
    message_uuid: str = Field(
        ..., description="UUID of the chart message to attach this chart to"
    )
    sql_query_uuid: Optional[str] = Field(
        None, description="UUID of the SQL query linked to this chart"
    )
    name: Optional[str] = Field(None, max_length=500)
    vega_chart: dict = Field(..., description="Vega chart specification")


class UpdateChartRequest(BaseModel):
    sql_query_uuid: Optional[str] = Field(
        None, description="UUID of the SQL query linked to this chart"
    )
    name: Optional[str] = Field(None, min_length=1, max_length=500)
    vega_chart: Optional[dict] = Field(None, description="Vega chart specification")


class ChartResponse(BaseModel):
    uuid: str
    sql_query_uuid: str | None
    name: str | None
    vega_chart: dict
    created_at: datetime
    updated_at: datetime
