"""Tests for spens.pricing (Portkey dataset parsing, matching and caching).

The dataset excerpts below are copied verbatim from
``https://configs.portkey.ai/pricing/google.json``.  Google is used for the
parsing cases because it exercises every awkward part of the format at once:
context-banded keys, per-model cache rates, and entries that carry no token
prices at all.
"""

from __future__ import annotations

import json
import os
import time
import urllib.error
from pathlib import Path

import pytest
from spens import pricing
from spens.pricing import (
    Price,
    PricingTable,
    Tier,
    best_match,
    build_table,
    load_provider_dataset,
    load_table,
    model_candidates,
    parse_dataset,
    refresh_datasets,
    unit_warning,
)

GOOGLE_EXCERPT: dict = {
    "default": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0},
                "response_token": {"price": 0},
            },
            "calculate": {
                "request": {
                    "operation": "multiply",
                    "operands": [{"value": "input_tokens"}, {"value": "rates.request_token"}],
                },
                "response": {
                    "operation": "multiply",
                    "operands": [{"value": "output_tokens"}, {"value": "rates.response_token"}],
                },
            },
            "currency": "USD",
        }
    },
    "gemini-2.5-flash-lte-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.00003},
                "response_token": {"price": 0.00025},
                "additional_units": {"web_search": {"price": 3.5}, "search": {"price": 3.5}},
                "cache_read_input_token": {"price": 0.000003},
            },
            "batch_config": {
                "request_token": {"price": 0.000015},
                "response_token": {"price": 0.000125},
            },
        }
    },
    "gemini-2.5-flash-gt-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.00003},
                "response_token": {"price": 0.00025},
                "additional_units": {"web_search": {"price": 3.5}, "search": {"price": 3.5}},
                "cache_read_input_token": {"price": 0.000003},
            },
            "batch_config": {
                "request_token": {"price": 0.000015},
                "response_token": {"price": 0.000125},
            },
        }
    },
    "gemini-2.5-pro-lte-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.000125},
                "response_token": {"price": 0.001},
                "additional_units": {"web_search": {"price": 3.5}, "search": {"price": 3.5}},
                "cache_read_input_token": {"price": 0.0000125},
            },
            "batch_config": {
                "request_token": {"price": 0.0000625},
                "response_token": {"price": 0.0005},
            },
        }
    },
    "gemini-2.5-pro-gt-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.00025},
                "response_token": {"price": 0.0015},
                "additional_units": {"web_search": {"price": 3.5}, "search": {"price": 3.5}},
                "cache_read_input_token": {"price": 0.000025},
            },
            "batch_config": {
                "request_token": {"price": 0.000125},
                "response_token": {"price": 0.00075},
            },
        }
    },
    "gemini-2.5-flash-image-lte-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.00003},
                "response_token": {"price": 0.00025},
                "additional_units": {
                    "web_search": {"price": 3.5},
                    "image_token": {"price": 0.003},
                    "search": {"price": 3.5},
                },
            },
            "batch_config": {
                "request_token": {"price": 0.000015},
                "response_token": {"price": 0.000125},
            },
        }
    },
    "gemini-3.1-pro-preview-lte-200k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.0002},
                "response_token": {"price": 0.0012},
                "cache_write_input_token": {"price": 0.00002},
                "cache_read_input_token": {"price": 0.00045},
                "additional_units": {"web_search": {"price": 3.5}, "search": {"price": 3.5}},
            },
            "batch_config": {
                "request_token": {"price": 0.0001},
                "response_token": {"price": 0.0006},
            },
        }
    },
    "gemini-3.1-pro-preview-gt-200k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.0004},
                "response_token": {"price": 0.0018},
                "cache_write_input_token": {"price": 0.00004},
                "cache_read_input_token": {"price": 0.00045},
                "additional_units": {"web_search": {"price": 3.5}, "search": {"price": 3.5}},
            },
            "batch_config": {
                "request_token": {"price": 0.0002},
                "response_token": {"price": 0.0009},
            },
        }
    },
    "gemini-3.1-pro-preview-lte-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.0002},
                "response_token": {"price": 0.0012},
                "cache_read_input_token": {"price": 0.00002},
                "additional_units": {"web_search": {"price": 1.4}, "search": {"price": 1.4}},
            },
            "batch_config": {
                "request_token": {"price": 0.0001},
                "response_token": {"price": 0.0006},
            },
        }
    },
    "gemini-3.1-pro-preview-gt-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0.0004},
                "response_token": {"price": 0.0018},
                "cache_read_input_token": {"price": 0.00004},
                "additional_units": {"web_search": {"price": 1.4}, "search": {"price": 1.4}},
            },
            "batch_config": {
                "request_token": {"price": 0.0002},
                "response_token": {"price": 0.0009},
            },
        }
    },
    "imagen-4.0-generate-001-lte-128k": {
        "pricing_config": {
            "pay_as_you_go": {
                "request_token": {"price": 0},
                "response_token": {"price": 0},
                "image": {"default": {"default": {"price": 4}}},
            }
        }
    },
}


def _google_table() -> PricingTable:
    return build_table([("google", parse_dataset(GOOGLE_EXCERPT))])


# -- unit conversion ---------------------------------------------------------


def test_prices_are_converted_to_usd_per_1m_tokens() -> None:
    """Pins PRICE_TO_USD_PER_1M against published Google prices.

    Gemini 2.5 Flash is $0.30 / $2.50 per 1M and 2.5 Pro (<=128k context) is
    $1.25 / $10.00.  Reading the dataset's figure as a per-token or a
    per-1M-token price puts the whole session's cost out by 100x either way.
    """
    table = _google_table()

    flash = table.lookup("gemini-2.5-flash")
    assert flash is not None
    assert flash.input_per_1m == pytest.approx(0.30)
    assert flash.output_per_1m == pytest.approx(2.50)

    pro = table.lookup("gemini-2.5-pro")
    assert pro is not None
    assert pro.input_per_1m == pytest.approx(1.25)
    assert pro.output_per_1m == pytest.approx(10.00)


def test_cache_rates_are_read_from_the_dataset() -> None:
    """Gemini 2.5 Flash cache reads are $0.03/1M, not a tenth of input."""
    flash = _google_table().lookup("gemini-2.5-flash")
    assert flash is not None
    assert flash.cache_read_per_1m == pytest.approx(0.03)
    # Absent from this entry, so the caller estimates it.
    assert flash.cache_write_per_1m is None


def test_cache_write_rate_is_read_when_present() -> None:
    price = _google_table().lookup("gemini-3.1-pro-preview", 150_000)
    assert price is not None
    assert price.cache_write_per_1m == pytest.approx(0.20)
    assert price.cache_read_per_1m == pytest.approx(4.50)


# -- entries that must not produce a price -----------------------------------


def test_default_entry_is_skipped() -> None:
    """The dataset's ``default`` entry prices everything at 0.

    Treating it as a match would make every unlisted model resolve
    "successfully" at $0 and shadow the fallback table.
    """
    assert "default" not in parse_dataset(GOOGLE_EXCERPT)


def test_entries_without_token_prices_are_skipped() -> None:
    """Imagen is priced per image, so it has no usable token price."""
    assert "imagen-4.0-generate-001-lte-128k" not in parse_dataset(GOOGLE_EXCERPT)


def test_malformed_entries_do_not_break_the_file() -> None:
    prices = parse_dataset(
        {
            "broken-a": None,
            "broken-b": {"pricing_config": []},
            "broken-c": {"pricing_config": {"pay_as_you_go": {"request_token": {"price": "free"}}}},
            "good": {"pricing_config": {"pay_as_you_go": {"request_token": {"price": 0.0003}}}},
        }
    )
    assert list(prices) == ["good"]
    assert prices["good"].input_per_1m == pytest.approx(3.00)


def test_parse_dataset_tolerates_a_non_dict_payload() -> None:
    assert parse_dataset(["nope"]) == {}


# -- context bands -----------------------------------------------------------


def test_band_is_picked_from_the_context_size() -> None:
    table = _google_table()

    small = table.lookup("gemini-2.5-pro", 10_000)
    large = table.lookup("gemini-2.5-pro", 500_000)
    assert small is not None and large is not None
    assert small.input_per_1m == pytest.approx(1.25)
    assert large.input_per_1m == pytest.approx(2.50)
    assert large.output_per_1m == pytest.approx(15.00)


def test_mixed_band_limits_resolve_to_the_tightest_bound() -> None:
    """One model can carry both 128k and 200k bands upstream.

    Gemini 3.1 Pro lists ``-lte-128k``/``-gt-128k`` *and*
    ``-lte-200k``/``-gt-200k`` with different cache rates, so the choice has
    to be deterministic rather than dict-order dependent.
    """
    table = _google_table()

    assert table.lookup("gemini-3.1-pro-preview", 0).cache_read_per_1m == pytest.approx(0.20)
    assert table.lookup("gemini-3.1-pro-preview", 150_000).cache_read_per_1m == pytest.approx(4.50)

    huge = table.lookup("gemini-3.1-pro-preview", 250_000)
    assert huge.input_per_1m == pytest.approx(4.00)
    assert huge.cache_read_per_1m == pytest.approx(4.50)


def test_only_gt_bands_still_price_a_small_context() -> None:
    table = build_table([("x", {"m-gt-128k": Price(1.0, 2.0)})])
    price = table.lookup("m", 10)
    assert price is not None
    assert price.input_per_1m == pytest.approx(1.0)


def test_pick_tier_ignores_a_duplicate_band() -> None:
    """A repeated (limit, side) pair must not multiply into the band list."""
    table = build_table(
        [
            ("a", {"m-lte-128k": Price(1.0, 2.0)}),
            ("b", {"m-lte-128k": Price(9.0, 9.0)}),
        ]
    )
    assert table.tiers["m"] == (Tier(limit=128_000, side="lte", price=Price(1.0, 2.0)),)


# -- model id matching -------------------------------------------------------


def test_a_longer_key_never_prices_a_shorter_model_id() -> None:
    """``gemini-2.5-flash`` must not be costed as ``gemini-2.5-flash-image``.

    Matching only ever goes dataset-key-into-model-id; the reverse direction
    would silently bill a cheap model at an image model's rate.
    """
    image_only = build_table(
        [("google", parse_dataset({"gemini-2.5-flash-image-lte-128k": GOOGLE_EXCERPT["gemini-2.5-flash-image-lte-128k"]}))]
    )
    assert image_only.lookup("gemini-2.5-flash") is None
    assert image_only.lookup("gemini-2.5-flash-image") is not None


def test_dated_model_ids_match_their_base_entry() -> None:
    table = build_table([("anthropic", {"claude-sonnet-4": Price(3.0, 15.0)})])
    price = table.lookup("claude-sonnet-4-20250514")
    assert price is not None
    assert price.input_per_1m == pytest.approx(3.0)


def test_vendor_prefixed_ids_resolve() -> None:
    table = build_table([("anthropic", {"claude-sonnet-4": Price(3.0, 15.0)})])
    assert table.lookup("anthropic/claude-sonnet-4") is not None


def test_bedrock_region_and_version_tagged_ids_resolve() -> None:
    table = build_table([("anthropic", {"claude-sonnet-4-20250514": Price(3.0, 15.0)})])
    assert table.lookup("us.anthropic.claude-sonnet-4-20250514-v1:0") is not None


def test_model_candidates_keeps_a_version_number_intact() -> None:
    """The ``2.`` in ``gemini-2.5-flash`` is not a vendor prefix."""
    assert model_candidates("gemini-2.5-flash") == ["gemini-2.5-flash"]


def test_model_candidates_normalises_case_and_separators() -> None:
    assert model_candidates("GPT_4o Mini")[0] == "gpt-4o-mini"


def test_model_candidates_of_nothing_is_empty() -> None:
    assert model_candidates(None) == []
    assert model_candidates("") == []


def test_best_match_prefers_exact_then_longest_prefix() -> None:
    keys = {"gpt-4o": 1, "gpt-4o-mini": 2, "gpt-4": 3}
    assert best_match(keys, "gpt-4o-mini") == "gpt-4o-mini"
    assert best_match(keys, "gpt-4o-2024-08-06") == "gpt-4o"
    assert best_match(keys, "unrelated-model") is None


def test_lookup_of_an_unknown_model_is_none() -> None:
    assert _google_table().lookup("some-future-model") is None
    assert _google_table().lookup(None) is None


# -- provider precedence -----------------------------------------------------


def test_the_first_provider_listing_a_model_prices_it() -> None:
    """A direct provider's rate wins over a reseller's for the same id."""
    table = build_table(
        [
            ("anthropic", {"claude-sonnet-4": Price(3.0, 15.0)}),
            ("openrouter", {"claude-sonnet-4": Price(99.0, 99.0)}),
        ]
    )
    assert table.lookup("claude-sonnet-4").input_per_1m == pytest.approx(3.0)


# -- fetching, caching and fallbacks -----------------------------------------


@pytest.fixture
def counting_fetch(monkeypatch):
    """Serve a canned dataset over the fetch path, counting requests."""
    calls: list[str] = []

    def _fake(url: str, timeout: float):
        calls.append(url)
        return {"claude-sonnet-4": {"pricing_config": {"pay_as_you_go": {"request_token": {"price": 0.0003}}}}}

    monkeypatch.setattr(pricing, "_http_get_json", _fake)
    pricing.set_live_fetch(True)
    return calls


def test_a_fetched_dataset_is_cached_and_reused(counting_fetch) -> None:
    first, first_source = load_provider_dataset("anthropic")
    second, second_source = load_provider_dataset("anthropic")

    assert first == second
    assert (first_source, second_source) == (pricing.SOURCE_LIVE, pricing.SOURCE_CACHE)
    assert len(counting_fetch) == 1


def test_a_stale_cache_is_refetched(counting_fetch) -> None:
    load_provider_dataset("anthropic")
    cached = pricing.cache_dir() / "anthropic.json"
    stale = time.time() - (pricing.ttl_seconds() + 60)
    os.utime(cached, (stale, stale))

    _, source = load_provider_dataset("anthropic")
    assert source == pricing.SOURCE_LIVE
    assert len(counting_fetch) == 2


def test_a_failed_fetch_falls_back_to_the_vendored_snapshot(monkeypatch) -> None:
    def _boom(url: str, timeout: float):
        raise urllib.error.URLError("no route to host")

    monkeypatch.setattr(pricing, "_http_get_json", _boom)
    pricing.set_live_fetch(True)
    vendored = pricing.vendored_path("anthropic")
    vendored.parent.mkdir(parents=True, exist_ok=True)
    vendored.write_text(json.dumps({"claude-sonnet-4": {}}), encoding="utf-8")

    data, source = load_provider_dataset("anthropic")
    assert source == pricing.SOURCE_VENDORED
    assert data == {"claude-sonnet-4": {}}


def test_offline_uses_the_vendored_snapshot_without_fetching(counting_fetch) -> None:
    pricing.set_live_fetch(False)
    vendored = pricing.vendored_path("anthropic")
    vendored.parent.mkdir(parents=True, exist_ok=True)
    vendored.write_text(json.dumps({"claude-sonnet-4": {}}), encoding="utf-8")

    _, source = load_provider_dataset("anthropic")
    assert source == pricing.SOURCE_VENDORED
    assert counting_fetch == []


def test_missing_everywhere_reports_unavailable() -> None:
    data, source = load_provider_dataset("anthropic")
    assert data is None
    assert source == pricing.SOURCE_UNAVAILABLE


def test_a_corrupt_cache_file_is_ignored(counting_fetch) -> None:
    cached = pricing.cache_dir() / "anthropic.json"
    cached.parent.mkdir(parents=True, exist_ok=True)
    cached.write_text("{not json", encoding="utf-8")

    _, source = load_provider_dataset("anthropic")
    assert source == pricing.SOURCE_LIVE


def test_load_table_records_where_each_provider_came_from(counting_fetch) -> None:
    table = load_table(("anthropic", "openai"))
    assert table.sources == {"anthropic": pricing.SOURCE_LIVE, "openai": pricing.SOURCE_LIVE}
    assert table.lookup("claude-sonnet-4") is not None


def test_load_table_is_memoised_for_the_default_providers(counting_fetch) -> None:
    """The viewer summarises sessions in a loop; the table is built once."""
    first = load_table()
    assert load_table() is first

    pricing.reset_table_cache()
    assert load_table() is not first


def test_an_empty_table_is_falsey() -> None:
    assert not PricingTable()
    assert _google_table()


# -- refresh and unit verification -------------------------------------------


def test_refresh_writes_the_cache_and_the_vendored_snapshot(monkeypatch, tmp_path: Path) -> None:
    monkeypatch.setattr(pricing, "VENDOR_DIR", tmp_path / "vendor")
    monkeypatch.setattr(
        pricing,
        "_http_get_json",
        lambda url, timeout: {"claude-sonnet-4": {"pricing_config": {"pay_as_you_go": {
            "request_token": {"price": 0.0003},
            "response_token": {"price": 0.0015},
        }}}},
    )

    statuses = refresh_datasets(("anthropic",), vendor=True)

    assert len(statuses) == 1
    status = statuses[0]
    assert status.source == pricing.SOURCE_LIVE
    assert status.models == 1
    assert not status.error and not status.warning
    assert status.vendored == tmp_path / "vendor" / "anthropic.json"
    assert status.vendored.exists()
    assert (pricing.cache_dir() / "anthropic.json").exists()


def test_refresh_reports_a_failed_provider(monkeypatch) -> None:
    def _boom(url: str, timeout: float):
        raise urllib.error.URLError("nope")

    monkeypatch.setattr(pricing, "_http_get_json", _boom)
    statuses = refresh_datasets(("bedrock",))
    assert statuses[0].source == pricing.SOURCE_UNAVAILABLE
    assert "nope" in statuses[0].error


def test_unit_warning_passes_on_correctly_scaled_prices() -> None:
    assert unit_warning("google", parse_dataset(GOOGLE_EXCERPT)) == ""
    assert unit_warning("anthropic", {"claude-sonnet-4": Price(3.00, 15.00)}) == ""


def test_unit_warning_catches_a_changed_price_unit() -> None:
    """If Portkey ever switches to per-token figures, say so loudly."""
    warning = unit_warning("anthropic", {"claude-sonnet-4": Price(300.0, 1500.0)})
    assert "PRICE_TO_USD_PER_1M" in warning


def test_unit_warning_skips_a_provider_without_a_reference() -> None:
    assert unit_warning("fireworks", {"whatever": Price(1.0, 2.0)}) == ""


def test_unit_warning_notes_a_missing_reference_model() -> None:
    warning = unit_warning("openai", {"gpt-4o-mini": Price(0.15, 0.60)})
    assert "unit check skipped" in warning
