"""Responses preparation contracts, retaining the original historical wire format."""

from datetime import datetime
from typing import Annotated, Literal
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from app.schemas.catalog import Input
from app.schemas.indexing import ChunkProvenance


class CompletionOptions(Input):
    """Read-only options for Chat Completions snapshots created before Responses."""

    model: str = Field(pattern=r"^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$")
    output_tokens: int = Field(512, ge=1, le=4096)
    token_parameter: Literal["max_completion_tokens", "max_tokens"] = "max_completion_tokens"
    instruction_role: Literal["system", "developer"] = "system"
    temperature: float | None = Field(None, ge=0, le=2)


class ResponsesOptions(Input):
    protocol: Literal["responses"]
    provider_profile: Literal["compatible", "openai", "deepseek", "bailian"] = "compatible"
    model: str = Field(pattern=r"^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$")
    output_tokens: int = Field(512, ge=16, le=4096)
    temperature: float | None = Field(None, ge=0, le=2)

    @model_validator(mode="after")
    def supported_temperature(self):
        if self.provider_profile == "bailian" and self.temperature == 2:
            raise ValueError("bailian temperature must be less than 2")
        return self


class PrepareContext(Input):
    retrieval_run_id: UUID
    completion: ResponsesOptions
    max_input_characters: int = Field(10000, ge=1024, le=24000)
    max_evidence: int = Field(5, ge=1, le=10)

    @field_validator("completion", mode="before")
    @classmethod
    def default_new_protocol(cls, value):
        if isinstance(value, dict):
            return {"protocol": "responses", **value}
        return value


class Evidence(ChunkProvenance):
    citation_id: str
    chunk_id: UUID
    chunk_ordinal: int
    retrieval_rank: int
    text: str
    text_sha256: str
    start_offset: int
    end_offset: int


class ExcludedEvidence(BaseModel):
    chunk_id: UUID
    reason: Literal["overlap", "budget", "evidence_limit"]


class ContextRead(BaseModel):
    id: UUID
    kind: Literal["prepared_request"] = "prepared_request"
    model_called: Literal[False] = False
    status: Literal["prepared", "insufficient_evidence"]
    reason: Literal["candidates_available", "no_candidates", "context_budget"]
    retrieval_run_id: UUID
    retrieval_mode: Literal["keyword_chunks", "vector", "hybrid"]
    query: str
    text_index_version_id: UUID
    text_generation: int
    catalog_revision: int
    vector_version_id: UUID | None
    vector_generation: int | None
    corpus_fingerprint: str
    prompt_version: Literal["nalati-evidence-v1"] = "nalati-evidence-v1"
    completion: ResponsesOptions | CompletionOptions
    max_input_characters: int
    input_characters: int
    max_evidence: int
    evidence: list[Evidence]
    excluded: list[ExcludedEvidence]
    request_body: dict | None
    request_sha256: str | None
    created_at: datetime


class CompletionUsage(BaseModel):
    prompt_tokens: int = Field(ge=0, strict=True)
    completion_tokens: int = Field(ge=0, strict=True)
    total_tokens: int = Field(ge=0, strict=True)


class ParsedCompletion(BaseModel):
    status: Literal["answer_candidate", "insufficient_evidence", "incomplete", "refused"]
    completion_id: str
    model: str
    answer: str | None
    citations: list[str]
    usage: CompletionUsage | None


TokenCount = Annotated[int, Field(ge=0, strict=True)]


class ResponseUsage(BaseModel):
    model_config = ConfigDict(extra="allow")

    input_tokens: TokenCount
    output_tokens: TokenCount
    total_tokens: TokenCount
    input_tokens_details: dict[str, TokenCount] | None = None
    output_tokens_details: dict[str, TokenCount] | None = None

    @model_validator(mode="after")
    def consistent_counts(self):
        if self.total_tokens != self.input_tokens + self.output_tokens:
            raise ValueError("inconsistent usage")
        for details, count in [
            (self.input_tokens_details, self.input_tokens),
            (self.output_tokens_details, self.output_tokens),
        ]:
            if details and any(value > count for value in details.values()):
                raise ValueError("invalid usage details")
        return self


class ParsedResponse(BaseModel):
    status: Literal["answer_candidate", "insufficient_evidence", "incomplete", "refused", "failed"]
    response_id: str
    model: str
    answer: str | None
    citations: list[str]
    usage: ResponseUsage | None