"""Session-token evidence ingested from State Hub aggregates."""

from __future__ import annotations

from datetime import date, datetime
from decimal import Decimal
from typing import Literal

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

MEASUREMENT_KINDS = ("measured", "allocated", "estimated", "superseded")
EVIDENCE_BASIS = {
    "measured": "measured",
    "allocated": "derived",
    "estimated": "estimated",
    "superseded": "excluded",
}


class SessionTokenSlice(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tokens_in: int = Field(ge=0)
    tokens_out: int = Field(ge=0)
    event_count: int = Field(ge=0)
    confidence: Decimal | None = None
    method: str | None = None

    @field_validator("confidence")
    @classmethod
    def validate_confidence(cls, value: Decimal | None) -> Decimal | None:
        if value is None:
            return None
        if not value.is_finite() or value < 0 or value > 1:
            raise ValueError("confidence must be between 0 and 1")
        return value


class SessionTokenWork(BaseModel):
    model_config = ConfigDict(extra="forbid")

    unit: Literal["task", "outcome"] = "task"
    count: int = Field(ge=0)
    source: str = Field(min_length=1)


class SessionTokenAssociation(BaseModel):
    model_config = ConfigDict(extra="forbid")

    scope: Literal["repo", "workplan", "model"]
    scope_id: str = Field(min_length=1, max_length=256)
    label: str | None = None
    measurement_kind: Literal["measured", "allocated", "estimated", "superseded"]
    tokens_in: int = Field(ge=0)
    tokens_out: int = Field(ge=0)
    event_count: int = Field(ge=0, default=1)


class SessionTokenEvidence(BaseModel):
    """State Hub session-token snapshot. Not booked spend and not usage_observation."""

    model_config = ConfigDict(extra="forbid")

    schema_version: Literal["0.1"] = "0.1"
    record_type: Literal["session_token_evidence"] = "session_token_evidence"
    record_id: str = Field(min_length=1, max_length=256)
    revision_of: str | None = None
    source_system: Literal["state-hub"] = "state-hub"
    period_start: date
    period_end: date
    captured_at: datetime
    source_evidence: list[str] = Field(min_length=1)
    totals: dict[str, SessionTokenSlice] = Field(default_factory=dict)
    associations: list[SessionTokenAssociation] = Field(default_factory=list)
    work: SessionTokenWork | None = None

    @field_validator("totals")
    @classmethod
    def validate_totals(cls, value: dict[str, SessionTokenSlice]) -> dict[str, SessionTokenSlice]:
        unknown = set(value) - set(MEASUREMENT_KINDS)
        if unknown:
            raise ValueError(f"unknown measurement_kind: {', '.join(sorted(unknown))}")
        return value

    @model_validator(mode="after")
    def validate_period_and_methods(self) -> "SessionTokenEvidence":
        if self.period_end < self.period_start:
            raise ValueError("period_end cannot precede period_start")
        allocated = self.totals.get("allocated")
        if (
            allocated is not None
            and allocated.event_count > 0
            and not (allocated.method and allocated.method.strip())
        ):
            raise ValueError("allocated totals with events require method")
        return self