from sqlmodel import SQLModel, Field, Relationship
from typing import Optional, List
from datetime import datetime
import uuid
from sqlalchemy import text


class MessageAsset(SQLModel, table=True):
    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        index=True,
        nullable=False,
        sa_column_kwargs={"server_default": text("gen_random_uuid()")},
    )
    message_id: uuid.UUID = Field(foreign_key="message.id", index=True)
    asset_type: str = Field(index=True)  # "image" | "file" | "audio" | "video" | etc.
    url: str = Field()                    # public or signed URL
    # Pretty good suggestions but not needed for now.
    # filename: Optional[str] = Field(default=None)
    # mime_type: Optional[str] = Field(default=None)
    # width: Optional[int] = Field(default=None)   # for images/videos
    # height: Optional[int] = Field(default=None)
    # size_bytes: Optional[int] = Field(default=None)
    # alt_text: Optional[str] = Field(default=None)  # user-provided description
    created_at: datetime = Field(default_factory=datetime.utcnow, nullable=False)

    # Relationship back to message
    message: "Message" = Relationship(back_populates="assets")

class Message(SQLModel, table=True):
    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        index=True,
        nullable=False,
        sa_column_kwargs={"server_default": text("gen_random_uuid()")},
    )
    chat_id: uuid.UUID = Field(foreign_key="chat.id")

    role: str = Field(
        index=True
    )  # "user" | "assistant" TODO: use enum?  I don't mind just defining like this tho.
    content: str
    created_at: datetime = Field(default_factory=datetime.utcnow, nullable=False)

    # Relationship
    chat: Optional["Chat"] = Relationship(back_populates="messages")
    assets: List[MessageAsset] = Relationship(back_populates="message")

class Chat(SQLModel, table=True):
    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        index=True,
        nullable=False,
        sa_column_kwargs={"server_default": text("gen_random_uuid()")},
    )

    title: str = Field(default="New Chat")
    created_at: datetime = Field(default_factory=datetime.utcnow, nullable=False)

    # Relationship
    messages: List[Message] = Relationship(
        back_populates="chat", sa_relationship_kwargs={"cascade": "all, delete-orphan"}
    )