from typing import Literal, Optional

from pydantic import BaseModel, Field


class ChunkHighlightTarget(BaseModel):
    document_id: str
    chunk_index: int
    sub_question_text: str
    sub_question_index: int


class HighlightBatchRequest(BaseModel):
    targets: list[ChunkHighlightTarget]


class RelevantSentence(BaseModel):
    sentence_index: int = Field(description="0-based index of the relevant sentence")
    reason: str = Field(
        description="Brief explanation of why this sentence is relevant",
        max_length=80,
    )


class ChunkHighlights(BaseModel):
    document_id: str
    chunk_index: int
    relevant_sentences: list[RelevantSentence] = Field(default_factory=list)


class HighlightBatchResult(BaseModel):
    results: list[ChunkHighlights]


class HighlightBatchResponse(BaseModel):
    status: Literal["completed", "partial", "failed"]
    cached_count: int = 0
    errors: list[str] = Field(default_factory=list)
    highlight_prompt: Optional[str] = None
    highlight_response_json: Optional[str] = None
    highlight_time_ms: int = 0