from multi_swarm.genome.hypothesis import ModelTier
from multi_swarm.llm.cost_tracker import CostTracker, estimate_cost


def test_estimate_cost_tier_c():
    cost = estimate_cost(input_tokens=1_000_000, output_tokens=1_000_000, tier=ModelTier.C)
    assert cost == 0.40 + 0.40


def test_estimate_cost_tier_b():
    cost = estimate_cost(input_tokens=1_000_000, output_tokens=1_000_000, tier=ModelTier.B)
    assert cost == 0.14 + 0.28


def test_tracker_accumulates():
    t = CostTracker()
    t.record(input_tokens=10_000, output_tokens=20_000, tier=ModelTier.C, run_id="r", agent_id="a")
    t.record(input_tokens=5_000, output_tokens=15_000, tier=ModelTier.C, run_id="r", agent_id="b")
    summary = t.summary()
    assert summary["calls"] == 2
    assert summary["input_tokens"] == 15_000
    assert summary["output_tokens"] == 35_000
    assert summary["cost_usd"] > 0


def test_tracker_per_tier_breakdown():
    t = CostTracker()
    t.record(input_tokens=10_000, output_tokens=10_000, tier=ModelTier.C, run_id="r", agent_id="a")
    t.record(input_tokens=10_000, output_tokens=10_000, tier=ModelTier.B, run_id="r", agent_id="b")
    summary = t.summary()
    assert "C" in summary["by_tier"]
    assert "B" in summary["by_tier"]


def test_estimate_cost_tier_s():
    cost = estimate_cost(input_tokens=1_000_000, output_tokens=1_000_000, tier=ModelTier.S)
    assert cost == 0.50 + 3.00


def test_estimate_cost_tier_a():
    cost = estimate_cost(input_tokens=1_000_000, output_tokens=1_000_000, tier=ModelTier.A)
    assert cost == 0.14 + 0.28


def test_estimate_cost_tier_d():
    cost = estimate_cost(input_tokens=1_000_000, output_tokens=1_000_000, tier=ModelTier.D)
    assert cost == 0.03 + 0.14


def test_tracker_summary_contains_all_five_tiers():
    t = CostTracker()
    for tier in (ModelTier.S, ModelTier.A, ModelTier.B, ModelTier.C, ModelTier.D):
        t.record(
            input_tokens=1_000,
            output_tokens=1_000,
            tier=tier,
            run_id="r",
            agent_id=f"a-{tier.value}",
        )
    summary = t.summary()
    for tier_letter in ("S", "A", "B", "C", "D"):
        assert tier_letter in summary["by_tier"]
        assert summary["by_tier"][tier_letter]["calls"] == 1


def test_tracker_default_call_kind_is_hypothesis():
    t = CostTracker()
    rec = t.record(input_tokens=10, output_tokens=10, tier=ModelTier.C, run_id="r", agent_id="a")
    assert rec.call_kind == "hypothesis"
    summary = t.summary()
    assert "hypothesis" in summary["by_call_kind"]
    assert summary["by_call_kind"]["hypothesis"]["calls"] == 1
    assert "mutation" not in summary["by_call_kind"]


def test_tracker_by_call_kind_breakdown():
    t = CostTracker()
    t.record(input_tokens=100, output_tokens=200, tier=ModelTier.C, run_id="r", agent_id="a")
    t.record(input_tokens=100, output_tokens=200, tier=ModelTier.C, run_id="r", agent_id="a")
    t.record(
        input_tokens=50, output_tokens=80, tier=ModelTier.B,
        run_id="r", agent_id="parent-x", call_kind="mutation",
    )
    summary = t.summary()
    assert summary["by_call_kind"]["hypothesis"]["calls"] == 2
    assert summary["by_call_kind"]["mutation"]["calls"] == 1
    assert summary["by_call_kind"]["mutation"]["input_tokens"] == 50
    assert summary["by_call_kind"]["mutation"]["output_tokens"] == 80