"""Query request/response Pydantic models."""

from typing import Any
from pydantic import BaseModel, Field


class QuestionRequest(BaseModel):
    """Single question for batch query."""

    body: str
    with_docs: bool = True


class BenchQueryRequest(BaseModel):
    """Bench mode query request."""

    environment: str  # ift, psi, prod
    questions: list[QuestionRequest]


class BackendQueryRequest(BaseModel):
    """Backend mode query request (one-by-one)."""

    environment: str  # ift, psi, prod
    questions: list[QuestionRequest]
    reset_session: bool = True


class Docs(BaseModel):
    """Documents from RAG."""

    research: list
    analytical_hub: list


class RagResponse(BaseModel):
    """Ответ от RAG на вопрос пользователя."""

    body_research: str = Field(description="Текст ответа от Research на вопрос")
    body_analytical_hub: str = Field(description="Текст ответа от Analytical Hub на вопрос")
    docs_from_vectorstore: Docs | None = None
    docs_to_llm: Docs | None = None


class RagResponseBench(RagResponse):
    """Ответ на вопрос + время обработки именно этого вопроса."""

    processing_time_sec: float = Field(
        description="Время обработки запроса в секундах",
        ge=0,
    )
    question: str = Field(description="Исходный вопрос")


class RagResponseBenchList(BaseModel):
    """Список ответов RAG в bench режиме."""

    answers: list[RagResponseBench]


class QueryResponse(BaseModel):
    """Query response with metadata."""

    request_id: str
    timestamp: str
    environment: str
    response: RagResponseBenchList | dict | list  # RagResponseBenchList для bench, dict/list для backend