"""Pydantic response models for the ModelPricing.ai API."""

from typing import Any

from pydantic import BaseModel, Field


class EstimateBreakdown(BaseModel):
    """Cost breakdown for a single direction (input or output).

    @param unit - Pricing unit (e.g. ``"per-1M-input"``).
    @param branch - Pricing tier that matched the model.
    @param qty - Number of tokens.
    @param rate - Per-unit rate in USD.
    @param subtotal - Computed cost for this direction.
    """

    unit: str
    branch: str
    qty: int
    rate: float
    subtotal: float


class EstimateBreakdownGroup(BaseModel):
    """Grouped input/output cost breakdowns.

    @param input - Breakdown for input tokens.
    @param output - Breakdown for output tokens.
    """

    input: EstimateBreakdown
    output: EstimateBreakdown


class EstimateResponse(BaseModel):
    """API response for a cost estimation request.

    @param total - Total estimated cost in USD.
    @param breakdown - Line-item breakdown of input/output costs.
    @param model - Canonical model identifier returned by the server.
    @param traceId - Optional pass-through trace data.
    """

    total: float
    breakdown: EstimateBreakdownGroup
    model: str
    traceId: dict[str, Any] | None = Field(default=None)
