"""Pricing, estimation and the spend ceiling (docs/06-cost.md).

Two jobs. Before a run: say what it will cost and what it will be able to
resolve, so the user can decline to spend $28 learning "inconclusive" —
spending money to learn nothing is the worst outcome and it is avoidable at
estimate time. During a run: enforce a hard ceiling against *actual* reported
costs, not the estimate, because the estimate is the thing most likely to be
wrong.

The agent's own reported `cost_usd` is the source of truth wherever it exists.
The planning figures here are only used before any trial has finished, and the
estimator recalibrates onto observed costs as soon as it has them.
"""

from __future__ import annotations

from dataclasses import dataclass
from datetime import date

#: Per-million-token prices, verified 2026-07-26 against the Anthropic
#: first-party API. Stale prices produce confidently wrong estimates, so the
#: verification date travels with the table and the CLI prints it.
PRICES_VERIFIED = date(2026, 7, 26)


@dataclass(frozen=True)
class ModelPrice:
    model_id: str
    label: str
    input_per_mtok: float
    output_per_mtok: float
    #: Planning cost of one trial before any real data — docs/06-cost.md's
    #: worked example of a mid-sized agentic task at ~8 turns with caching on.
    planning_trial_usd: float
    intro_input_per_mtok: float | None = None
    intro_output_per_mtok: float | None = None
    intro_until: date | None = None

    def rates(self, on: date | None = None) -> tuple[float, float]:
        """Input/output rates, honouring introductory pricing while it lasts."""
        today = on or date.today()
        if self.intro_until and today <= self.intro_until:
            return (
                self.intro_input_per_mtok or self.input_per_mtok,
                self.intro_output_per_mtok or self.output_per_mtok,
            )
        return self.input_per_mtok, self.output_per_mtok


#: Cache reads bill at ~0.1x input; cache writes at 1.25x (5-min TTL) or
#: 2x (1-hour TTL). The runner keeps the invariant prompt prefix byte-identical
#: across trials so reads dominate — getting that wrong silently multiplies the
#: bill roughly 3x, and getting it right is free.
CACHE_READ_MULTIPLIER = 0.1
CACHE_WRITE_MULTIPLIER = 1.25

MODELS: dict[str, ModelPrice] = {
    "claude-opus-5": ModelPrice(
        "claude-opus-5", "Opus 5", 5.0, 25.0, planning_trial_usd=0.90
    ),
    "claude-sonnet-5": ModelPrice(
        "claude-sonnet-5",
        "Sonnet 5",
        3.0,
        15.0,
        planning_trial_usd=0.35,
        intro_input_per_mtok=2.0,
        intro_output_per_mtok=10.0,
        intro_until=date(2026, 8, 31),
    ),
    "claude-haiku-4-5-20251001": ModelPrice(
        "claude-haiku-4-5-20251001", "Haiku 4.5", 1.0, 5.0, planning_trial_usd=0.10
    ),
}

DEFAULT_MODEL = "claude-sonnet-5"

#: `--cheap` maps here: ~3.5x cheaper, and the model the ablation screening
#: stage runs on. Any report generated on a non-target model carries a banner,
#: because config sensitivity may not transfer across model tiers.
CHEAP_MODEL = "claude-haiku-4-5-20251001"

#: docs/06-cost.md. Not a default the tool picks silently — `ablate compare`
#: always shows the estimate, and CI must set this explicitly.
DEFAULT_MAX_SPEND_USD = 25.0

#: Median seconds per trial, for the wall-clock estimate.
PLANNING_SECONDS_PER_TRIAL = 75.0


def resolve_model(model: str | None, *, cheap: bool = False) -> ModelPrice:
    if cheap and model is None:
        return MODELS[CHEAP_MODEL]
    key = model or DEFAULT_MODEL
    if key in MODELS:
        return MODELS[key]
    # An unknown model id is usable — the agent CLI may support models this
    # table has never heard of — but it cannot be priced, so it inherits the
    # default's planning figure and the report says the estimate is a guess.
    return ModelPrice(key, key, *MODELS[DEFAULT_MODEL].rates(), planning_trial_usd=MODELS[DEFAULT_MODEL].planning_trial_usd)


def token_cost_usd(
    price: ModelPrice,
    *,
    input_tokens: int = 0,
    output_tokens: int = 0,
    cache_read: int = 0,
    cache_write: int = 0,
    on: date | None = None,
) -> float:
    """Dollar cost of a token bundle at this model's rates."""
    rate_in, rate_out = price.rates(on)
    return (
        input_tokens * rate_in
        + cache_read * rate_in * CACHE_READ_MULTIPLIER
        + cache_write * rate_in * CACHE_WRITE_MULTIPLIER
        + output_tokens * rate_out
    ) / 1_000_000


def config_cost_per_turn(tokens: int, price: ModelPrice, *, on: date | None = None) -> float:
    """What a config's own text costs on every turn, as a cache read.

    This is the deterministic floor of the whole product (ADR-0010): a config
    is re-sent on every turn of every session forever, and that number is
    arithmetic rather than an experiment. It cannot come back "inconclusive."
    """
    return token_cost_usd(price, cache_read=tokens, on=on)


@dataclass(frozen=True)
class RunEstimate:
    n_runs: int
    per_trial_usd: float
    total_usd: float
    seconds: float
    model: ModelPrice
    calibrated: bool

    @property
    def minutes(self) -> float:
        return self.seconds / 60.0


def estimate_run(
    n_runs: int,
    price: ModelPrice,
    *,
    observed_costs: list[float] | None = None,
    concurrency: int = 4,
) -> RunEstimate:
    """Projected cost and wall clock for `n_runs` trials.

    Once trials start reporting real costs the estimate switches to their mean.
    Observed cost on real tasks ranges roughly 0.4x-2.5x the planning figure,
    so recalibrating early matters more than the table being precise.

    Pass warm-cache costs where you have them (Q7): the first trial of each
    (task, arm) pays a prompt-cache write and is systematically dearer, so
    recalibrating on a blended mean overestimates the trials still to come.
    """
    usable = [c for c in (observed_costs or []) if c is not None and c > 0]
    calibrated = len(usable) >= 3
    per_trial = sum(usable) / len(usable) if calibrated else price.planning_trial_usd

    return RunEstimate(
        n_runs=n_runs,
        per_trial_usd=per_trial,
        total_usd=per_trial * n_runs,
        seconds=PLANNING_SECONDS_PER_TRIAL * n_runs / max(concurrency, 1),
        model=price,
        calibrated=calibrated,
    )


class SpendCeilingExceeded(Exception):
    """Raised when the running total of actual costs breaches the ceiling.

    The run stops cleanly rather than being killed: because the work queue was
    shuffled, whatever completed is a roughly balanced partial design and is
    still analyzable.
    """

    def __init__(self, spent: float, ceiling: float, completed: int, total: int) -> None:
        self.spent = spent
        self.ceiling = ceiling
        self.completed = completed
        self.total = total
        super().__init__(
            f"spend ceiling reached: ${spent:.2f} of ${ceiling:.2f} after "
            f"{completed}/{total} trials"
        )


class SpendTracker:
    """Running total against a hard ceiling, checked after every trial."""

    def __init__(self, ceiling: float | None, total_trials: int) -> None:
        self.ceiling = ceiling
        self.total_trials = total_trials
        self.spent = 0.0
        self.completed = 0
        self.unpriced = 0

    def record(self, cost_usd: float | None) -> None:
        self.completed += 1
        if cost_usd is None:
            # An agent that does not report cost cannot be budgeted against.
            # Counted and surfaced rather than assumed to be free.
            self.unpriced += 1
            return
        self.spent += cost_usd

    def check(self) -> None:
        if self.ceiling is not None and self.spent >= self.ceiling:
            raise SpendCeilingExceeded(
                self.spent, self.ceiling, self.completed, self.total_trials
            )

    @property
    def remaining(self) -> float | None:
        return None if self.ceiling is None else max(self.ceiling - self.spent, 0.0)
