from agentcost.pricing import PricingDB


class TestPricingDB:
    def test_loads_default_llm_prices(self):
        db = PricingDB()
        assert db.llm_cost("gpt-4o", 1_000_000, 0) > 0

    def test_loads_default_tool_prices(self):
        db = PricingDB()
        assert db.tool_cost("stripe_api") == 0.02

    def test_unknown_tool_returns_zero(self):
        db = PricingDB()
        assert db.tool_cost("nonexistent_tool") == 0.0

    def test_unknown_llm_returns_zero(self):
        db = PricingDB()
        assert db.llm_cost("nonexistent-model", 1000, 1000) == 0.0

    def test_llm_cost_calculation(self):
        db = PricingDB()
        # gpt-4o: $2.50/M input, $10.00/M output
        cost = db.llm_cost("gpt-4o", 1_000_000, 1_000_000)
        assert cost == 12.50

    def test_llm_cost_small_tokens(self):
        db = PricingDB()
        # gpt-4o: 500 input, 200 output
        cost = db.llm_cost("gpt-4o", 500, 200)
        expected = (500 / 1_000_000) * 2.50 + (200 / 1_000_000) * 10.00
        assert abs(cost - expected) < 1e-8

    def test_llm_prefix_match(self):
        db = PricingDB()
        # "gpt-4o-2024-08-06" should match "gpt-4o"
        cost = db.llm_cost("gpt-4o-2024-08-06", 1_000_000, 0)
        assert cost == 2.50

    def test_pricing_overrides(self):
        db = PricingDB(overrides={
            "my_tool": 0.05,
            "my-model": {"input": 1.00, "output": 2.00},
        })
        assert db.tool_cost("my_tool") == 0.05
        assert db.llm_cost("my-model", 1_000_000, 1_000_000) == 3.00

    def test_set_tool_price(self):
        db = PricingDB()
        db.set_tool_price("custom_api", 0.10)
        assert db.tool_cost("custom_api") == 0.10

    def test_set_llm_price(self):
        db = PricingDB()
        db.set_llm_price("custom-model", 5.00, 15.00)
        cost = db.llm_cost("custom-model", 1_000_000, 1_000_000)
        assert cost == 20.00

    def test_override_existing_tool_price(self):
        db = PricingDB()
        original = db.tool_cost("stripe_api")
        db.set_tool_price("stripe_api", 0.99)
        assert db.tool_cost("stripe_api") == 0.99
        assert db.tool_cost("stripe_api") != original
