app/db/schemas.py

# app/db/schemas.py

import uuid
from datetime import datetime
from pydantic import BaseModel, Field
import enum as py_enum # Use alias
from typing import List, Dict, Any # Added List, Dict, Any

# --- Define Enums Here ---
class MessageRole(str, py_enum.Enum):
    """Enum for message roles."""
    USER = "user"
    BOT = "bot"

class FeedbackTypeEnum(str, py_enum.Enum):
    """Enum for feedback types."""
    ACCEPT = "accept"
    REJECT = "reject"
# --- End Enums ---


# --- Base Models (for shared fields) ---
class MessageBase(BaseModel):
    """Base schema for message properties."""
    role: MessageRole
    content: str

class SessionBase(BaseModel):
    """Base schema for session properties."""
    name: str | None = None

# --- Schemas for specific data structures ---
class TopResultPreview(BaseModel):
    """Schema for the preview of the top search result."""
    id: str
    score: float
    content_preview: str
    original_file: str | None = None

class QueryResultDetail(BaseModel):
    """Schema for detailed information about a query result."""
    session_id: uuid.UUID
    message_id: uuid.UUID | None = None # Bot message ID
    attempt_number: int = 1
    retrieved_ids: list[str] = []
    search_scores: list[float] = []
    original_file: str | None = None # Top result's original file

# --- Schemas for Creation (request bodies) ---
class MessageCreate(MessageBase):
    """Schema for creating a new message (used internally)."""
    pass

class SessionCreate(SessionBase):
    """Schema for creating a new session."""
    name: str = Field(..., min_length=1, example="My New Chat Session")

class FeedbackCreate(BaseModel):
    """Schema for submitting feedback."""
    message_id: uuid.UUID
    feedback_type: FeedbackTypeEnum
    feedback_comment: str | None = None

class QueryRequest(BaseModel):
    """Schema for the user query request."""
    query: str
    session_id: uuid.UUID | None = None # Session ID is optional


# --- Schemas for Reading (responses, includes fields from DB) ---
class ChatMessage(MessageBase):
    """Schema for reading chat messages from the database."""
    id: uuid.UUID
    session_id: uuid.UUID
    created_at: datetime
    original_query: str | None = None
    retrieved_context_ids: list[str] | None = None
    used_context_ids: list[str] | None = None # IDs used for this specific message
    attempt_number: int = 1
    cumulative_used_context_ids: list[str] | None = None # All IDs used up to this attempt

    class Config:
        from_attributes = True # Use from_attributes for Pydantic v2
        # orm_mode = True # Use this for Pydantic v1

class FeedbackLog(FeedbackCreate):
     """Schema for reading feedback logs from the database."""
     id: uuid.UUID
     created_at: datetime

     class Config:
        from_attributes = True # Use from_attributes for Pydantic v2
        # orm_mode = True # Use this for Pydantic v1

class FeedbackLogResponse(FeedbackLog):
    """Explicit response schema for feedback log entries."""
    pass

class ChatSession(SessionBase):
    """Schema for reading chat sessions from the database."""
    id: uuid.UUID
    created_at: datetime
    # Optionally include messages: list[ChatMessage] = []

    class Config:
        from_attributes = True # Use from_attributes for Pydantic v2
        # orm_mode = True # Use this for Pydantic v1

# --- Schemas for Endpoint Responses ---
class QueryResponse(BaseModel):
    """Schema for the standard query response."""
    llm_answer: str
    context_used_preview: str | None = None
    top_result_preview: TopResultPreview | None = None
    details: QueryResultDetail

class RegeneratedResponse(BaseModel):
    """Schema for the response when regeneration occurs."""
    new_llm_answer: str
    new_context_used_preview: str | None = None
    new_top_result_preview: TopResultPreview | None = None
    new_details: QueryResultDetail # Uses the same detail structure

# --- Schemas for Updating ---
class ChatSessionUpdate(BaseModel):
    """Schema for updating a chat session (e.g., renaming)."""
    name: str = Field(..., min_length=1, example="Renamed Session")