from sqlmodel import SQLModel, Field, Relationship
from typing import Optional, List, TYPE_CHECKING # Ensure TYPE_CHECKING is imported
from datetime import datetime

# Forward reference for type hinting to avoid circular imports
if TYPE_CHECKING:
    from .user import User # For the 'user' attribute in ChatSession

class ChatSessionBase(SQLModel):
    user_id: int = Field(foreign_key="user.id") # Foreign key to the 'user' table's 'id' column
    start_time: datetime = Field(default_factory=datetime.utcnow)
    title: Optional[str] = None

class ChatSession(ChatSessionBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    
    # ------------ RELATIONSHIP DEFINITION ------------
    # A ChatSession belongs to one User.
    # `back_populates="chat_sessions"` links this to the `chat_sessions` attribute in the User model.
    user: Optional["User"] = Relationship(back_populates="chat_sessions")
    
    # A ChatSession can have many ChatMessage objects.
    messages: List["ChatMessage"] = Relationship(back_populates="session")
    # -------------------------------------------------

class ChatMessageBase(SQLModel):
    # Foreign key to the 'chatsession' table's 'id' column.
    # SQLModel infers table names from class names (lowercase).
    session_id: int = Field(foreign_key="chatsession.id")
    role: str
    content: str
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    tool_call_id: Optional[str] = None
    tool_name: Optional[str] = None

class ChatMessage(ChatMessageBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    
    # A ChatMessage belongs to one ChatSession.
    session: Optional["ChatSession"] = Relationship(back_populates="messages")

class ChatMessageCreate(ChatMessageBase):
    pass