"""Public list pricing tables and per-provider cost helpers.

Pricing is matched by substring against the model id recorded in the session
transcript, so keys must reflect the *actual* model id form emitted by each
provider:

  * Anthropic — `claude-fable-5*`, `claude-opus-5*`, `claude-opus-4-*`, `claude-sonnet-5*`, `claude-sonnet-4-*`,
    `claude-haiku-4-5-*`, `claude-3-5-sonnet-*`, `claude-3-5-haiku-*`,
    `claude-3-opus-*`, `claude-3-haiku-*`.
  * OpenAI / Codex — `gpt-5*`, `gpt-4o*`, `gpt-4*`.
  * Google / Gemini — `gemini-3.6-flash*`, `gemini-3.5-flash*`, `gemini-3.1-pro*`, `gemini-3-flash*`,
    `gemini-2.5-pro*`, `gemini-2.5-flash*`, `gemini-2.0-flash*`.

Matching prefers the longest key that is a substring of the model id (most
specific wins), so table ordering does not affect the result. Update when
providers change list pricing.

Sources (last verified 2026-08-03, public list prices, USD per 1M tokens):
  * Anthropic: https://www.anthropic.com/pricing
  * OpenAI:    https://openai.com/api/pricing
  * Google:    https://ai.google.dev/gemini-api/docs/pricing
"""
from __future__ import annotations

from okstra_ctl.models import catalog_pricing


# Anthropic billing ratios relative to base input: cache_creation (5m) = 1.25x,
# cache_creation (1h) = 2x, cache_read = 0.1x, output = 5x. The CLAUDE_PRICING
# entries below carry the 5m tier; the 1h price is derived as base_input * 2x
# at call time so the table stays compact.
CLAUDE_PRICING = {
    # model substring -> (input, cache_creation_5m, cache_read, output) USD/1M.
    #
    # _match_pricing picks the longest matching key (most specific), so ordering
    # here is for readability, not correctness.
    #
    # For the newer 4.x point releases (Opus 4.7, Sonnet 4.6, Haiku 4.5),
    # Anthropic's public price page only lists input/output. Cache-write and
    # cache-read are filled in using Anthropic's published billing ratios
    # (5m cache_creation = 1.25x input, cache_read = 0.1x input), which have
    # been consistent across the Claude 3 / 4 families.

    # Claude 3 series (legacy).
    "3-7-sonnet": (3.0, 3.75, 0.30, 15.0),     # Sonnet 3.7
    "3-5-sonnet": (3.0, 3.75, 0.30, 15.0),     # Sonnet 3.5
    "3-5-haiku":  (0.80, 1.0, 0.08, 4.0),      # Haiku 3.5
    "3-opus":     (15.0, 18.75, 1.50, 75.0),   # Opus 3
    "3-sonnet":   (3.0, 3.75, 0.30, 15.0),     # legacy 3 Sonnet
    "3-haiku":    (0.25, 0.30, 0.03, 1.25),    # Haiku 3

    # Claude Fable 5 (tier above Opus).
    "fable-5":    (10.0, 12.5, 1.0, 50.0),     # Fable 5 (cache prices derived from ratios)

    # Claude Opus 5.
    "opus-5":     (5.0, 6.25, 0.50, 25.0),     # Opus 5 (cache prices derived from ratios)

    # Claude Sonnet 5.
    "sonnet-5":   (3.0, 3.75, 0.30, 15.0),     # Sonnet 5 (cache prices derived from ratios)

    # Claude 4 point releases (explicit so future divergence is easy to see).
    "opus-4-8":   (5.0, 6.25, 0.50, 25.0),     # Opus 4.8 (cache prices derived from ratios)
    "opus-4-7":   (5.0, 6.25, 0.50, 25.0),     # Opus 4.7 (cache prices derived from ratios)
    "opus-4-6":   (5.0, 6.25, 0.50, 25.0),     # Opus 4.6 (legacy; pricing matches 4.7 per Anthropic)
    "sonnet-4-6": (3.0, 3.75, 0.30, 15.0),     # Sonnet 4.6 (cache prices derived from ratios)
    "haiku-4-5":  (1.0, 1.25, 0.10, 5.0),      # Haiku 4.5  (cache prices derived from ratios)

    # Claude 4 family fallbacks (Opus 4 / Sonnet 4 / Haiku 4 base).
    "opus-4":     (15.0, 18.75, 1.50, 75.0),
    "sonnet-4":   (3.0, 3.75, 0.30, 15.0),
    "haiku-4":    (1.0, 1.25, 0.10, 5.0),
}

# Anthropic 1h ephemeral cache_creation multiplier on the base input rate.
CLAUDE_CACHE_CREATE_1H_MULT = 2.0

# Legacy / non-selectable codex prices. Selectable-model prices (the aliases the
# picker offers) live in the model catalog and are merged in below via
# `catalog_pricing("codex")`. Substring match order is now longest-key-first
# (see `_match_pricing`), so this dict's ordering no longer affects matching.
# For models with no published cached-input rate (o1-pro, o3-pro), cached is set
# equal to input as a conservative no-discount default.
_LEGACY_CODEX_PRICING = {
    # model substring -> (input USD/1M, cached_input USD/1M, output USD/1M).

    # GPT-5 series.
    "gpt-5.2-pro":  (21.0,  2.10,  168.0),
    "gpt-5.1":      (1.25,  0.125, 10.0),
    "gpt-5-mini":   (0.25,  0.025, 2.00),
    "gpt-5-nano":   (0.05,  0.005, 0.40),
    "gpt-5":        (1.25,  0.125, 10.0),  # base GPT-5 (also matches gpt-5-codex)

    # O-series reasoning models.
    "o1-pro":  (150.0, 150.0, 600.0),  # no cached rate published
    "o3-pro":  (20.0,  20.0,  80.0),   # no cached rate published
    "o4-mini": (1.10,  0.275, 4.40),
    "o3-mini": (1.10,  0.275, 4.40),
    "o1":      (15.0,  7.50,  60.0),
    "o3":      (2.00,  1.00,  8.00),

    # GPT-4 series.
    "gpt-4.1-nano": (0.10, 0.01,  0.40),
    "gpt-4.1-mini": (0.40, 0.04,  1.60),
    "gpt-4.1":      (2.00, 0.20,  8.00),
    "gpt-4o-mini":  (0.15, 0.075, 0.60),
    "gpt-4o":       (2.50, 1.25,  10.0),
    "gpt-4":        (2.50, 0.625, 10.0),  # legacy gpt-4 fallback
}

# catalog + legacy: catalog owns selectable prices, legacy covers the rest. No
# key overlaps today; if one ever collided, the legacy (right) side would win.
CODEX_PRICING = {**catalog_pricing("codex"), **_LEGACY_CODEX_PRICING}
GROK_PRICING = catalog_pricing("grok")
KIMI_PRICING = catalog_pricing("kimi")

GEMINI_PRICING = {
    # model substring -> (input USD/1M, output USD/1M).
    #
    # Cached-input prices exist for some models but are not separately priced
    # here because the Gemini transcript collector does not yet record cached
    # input tokens. Models with two-tier context pricing (Gemini 2.5 Pro,
    # Gemini 3.1 Pro) are charged at the ≤200K rate; runs above 200K input
    # will be slightly undercounted.
    #
    # Both dotted (`gemini-3.1-pro`) and hyphenated (`gemini-3-1-pro`) id
    # forms appear in the wild, so include both for the new 3.x families.

    # Gemini 3 series.
    "3.6-flash":       (1.50, 7.50),
    "3-6-flash":       (1.50, 7.50),
    "3.5-flash":       (1.50, 9.00),
    "3-5-flash":       (1.50, 9.00),
    "3.1-pro":         (2.00, 12.0),
    "3-1-pro":         (2.00, 12.0),
    "3-flash":         (0.50, 3.00),

    # Gemini 2.5 series.
    "2.5-flash-lite":  (0.10, 0.40),
    "2.5-flash":       (0.30, 2.50),
    "2.5-pro":         (1.25, 10.0),

    # Gemini 2.0 series.
    "2.0-flash-lite":  (0.075, 0.30),
    "2.0-flash":       (0.10,  0.40),

    # Fallbacks for unspecified family names.
    "flash-lite":      (0.10, 0.40),   # assume 2.5 Flash-Lite
    "pro":             (1.25, 10.0),   # assume 2.5 Pro
    "flash":           (0.30, 2.50),   # assume 2.5 Flash
    "auto":            (1.25, 10.0),   # treat unknown/auto as 2.5 Pro
}


def _match_pricing(model: str | None, table: dict) -> tuple | None:
    if not model:
        return None
    m = model.lower()
    for key, val in sorted(table.items(), key=lambda kv: len(kv[0]), reverse=True):
        if key in m:
            return val
    return None


def claude_billable_equivalent(
    input_t: int,
    cache_create_t: int,
    cache_read_t: int,
    output_t: int,
    cache_create_1h_t: int = 0,
) -> int:
    """Sum normalized to base-input units.

    Ratios: cache_creation_5m=1.25x, cache_creation_1h=2x, cache_read=0.1x,
    output=5x. `cache_create_t` is the total cache_creation tokens; pass the
    1h portion separately via `cache_create_1h_t` so the 5m vs 1h tiers are
    weighted correctly (the 5m portion is the difference).
    """
    cc_1h = max(0, cache_create_1h_t)
    cc_5m = max(0, cache_create_t - cc_1h)
    return int(round(
        input_t
        + 1.25 * cc_5m
        + CLAUDE_CACHE_CREATE_1H_MULT * cc_1h
        + 0.1 * cache_read_t
        + 5.0 * output_t
    ))


def claude_cost_usd(
    model: str | None,
    input_t: int,
    cache_create_t: int,
    cache_read_t: int,
    output_t: int,
    cache_create_1h_t: int = 0,
) -> float | None:
    p = _match_pricing(model, CLAUDE_PRICING)
    if p is None:
        return None
    pi, pcc, pcr, po = p
    cc_1h = max(0, cache_create_1h_t)
    cc_5m = max(0, cache_create_t - cc_1h)
    pcc_1h = pi * CLAUDE_CACHE_CREATE_1H_MULT
    return round((
        input_t * pi
        + cc_5m * pcc
        + cc_1h * pcc_1h
        + cache_read_t * pcr
        + output_t * po
    ) / 1_000_000, 4)


def codex_cost_usd(model: str | None, input_t: int, cached_input_t: int, output_t: int) -> float | None:
    p = _match_pricing(model, CODEX_PRICING)
    if p is None:
        return None
    pi, pci, po = p
    # Codex `input_tokens` already includes cached; subtract cached and price the rest at full input.
    fresh = max(0, input_t - cached_input_t)
    return round((fresh * pi + cached_input_t * pci + output_t * po) / 1_000_000, 4)


def provider_cost_usd(
    provider: str,
    model: str | None,
    input_t: int,
    cached_input_t: int,
    output_t: int,
) -> float | None:
    """Price a three-rate provider without fabricating an unknown model rate."""
    tables = {
        "codex": CODEX_PRICING,
        "grok": GROK_PRICING,
        "kimi": KIMI_PRICING,
    }
    pricing = _match_pricing(model, tables.get(provider, {}))
    if pricing is None:
        return None
    input_price, cached_price, output_price = pricing
    fresh_input = max(0, input_t - cached_input_t)
    return round((
        fresh_input * input_price
        + cached_input_t * cached_price
        + output_t * output_price
    ) / 1_000_000, 4)


def antigravity_cost_usd(model: str | None, input_t: int, output_t: int) -> float | None:
    p = _match_pricing(model, GEMINI_PRICING)
    if p is None:
        return None
    pi, po = p
    return round((input_t * pi + output_t * po) / 1_000_000, 4)