from typing import List, Literal, Union

from pydantic import BaseModel

from app.models.common import SourceMetadata


class QueryRequest(BaseModel):
    question: str


class SubQuestionSources(BaseModel):
    sub_question_index: int
    sub_question_text: str
    sources: List[SourceMetadata]


class SubQuestionResult(BaseModel):
    sub_question_index: int
    sub_question_text: str
    answer_section: str


class QueryResponse(BaseModel):
    extracted_questions: List[str]
    answer: str
    sub_question_sources: List[SubQuestionSources] = []
    sources: List[SourceMetadata] = []


class DecomposedEvent(BaseModel):
    phase: Literal["decomposed"]
    extracted_questions: List[str]


class RetrievingEvent(BaseModel):
    phase: Literal["retrieving"]


class FilteringEvent(BaseModel):
    phase: Literal["filtering"]


class GeneratingEvent(BaseModel):
    phase: Literal["generating"]


class GeneratingSubquestionEvent(BaseModel):
    phase: Literal["generating_subquestion"]
    sub_question_index: int
    sub_question_text: str
    answer_section: str = ""


class CompletedEvent(BaseModel):
    phase: Literal["completed"]
    answer: str
    sub_question_sources: List[SubQuestionSources] = []
    sources: List[SourceMetadata] = []


class ErrorEvent(BaseModel):
    phase: Literal["error"]
    message: str


StreamingQueryEvent = Union[
    DecomposedEvent, RetrievingEvent, FilteringEvent,
    GeneratingEvent, GeneratingSubquestionEvent,
    CompletedEvent, ErrorEvent
]