from __future__ import annotations

from typing import Any

from pydantic import BaseModel, Field


class HealthResponse(BaseModel):
    status: str


class IndexRequest(BaseModel):
    source_ids: list[int] | None = None
    law_types: list[str] | None = None
    reset_collection: bool = True
    batch_size: int | None = None


class IndexResponse(BaseModel):
    indexed_chunks: int
    indexed_sources: int
    collection_name: str


class SearchRequest(BaseModel):
    question: str
    category: str | None = None
    region: str | None = None
    user_type: str | None = None
    history: list[dict[str, str]] = Field(default_factory=list)
    law_types: list[str] | None = None
    top_k: int = 5


class RetrievedChunk(BaseModel):
    chunk_id: int
    source_id: int
    source_title: str
    source_url: str | None = None
    law_type: str | None = None
    article_number: str | None = None
    article_title: str | None = None
    chunk_text: str
    metadata: dict[str, Any] = Field(default_factory=dict)
    score: float


class ClassificationResult(BaseModel):
    legal_domain: str
    issue_type: str
    jurisdiction: str = "RU"
    region: str | None = None
    needs_clarification: bool = False
    clarification_questions: list[str] = Field(default_factory=list)
    search_queries: list[str] = Field(default_factory=list)
    filters: dict[str, Any] = Field(default_factory=dict)


class StructuredInitialAnswer(BaseModel):
    short_conclusion: str
    legal_points: list[str] = Field(default_factory=list)
    action_steps: list[str] = Field(default_factory=list)
    risks: list[str] = Field(default_factory=list)


class SearchResponse(BaseModel):
    classification: ClassificationResult
    generated_queries: list[str]
    retrieved_chunks: list[RetrievedChunk]


class AnswerRequest(BaseModel):
    user_id: int | None = None
    consultation_id: int | None = None
    save_history: bool = False
    question: str
    category: str | None = None
    region: str | None = None
    user_type: str | None = None
    history: list[dict[str, str]] = Field(default_factory=list)
    law_types: list[str] | None = None
    top_k: int = 5


class AnswerResponse(BaseModel):
    classification: ClassificationResult
    generated_queries: list[str]
    retrieved_chunks: list[RetrievedChunk]
    answer: str
    consultation_id: int | None = None
    user_message_id: int | None = None
    assistant_message_id: int | None = None