from __future__ import annotations

from pydantic import Field

from app.models.common import StrictModel
from app.models.internal_retrieve import SourceType


class Citation(StrictModel):
    index: int
    record_id: str
    source_type: SourceType
    score: float
    file_path: str | None = None
    doc_path: str | None = None
    commit_sha: str | None = None
    symbol_name: str | None = None
    section_title: str | None = None


class ContextChunk(StrictModel):
    index: int
    source_type: SourceType
    record_id: str
    label: str
    text: str
    score: float
    file_path: str | None = None
    doc_path: str | None = None
    commit_sha: str | None = None
    symbol_name: str | None = None
    section_title: str | None = None


class BuiltContext(StrictModel):
    chunks: list[ContextChunk] = Field(default_factory=list)
    prompt_text: str = ""
    citations: list[Citation] = Field(default_factory=list)
    total_tokens: int = 0
    trimmed_count: int = 0