from __future__ import annotations

from typing import Literal

from pydantic import Field

from app.models.common import StrictModel
from app.models.context import Citation
from app.models.internal_retrieve import CodeSnippet, DocExcerpt, RelatedCommit
from app.models.intent import QueryIntent


class SSEStartPayload(StrictModel):
    request_id: str
    repo_id: str
    snapshot_id: str
    intent: QueryIntent
    answer_mode: Literal["rag", "general"] | None = None


class SSETokenPayload(StrictModel):
    delta: str


class SSESourcesPayload(StrictModel):
    code_snippets: list[CodeSnippet] = Field(default_factory=list)
    doc_excerpts: list[DocExcerpt] = Field(default_factory=list)
    related_commits: list[RelatedCommit] = Field(default_factory=list)
    citations: list[Citation] = Field(default_factory=list)


class SSEDonePayload(StrictModel):
    latency_ms: int
    model: str | None = None
    answer: str
    answer_length: int
    answer_mode: Literal["rag", "general"] | None = None


class SSEErrorPayload(StrictModel):
    message: str
    status_code: int | None = None