"""Metrics model for ATIF trajectories."""

from typing import Any

from pydantic import BaseModel, Field


class Metrics(BaseModel):
    """Token usage and cost statistics for a single step (ATIF v1.6).

    Core ATIF fields use int = 0 defaults in VibeLens for safe arithmetic
    without None-guards. ATIF spec uses Optional[int] = None.
    """

    prompt_tokens: int = Field(
        default=0, description="Total input tokens including cached and non-cached."
    )
    completion_tokens: int = Field(
        default=0, description="Total tokens generated by the LLM response."
    )
    cached_tokens: int = Field(
        default=0, description="Subset of prompt_tokens that were cache hits."
    )
    cache_creation_tokens: int = Field(
        default=0,
        description="[VibeLens] Tokens written into the prompt cache (Anthropic-specific).",
    )
    cost_usd: float | None = Field(
        default=None, description="Monetary cost of the API call in USD."
    )
    prompt_token_ids: list[int] | None = Field(
        default=None, description="Token IDs for prompt tokens sent to the LLM (ATIF v1.4)."
    )
    completion_token_ids: list[int] | None = Field(
        default=None,
        description="Token IDs for completion tokens, enabling RL training (ATIF v1.3).",
    )
    logprobs: list[float] | None = Field(
        default=None, description="Log probability assigned to each generated token."
    )
    extra: dict[str, Any] | None = Field(
        default=None, description="Custom metrics not covered by ATIF core fields."
    )
