import uuid
from datetime import datetime, timezone
from typing import Optional, Any, TYPE_CHECKING
from sqlmodel import Field, Relationship, SQLModel, JSON
from sqlalchemy import Column, Text, TIMESTAMP

if TYPE_CHECKING:
    from .chat import Chat

# Helper function for aware UTC timestamps
def aware_utcnow():
    return datetime.now(timezone.utc)

class Message(SQLModel, table=True):
    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    chat_id: uuid.UUID = Field(foreign_key="chat.id", index=True)
    role: str = Field(nullable=False)

    # Stores the RAW, full text content.
    # e.g., "The answer is... Links: [http://...]"
    content: Optional[str] = Field(default=None, sa_column=Column(Text))

    # Stores the CLEAN, parsed answer for display on the frontend.
    # Populated only for the final AI message of a turn.
    answer: Optional[str] = Field(default=None, sa_column=Column(Text))

    # Stores the parsed source links for the frontend.
    links: Optional[list[str]] = Field(default=None, sa_column=Column(JSON))

    # Fields for tool-calling content
    tool_calls: Optional[list[dict[str, Any]]] = Field(default=None, sa_column=Column(JSON))
    tool_call_id: Optional[str] = Field(default=None, index=True)

    # Use aware timestamps and TIMESTAMP WITH TIME ZONE
    created_at: datetime = Field(
        default_factory=aware_utcnow,
        sa_column=Column(TIMESTAMP(timezone=True), nullable=False, index=True)
    )

    chat: Optional["Chat"] = Relationship(back_populates="messages")