from pydantic import BaseModel, Field
from typing import List, Optional


class BaleStartMessageForm(BaseModel):
    id: int
    is_bot: bool = False
    first_name: str
    last_name: Optional[str] = None
    username: Optional[str] = None


class BaleStartMessageChat(BaseModel):
    id: int
    type: str
    username: Optional[str] = None
    first_name: Optional[str] = None


class BaleStartMessage(BaseModel):
    message_id: int
    from_user: BaleStartMessageForm = Field(..., alias="from")
    date: int
    chat: BaleStartMessageChat
    text: str
    entities: List[dict] = []

    class Config:
        populate_by_name = True


class BaleCallbackFrom(BaseModel):
    id: int
    is_bot: bool
    first_name: str
    username: Optional[str] = None


class BaleCallbackMessage(BaseModel):
    message_id: int
    chat: BaleStartMessageChat
    text: Optional[str]


class BaleCallbackQuery(BaseModel):
    id: str
    from_user: BaleCallbackFrom = Field(..., alias="from")
    message: BaleCallbackMessage
    data: str

    class Config:
        populate_by_name = True


class BaleUpdate(BaseModel):
    update_id: int
    message: Optional[BaleStartMessage] = None
    callback_query: Optional[BaleCallbackQuery] = None

    class Config:
        exclude_none = True


class QaChatSingle(BaseModel):
    id: str
    chat_id: int
    user_query: str
    model_key: str
    model_effort: str

    retrived_passage: str
    retrived_ref_ids: str
    retrived_duration: Optional[int] = 0
    prompt_type: str = "question"
    llm_duration: int
    full_duration: Optional[int] = 0
    time_create: Optional[int] = 0
    used_ref_ids: Optional[str] = ""
    status_text: Optional[str] = ""
    status: Optional[int] = 0
    prompt_answer: str
    other_info: dict | None


class QaChatBlock(BaseModel):
    id: str
    title: str
    user_id: str
    is_premium: bool
    chat: QaChatSingle
    total_token: int
    is_end: bool


class QaChat(BaseModel):
    id: str
    chat_id: int
    title: Optional[str] = ""
    user_id: str
    user_query: str
    query_type: str = "question" # llm -> greeting, other, legal_question | rag -> question 
    full_duration: Optional[float] = 0
    other_info: Optional[dict] = ""

    ss_ref_ids: Optional[List[str]] = ""
    ss_model_key: Optional[str] = ""
    ss_duration: Optional[float] = 0
    ss_answer: Optional[str] = ""

    llm_ref_ids: Optional[List[str]] = []
    llm_model_key: Optional[str] = ""
    llm_duration: Optional[float] = 0
    llm_answer: Optional[str] = ""
    
    status_text: Optional[str] = ""
    status: Optional[int] = 0