from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
from datetime import datetime
from bson import ObjectId
from .user import PyObjectId

class ChatMessageBase(BaseModel):
    user_id: PyObjectId
    index_id: str
    query: str
    response: str
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None
    deleted_by_user: bool = False

class ChatMessageCreate(ChatMessageBase):
    pass

class ChatMessageInDB(ChatMessageBase):
    id: PyObjectId = Field(default_factory=PyObjectId, alias="_id")
    debug_info: Dict[str, Any] = Field(default_factory=dict)
    response_time: float = 0.0
    cached: bool = False
    sources: List[Dict[str, Any]] = Field(default_factory=list)
    context_used: Optional[str] = None
    
    class Config:
        populate_by_name = True
        arbitrary_types_allowed = True
        json_encoders = {ObjectId: str}

class ChatMessage(ChatMessageBase):
    id: PyObjectId = Field(default_factory=PyObjectId, alias="_id")
    debug_info: Dict[str, Any] = Field(default_factory=dict)
    response_time: float = 0.0
    cached: bool = False
    sources: List[Dict[str, Any]] = Field(default_factory=list)
    context_used: Optional[str] = None
    
    class Config:
        populate_by_name = True
        arbitrary_types_allowed = True
        json_encoders = {ObjectId: str}

class ChatQuery(BaseModel):
    query: str
    index_id: str

class ChatResponse(BaseModel):
    response: str
    debug_info: Dict[str, Any] = Field(default_factory=dict)
    cached: bool = False
    response_time: float = 0.0