"""Typed, serializable records produced by baseline runs."""

import math
from typing import Annotated, Any, Literal

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    FiniteFloat,
    StrictInt,
    field_validator,
    model_validator,
)


class StrictModel(BaseModel):
    """Base model that rejects fields not in the public artifact contract."""

    model_config = ConfigDict(extra="forbid")


NonNegativeToken = Annotated[StrictInt, Field(ge=0)]


class SampleRecord(StrictModel):
    sample_id: StrictInt = Field(ge=0)
    text: str = Field(min_length=1)
    token_ids: list[Annotated[StrictInt, Field(ge=0)]] = Field(min_length=1)
    seed: StrictInt
    generation_seconds: Annotated[FiniteFloat, Field(ge=0)]

    @field_validator("text")
    @classmethod
    def require_non_whitespace_text(cls, value: str) -> str:
        if not value.strip():
            raise ValueError("text must not be empty or whitespace")
        return value

    @field_validator("generation_seconds", mode="before")
    @classmethod
    def require_builtin_finite_float(cls, value: object) -> object:
        if type(value) is not float or not math.isfinite(value) or value < 0:
            raise ValueError("generation_seconds must be a finite non-negative built-in float")
        return value


class ConditionalSampleRecord(StrictModel):
    """One fixed-prefix conditional sample publication record."""

    sample_id: StrictInt = Field(ge=0)
    prompt_id: StrictInt = Field(ge=0, le=1023)
    completion_id: StrictInt = Field(ge=0, le=4)
    source_index: StrictInt = Field(ge=0)
    prefix_token_ids: list[NonNegativeToken]
    continuation_token_ids: list[NonNegativeToken] = Field(min_length=1)
    reference_token_ids: list[NonNegativeToken]
    full_token_ids: list[NonNegativeToken] = Field(min_length=65)
    prefix_text: str
    continuation_text: str
    reference_text: str
    full_text: str
    seed: StrictInt
    generation_seconds: Annotated[FiniteFloat, Field(ge=0)]
    prefix_exact_match: Literal[True]

    @model_validator(mode="after")
    def validate_slices(self) -> "ConditionalSampleRecord":
        if len(self.prefix_token_ids) != 64 or len(self.reference_token_ids) != 64:
            raise ValueError("prefix and reference must each contain 64 tokens")
        if self.full_token_ids[:64] != self.prefix_token_ids:
            raise ValueError("full token prefix differs from fixed prefix")
        if self.full_token_ids[64:] != self.continuation_token_ids:
            raise ValueError("continuation differs from full token suffix")
        return self


class RunMetadata(StrictModel):
    run_id: str = Field(min_length=1)
    model_id: str = Field(min_length=1)
    dataset_id: str = Field(min_length=1)
    step_count: int = Field(gt=0)
    seed: int
    command: list[str] = Field(min_length=1)
    config_sha256: str = Field(min_length=1)


class MetricRecord(StrictModel):
    run_id: str = Field(min_length=1)
    metric: str = Field(min_length=1)
    value: float
    sample_count: int = Field(ge=0)


class FailureRecord(StrictModel):
    run_id: str = Field(min_length=1)
    stage: str = Field(min_length=1)
    message: str = Field(min_length=1)
    details: dict[str, Any] = Field(default_factory=dict)