from __future__ import annotations

# custom
from sunwaee.modules.gen.engine.model import Model, compute_cost
from sunwaee.modules.gen.engine.models import get_model, list_models
from sunwaee.modules.gen.engine.types import Cost, Usage


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _model(
    input_price: float = 1.0,
    output_price: float = 2.0,
    cache_read: float = 0.1,
    cache_write: float = 0.2,
    input_128k: float | None = None,
    output_128k: float | None = None,
    input_200k: float | None = None,
    output_200k: float | None = None,
    cache_read_200k: float | None = None,
    cache_write_200k: float | None = None,
) -> Model:
    return Model(
        name="test-model",
        provider="test",
        display_name="Test Model",
        context_window=200_000,
        input_price_per_mtok=input_price,
        output_price_per_mtok=output_price,
        cache_read_price_per_mtok=cache_read,
        cache_write_price_per_mtok=cache_write,
        input_price_per_mtok_128k=input_128k,
        output_price_per_mtok_128k=output_128k,
        input_price_per_mtok_200k=input_200k,
        output_price_per_mtok_200k=output_200k,
        cache_read_price_per_mtok_200k=cache_read_200k,
        cache_write_price_per_mtok_200k=cache_write_200k,
    )


# ---------------------------------------------------------------------------
# compute_cost — base tier (≤128k tokens)
# ---------------------------------------------------------------------------


def test_compute_cost_base_tier_no_cache():
    model = _model(input_price=1.0, output_price=2.0, cache_read=0.0, cache_write=0.0)
    usage = Usage(input_tokens=1_000_000, output_tokens=500_000)
    cost = compute_cost(usage, model)
    assert isinstance(cost, Cost)
    assert abs(cost.input - 1.0) < 1e-9
    assert abs(cost.output - 1.0) < 1e-9
    assert abs(cost.cache_read - 0.0) < 1e-9
    assert abs(cost.cache_write - 0.0) < 1e-9
    assert abs(cost.total - 2.0) < 1e-9


def test_compute_cost_base_tier_with_cache():
    model = _model(input_price=3.0, output_price=15.0, cache_read=0.3, cache_write=3.75)
    # 100k input, 20k cache_read, 10k cache_write → regular = 70k
    usage = Usage(
        input_tokens=100_000,
        output_tokens=50_000,
        cache_read_tokens=20_000,
        cache_write_tokens=10_000,
    )
    cost = compute_cost(usage, model)
    expected_input = (70_000 / 1_000_000) * 3.0
    expected_output = (50_000 / 1_000_000) * 15.0
    expected_cr = (20_000 / 1_000_000) * 0.3
    expected_cw = (10_000 / 1_000_000) * 3.75
    assert abs(cost.input - expected_input) < 1e-9
    assert abs(cost.output - expected_output) < 1e-9
    assert abs(cost.cache_read - expected_cr) < 1e-9
    assert abs(cost.cache_write - expected_cw) < 1e-9
    assert abs(cost.total - (expected_input + expected_output + expected_cr + expected_cw)) < 1e-9


def test_compute_cost_cache_tokens_clamped_to_zero():
    # cache tokens > input_tokens should not produce negative regular input cost
    model = _model(input_price=1.0, output_price=1.0)
    usage = Usage(input_tokens=10, cache_read_tokens=100, cache_write_tokens=0)
    cost = compute_cost(usage, model)
    assert cost.input >= 0.0


# ---------------------------------------------------------------------------
# _price_tier — 128k tier
# ---------------------------------------------------------------------------


def test_compute_cost_128k_tier_uses_tiered_prices():
    model = _model(
        input_price=1.0,
        output_price=2.0,
        input_128k=5.0,
        output_128k=10.0,
    )
    usage = Usage(input_tokens=150_000, output_tokens=10_000)
    cost = compute_cost(usage, model)
    expected_input = (150_000 / 1_000_000) * 5.0
    expected_output = (10_000 / 1_000_000) * 10.0
    assert abs(cost.input - expected_input) < 1e-9
    assert abs(cost.output - expected_output) < 1e-9


def test_compute_cost_128k_tier_falls_back_to_base_when_no_128k():
    # >128k tokens but no 128k pricing → uses base rates
    model = _model(input_price=1.0, output_price=2.0)
    usage = Usage(input_tokens=150_000, output_tokens=10_000)
    cost = compute_cost(usage, model)
    expected_input = (150_000 / 1_000_000) * 1.0
    assert abs(cost.input - expected_input) < 1e-9


def test_compute_cost_128k_tier_output_falls_back_to_base():
    # 128k tier defined for input but no output_128k → falls back to base output price
    model = _model(input_price=1.0, output_price=2.0, input_128k=5.0, output_128k=None)
    usage = Usage(input_tokens=150_000, output_tokens=10_000)
    cost = compute_cost(usage, model)
    expected_output = (10_000 / 1_000_000) * 2.0
    assert abs(cost.output - expected_output) < 1e-9


# ---------------------------------------------------------------------------
# _price_tier — 200k tier
# ---------------------------------------------------------------------------


def test_compute_cost_200k_tier_uses_tiered_prices():
    model = _model(
        input_price=1.0,
        output_price=2.0,
        cache_read=0.1,
        cache_write=0.2,
        input_200k=6.0,
        output_200k=22.5,
        cache_read_200k=0.6,
        cache_write_200k=7.5,
    )
    usage = Usage(
        input_tokens=250_000,
        output_tokens=5_000,
        cache_read_tokens=10_000,
        cache_write_tokens=0,
    )
    cost = compute_cost(usage, model)
    regular = 250_000 - 10_000
    expected_input = (regular / 1_000_000) * 6.0
    expected_output = (5_000 / 1_000_000) * 22.5
    expected_cr = (10_000 / 1_000_000) * 0.6
    assert abs(cost.input - expected_input) < 1e-9
    assert abs(cost.output - expected_output) < 1e-9
    assert abs(cost.cache_read - expected_cr) < 1e-9


def test_compute_cost_200k_tier_falls_back_to_base_when_no_200k():
    model = _model(input_price=3.0, output_price=15.0)
    usage = Usage(input_tokens=250_000, output_tokens=5_000)
    cost = compute_cost(usage, model)
    expected_input = (250_000 / 1_000_000) * 3.0
    assert abs(cost.input - expected_input) < 1e-9


def test_compute_cost_200k_tier_output_falls_back_to_base():
    # 200k defined for input but not output → falls back through or-chain to base output
    model = _model(
        input_price=1.0,
        output_price=2.0,
        input_200k=6.0,
        output_200k=None,
    )
    usage = Usage(input_tokens=250_000, output_tokens=5_000)
    cost = compute_cost(usage, model)
    expected_output = (5_000 / 1_000_000) * 2.0
    assert abs(cost.output - expected_output) < 1e-9


def test_compute_cost_200k_tier_cache_falls_back_to_base():
    # 200k tier but no 200k cache prices → fall back to base cache prices
    model = _model(
        input_price=1.0,
        output_price=2.0,
        cache_read=0.1,
        cache_write=0.2,
        input_200k=6.0,
        output_200k=22.5,
        cache_read_200k=None,
        cache_write_200k=None,
    )
    usage = Usage(input_tokens=250_000, output_tokens=0, cache_read_tokens=50_000)
    cost = compute_cost(usage, model)
    expected_cr = (50_000 / 1_000_000) * 0.1
    assert abs(cost.cache_read - expected_cr) < 1e-9


# ---------------------------------------------------------------------------
# _price_tier — 272k tier
# ---------------------------------------------------------------------------


def test_compute_cost_272k_tier_uses_tiered_prices():
    model = Model(
        name="test-model",
        provider="test",
        display_name="Test Model",
        input_price_per_mtok=1.0,
        output_price_per_mtok=2.0,
        cache_read_price_per_mtok=0.1,
        cache_write_price_per_mtok=0.2,
        input_price_per_mtok_272k=8.0,
        output_price_per_mtok_272k=24.0,
        cache_read_price_per_mtok_272k=0.8,
    )
    usage = Usage(
        input_tokens=300_000,
        output_tokens=5_000,
        cache_read_tokens=10_000,
        cache_write_tokens=0,
    )
    cost = compute_cost(usage, model)
    regular = 300_000 - 10_000
    expected_input = (regular / 1_000_000) * 8.0
    expected_output = (5_000 / 1_000_000) * 24.0
    expected_cr = (10_000 / 1_000_000) * 0.8
    assert abs(cost.input - expected_input) < 1e-9
    assert abs(cost.output - expected_output) < 1e-9
    assert abs(cost.cache_read - expected_cr) < 1e-9


def test_compute_cost_272k_tier_falls_back_to_base_when_no_272k():
    # >272k tokens but no 272k pricing → uses base rates
    model = _model(input_price=1.0, output_price=2.0)
    usage = Usage(input_tokens=300_000, output_tokens=5_000)
    cost = compute_cost(usage, model)
    expected_input = (300_000 / 1_000_000) * 1.0
    assert abs(cost.input - expected_input) < 1e-9


def test_compute_cost_272k_tier_output_falls_back_to_base():
    # 272k input defined but no output_272k → falls back to base output price
    model = Model(
        name="test-model",
        provider="test",
        display_name="Test Model",
        input_price_per_mtok=1.0,
        output_price_per_mtok=2.0,
        input_price_per_mtok_272k=8.0,
        output_price_per_mtok_272k=None,
    )
    usage = Usage(input_tokens=300_000, output_tokens=5_000)
    cost = compute_cost(usage, model)
    expected_output = (5_000 / 1_000_000) * 2.0
    assert abs(cost.output - expected_output) < 1e-9


def test_compute_cost_272k_tier_cache_read_falls_back_to_base():
    # 272k input defined but no cache_read_272k → falls back to base cache_read price
    model = Model(
        name="test-model",
        provider="test",
        display_name="Test Model",
        input_price_per_mtok=1.0,
        output_price_per_mtok=2.0,
        cache_read_price_per_mtok=0.1,
        input_price_per_mtok_272k=8.0,
        cache_read_price_per_mtok_272k=None,
    )
    usage = Usage(input_tokens=300_000, output_tokens=0, cache_read_tokens=50_000)
    cost = compute_cost(usage, model)
    expected_cr = (50_000 / 1_000_000) * 0.1
    assert abs(cost.cache_read - expected_cr) < 1e-9


def test_compute_cost_272k_tier_uses_base_cache_write():
    # No 272k cache_write rate exists; 272k tier always uses base cache_write price
    model = Model(
        name="test-model",
        provider="test",
        display_name="Test Model",
        input_price_per_mtok=1.0,
        output_price_per_mtok=2.0,
        cache_write_price_per_mtok=0.2,
        input_price_per_mtok_272k=8.0,
    )
    usage = Usage(input_tokens=300_000, output_tokens=0, cache_write_tokens=20_000)
    cost = compute_cost(usage, model)
    expected_cw = (20_000 / 1_000_000) * 0.2
    assert abs(cost.cache_write - expected_cw) < 1e-9


# ---------------------------------------------------------------------------
# models registry
# ---------------------------------------------------------------------------


def test_list_models_returns_nonempty():
    models = list_models()
    assert len(models) > 0
    assert all(isinstance(m, Model) for m in models)


def test_list_models_sorted_by_provider_then_date():
    models = list_models()
    providers = [m.provider for m in models]
    assert providers == sorted(providers)


def test_get_model_known():
    models = list_models()
    first = models[0]
    result = get_model(first.name)
    assert result is not None
    assert result.name == first.name


def test_get_model_unknown_returns_none():
    result = get_model("this-model-does-not-exist")
    assert result is None
