from __future__ import annotations

# standard
from dataclasses import dataclass

# third party
# custom
from .types import Cost, Usage


@dataclass
class Model:
    name: str
    provider: str
    display_name: str
    description: str | None = None

    # basic specs
    context_window: int = 0
    max_output_tokens: int | None = None

    # features
    supports_vision: bool = False
    supports_tools: bool = False
    supports_thinking: bool = False
    supports_reasoning_tokens: bool = False

    # pricing
    input_price_per_mtok: float | None = None
    output_price_per_mtok: float | None = None
    cache_read_price_per_mtok: float | None = None
    cache_write_price_per_mtok: float | None = None
    # NOTE xai only
    input_price_per_mtok_128k: float | None = None
    output_price_per_mtok_128k: float | None = None
    # NOTE used by most normal providers
    input_price_per_mtok_200k: float | None = None
    output_price_per_mtok_200k: float | None = None
    cache_read_price_per_mtok_200k: float | None = None
    cache_write_price_per_mtok_200k: float | None = None
    # NOTE hi there we're openai and we love to fk dev experience up
    input_price_per_mtok_272k: float | None = None
    output_price_per_mtok_272k: float | None = None
    cache_read_price_per_mtok_272k: float | None = None

    # dates
    release_date: str | None = None
    deprecated_at: str | None = None
    sunset_at: str | None = None


def _price_tier(usage: Usage, model: Model) -> tuple[float, float, float, float]:
    """Return (p_input, p_output, p_cache_read, p_cache_write) for the correct tier.

    Tiers are selected by total input tokens:
      > 272k → use *_272k rates when available (no 272k cache_write rates defined)
      > 200k → use *_200k rates when available
      > 128k → use *_128k rates when available (no 128k cache rates defined)
      else   → use base rates
    """

    tok = usage.input_tokens
    if tok > 272_000 and model.input_price_per_mtok_272k is not None:
        return (
            model.input_price_per_mtok_272k,
            model.output_price_per_mtok_272k or model.output_price_per_mtok or 0,
            model.cache_read_price_per_mtok_272k
            or model.cache_read_price_per_mtok
            or 0,
            model.cache_write_price_per_mtok or 0,
        )
    if tok > 200_000 and model.input_price_per_mtok_200k is not None:
        return (
            model.input_price_per_mtok_200k,
            model.output_price_per_mtok_200k or model.output_price_per_mtok or 0,
            model.cache_read_price_per_mtok_200k
            or model.cache_read_price_per_mtok
            or 0,
            model.cache_write_price_per_mtok_200k
            or model.cache_write_price_per_mtok
            or 0,
        )
    if tok > 128_000 and model.input_price_per_mtok_128k is not None:
        return (
            model.input_price_per_mtok_128k,
            model.output_price_per_mtok_128k or model.output_price_per_mtok or 0,
            model.cache_read_price_per_mtok or 0,
            model.cache_write_price_per_mtok or 0,
        )
    return (
        model.input_price_per_mtok or 0,
        model.output_price_per_mtok or 0,
        model.cache_read_price_per_mtok or 0,
        model.cache_write_price_per_mtok or 0,
    )


def compute_cost(usage: Usage, model: Model) -> Cost:
    """Compute the cost for a single API response given usage and model pricing.

    API responses include cache_read and cache_write tokens inside input_tokens,
    so regular (non-cached) input = input_tokens - cache_read_tokens - cache_write_tokens.
    Each bucket is billed at its own rate; tier is selected by input token count.
    """

    mtok = 1_000_000
    p_in, p_out, p_cr, p_cw = _price_tier(usage, model)
    regular_input = max(
        0, usage.input_tokens - usage.cache_read_tokens - usage.cache_write_tokens
    )
    input_cost = (regular_input / mtok) * p_in
    output_cost = (usage.output_tokens / mtok) * p_out
    cache_read_cost = (usage.cache_read_tokens / mtok) * p_cr
    cache_write_cost = (usage.cache_write_tokens / mtok) * p_cw
    return Cost(
        input=input_cost,
        output=output_cost,
        cache_read=cache_read_cost,
        cache_write=cache_write_cost,
        total=input_cost + output_cost + cache_read_cost + cache_write_cost,
    )
