"""Tests for pricing, estimation and the spend ceiling.

The ceiling is a safety feature: a runaway loop bills a real credit card, and
in CI it does so unattended. These tests treat it as one.
"""

from __future__ import annotations

from datetime import date

import pytest

from ablate.config.tokens import estimate_tokens
from ablate.cost import (
    CHEAP_MODEL,
    MODELS,
    SpendCeilingExceeded,
    SpendTracker,
    config_cost_per_turn,
    estimate_run,
    resolve_model,
    token_cost_usd,
)

SONNET = MODELS["claude-sonnet-5"]
HAIKU = MODELS[CHEAP_MODEL]


# --- pricing ----------------------------------------------------------------


def test_introductory_pricing_applies_before_it_expires():
    assert SONNET.rates(date(2026, 8, 1)) == (2.0, 10.0)


def test_list_pricing_applies_after_the_intro_window():
    assert SONNET.rates(date(2026, 9, 1)) == (3.0, 15.0)


def test_cache_reads_are_an_order_of_magnitude_cheaper_than_fresh_input():
    fresh = token_cost_usd(SONNET, input_tokens=1_000_000, on=date(2026, 9, 1))
    cached = token_cost_usd(SONNET, cache_read=1_000_000, on=date(2026, 9, 1))
    assert cached == pytest.approx(fresh * 0.1)


def test_cache_writes_cost_more_than_fresh_input():
    fresh = token_cost_usd(SONNET, input_tokens=1_000_000, on=date(2026, 9, 1))
    written = token_cost_usd(SONNET, cache_write=1_000_000, on=date(2026, 9, 1))
    assert written > fresh


def test_config_cost_per_turn_is_a_cache_read():
    assert config_cost_per_turn(4_000, SONNET, on=date(2026, 9, 1)) == pytest.approx(
        4_000 * 3.0 * 0.1 / 1_000_000
    )


def test_cheap_resolves_to_haiku():
    assert resolve_model(None, cheap=True).model_id == CHEAP_MODEL


def test_explicit_model_beats_cheap():
    assert resolve_model("claude-opus-5", cheap=True).model_id == "claude-opus-5"


def test_unknown_model_is_usable_but_priced_off_the_default():
    """The agent CLI may support models this table has never heard of."""
    price = resolve_model("some-future-model")
    assert price.model_id == "some-future-model"
    assert price.planning_trial_usd == SONNET.planning_trial_usd


def test_haiku_is_meaningfully_cheaper_than_sonnet():
    assert HAIKU.planning_trial_usd < SONNET.planning_trial_usd / 3


# --- estimation -------------------------------------------------------------


def test_estimate_uses_planning_figures_before_any_data():
    estimate = estimate_run(80, SONNET)
    assert not estimate.calibrated
    assert estimate.total_usd == pytest.approx(80 * 0.35)


def test_estimate_recalibrates_onto_observed_costs():
    """Real per-trial cost ranges 0.4x-2.5x the planning figure, so switching
    to measured costs matters more than the table being precise."""
    estimate = estimate_run(80, SONNET, observed_costs=[0.70, 0.72, 0.68])
    assert estimate.calibrated
    assert estimate.per_trial_usd == pytest.approx(0.70)
    assert estimate.total_usd == pytest.approx(56.0)


def test_estimate_waits_for_enough_samples_before_recalibrating():
    assert not estimate_run(80, SONNET, observed_costs=[0.70]).calibrated


def test_estimate_ignores_missing_and_zero_costs():
    assert not estimate_run(80, SONNET, observed_costs=[None, 0.0, None]).calibrated


def test_wall_clock_estimate_scales_down_with_concurrency():
    serial = estimate_run(80, SONNET, concurrency=1)
    parallel = estimate_run(80, SONNET, concurrency=4)
    assert parallel.minutes == pytest.approx(serial.minutes / 4)


# --- the ceiling ------------------------------------------------------------


def test_tracker_raises_once_the_ceiling_is_reached():
    tracker = SpendTracker(ceiling=1.0, total_trials=10)
    for _ in range(4):
        tracker.record(0.20)
        tracker.check()

    tracker.record(0.20)
    with pytest.raises(SpendCeilingExceeded) as excinfo:
        tracker.check()
    assert excinfo.value.completed == 5
    assert excinfo.value.spent == pytest.approx(1.0)


def test_tracker_without_a_ceiling_never_raises():
    tracker = SpendTracker(ceiling=None, total_trials=3)
    tracker.record(1000.0)
    tracker.check()
    assert tracker.remaining is None


def test_unpriced_trials_are_counted_not_assumed_free():
    """An agent that reports no cost cannot be budgeted against, and pretending
    those trials were free would silently disable the ceiling."""
    tracker = SpendTracker(ceiling=5.0, total_trials=4)
    tracker.record(None)
    tracker.record(0.5)
    assert tracker.unpriced == 1
    assert tracker.completed == 2
    assert tracker.spent == pytest.approx(0.5)


def test_remaining_budget_never_goes_negative():
    tracker = SpendTracker(ceiling=1.0, total_trials=2)
    tracker.record(5.0)
    assert tracker.remaining == 0.0


# --- token estimation -------------------------------------------------------


def test_empty_config_costs_nothing():
    assert estimate_tokens("") == 0
    assert estimate_tokens(None) == 0


def test_token_estimate_is_deterministic():
    text = "# Conventions\n\nAlways run the tests.\n"
    assert estimate_tokens(text) == estimate_tokens(text)


def test_code_counts_denser_than_prose():
    """Same character count, different tokenization density."""
    prose = "a" * 400
    code = "```\n" + "a" * 400 + "\n```"
    assert estimate_tokens(code) > estimate_tokens(prose)


def test_token_estimate_is_in_the_documented_ballpark():
    """docs/05-statistics.md: a 400-line CLAUDE.md is ~4,000 tokens."""
    config = "\n".join("Some guidance about how to write code here." for _ in range(400))
    assert 3_000 <= estimate_tokens(config) <= 5_500


# --- steady-state cost (Q7) -------------------------------------------------


def _priced(task: str, arm: str, started: str, cost: float):
    from ablate.record import AssertionRecord, TelemetryRecord, TrialRecord

    return TrialRecord(
        trial_id=f"{task}-{arm}-{started}",
        run_id="r",
        task_id=task,
        variant_id=arm,
        variant_hash="sha256:x",
        replicate=1,
        agent="claude-code",
        started_at=started,
        duration_s=1.0,
        telemetry=TelemetryRecord(input_tokens=10, turns=2, cost_usd=cost),
        assertions=[AssertionRecord(id="a", type="exit_code", passed=True)],
        passed=True,
    )


def test_steady_state_cost_excludes_the_cache_warming_trial():
    """The first trial of each (task, arm) pays a prompt-cache write. Averaging
    it in overstates what the next run will cost."""
    from ablate.report.analysis import analyze

    records = [
        _priced("t1", "A", "2026-07-27T00:00:00Z", 1.00),  # cold
        _priced("t1", "A", "2026-07-27T00:01:00Z", 0.10),
        _priced("t1", "A", "2026-07-27T00:02:00Z", 0.10),
        _priced("t1", "B", "2026-07-27T00:03:00Z", 1.00),  # cold
        _priced("t1", "B", "2026-07-27T00:04:00Z", 0.10),
    ]
    analysis = analyze(records, ("A", "B"))

    assert analysis.total_cost_usd == pytest.approx(2.30)
    assert analysis.steady_state_cost_usd == pytest.approx(0.10)


def test_steady_state_uses_chronological_order_not_replicate_number():
    """The queue is shuffled, so the trial that warms the cache is whichever
    ran first — not the one labelled replicate 1."""
    from ablate.report.analysis import analyze

    records = [
        _priced("t1", "A", "2026-07-27T00:05:00Z", 0.10),
        _priced("t1", "A", "2026-07-27T00:01:00Z", 1.00),  # earliest, so cold
        _priced("t1", "B", "2026-07-27T00:02:00Z", 1.00),
        _priced("t1", "B", "2026-07-27T00:06:00Z", 0.10),
    ]
    assert analyze(records, ("A", "B")).steady_state_cost_usd == pytest.approx(0.10)


def test_steady_state_is_none_when_every_cell_has_one_trial():
    from ablate.report.analysis import analyze

    records = [
        _priced("t1", "A", "2026-07-27T00:00:00Z", 1.0),
        _priced("t1", "B", "2026-07-27T00:01:00Z", 1.0),
    ]
    assert analyze(records, ("A", "B")).steady_state_cost_usd is None


def test_steady_state_is_none_when_no_cost_is_reported():
    """An agent that reports no cost cannot be budgeted against at all."""
    from ablate.record import AssertionRecord, TelemetryRecord, TrialRecord
    from ablate.report.analysis import analyze

    def unpriced(arm: str, started: str):
        return TrialRecord(
            trial_id=f"{arm}-{started}", run_id="r", task_id="t1", variant_id=arm,
            variant_hash="h", replicate=1, agent="codex", started_at=started,
            duration_s=1.0, telemetry=TelemetryRecord(turns=2, cost_usd=None),
            assertions=[AssertionRecord(id="a", type="exit_code", passed=True)], passed=True,
        )

    records = [unpriced("A", "2026-07-27T00:00:00Z"), unpriced("B", "2026-07-27T00:01:00Z")]
    assert analyze(records, ("A", "B")).steady_state_cost_usd is None
