"""Tests for BudgetConfig and BudgetTracker."""

import time

import pytest

from .budget import BudgetConfig, BudgetTracker, MODEL_PRICING, _DEFAULT_PRICING


class TestBudgetConfig:
    def test_defaults_are_unlimited(self) -> None:
        cfg = BudgetConfig()
        assert cfg.max_tokens == 0
        assert cfg.max_cost_usd == 0.0
        assert cfg.max_tool_calls == 0
        assert cfg.max_duration_s == 0.0

    def test_custom_limits(self) -> None:
        cfg = BudgetConfig(max_tokens=1000, max_cost_usd=0.50, max_tool_calls=5, max_duration_s=30.0)
        assert cfg.max_tokens == 1000
        assert cfg.max_cost_usd == 0.50
        assert cfg.max_tool_calls == 5
        assert cfg.max_duration_s == 30.0


class TestBudgetTracker:
    def test_cost_calculation_gpt4o_mini(self) -> None:
        """gpt-4o-mini: 1000 input * 0.15/1M + 500 output * 0.60/1M = 0.00045."""
        tracker = BudgetTracker(model="gpt-4o-mini")

        class Usage:
            input_tokens = 1000
            output_tokens = 500
            total_tokens = 1500

        tracker.record_usage(Usage())
        summary = tracker.summary
        assert summary["input_tokens"] == 1000
        assert summary["output_tokens"] == 500
        assert summary["total_tokens"] == 1500
        expected_cost = 1000 * 0.15 / 1_000_000 + 500 * 0.60 / 1_000_000
        assert abs(summary["cost_usd"] - expected_cost) < 1e-10

    def test_cost_calculation_gpt4o(self) -> None:
        """gpt-4o: 1000 input * 2.50/1M + 500 output * 10.00/1M = 0.0075."""
        tracker = BudgetTracker(model="gpt-4o")

        class Usage:
            input_tokens = 1000
            output_tokens = 500
            total_tokens = 1500

        tracker.record_usage(Usage())
        expected_cost = 1000 * 2.50 / 1_000_000 + 500 * 10.00 / 1_000_000
        assert abs(tracker.summary["cost_usd"] - expected_cost) < 1e-10

    def test_cost_unknown_model_uses_default(self) -> None:
        """Unknown model should use default pricing (gpt-4o rates)."""
        tracker = BudgetTracker(model="some-future-model")

        class Usage:
            input_tokens = 1000
            output_tokens = 500
            total_tokens = 1500

        tracker.record_usage(Usage())
        input_price, output_price = _DEFAULT_PRICING
        expected_cost = 1000 * input_price / 1_000_000 + 500 * output_price / 1_000_000
        assert abs(tracker.summary["cost_usd"] - expected_cost) < 1e-10

    def test_check_budget_tokens_exceeded(self) -> None:
        tracker = BudgetTracker(model="gpt-4o-mini", config=BudgetConfig(max_tokens=100))

        class Usage:
            input_tokens = 80
            output_tokens = 50
            total_tokens = 130

        tracker.record_usage(Usage())
        result = tracker.check_budget()
        assert result is not None
        assert "token limit" in result
        assert "100" in result
        assert "130" in result

    def test_check_budget_cost_exceeded(self) -> None:
        tracker = BudgetTracker(model="gpt-4o", config=BudgetConfig(max_cost_usd=0.001))

        class Usage:
            input_tokens = 10000
            output_tokens = 5000
            total_tokens = 15000

        tracker.record_usage(Usage())
        result = tracker.check_budget()
        assert result is not None
        assert "cost limit" in result

    def test_check_budget_tool_calls_exceeded(self) -> None:
        tracker = BudgetTracker(model="gpt-4o-mini", config=BudgetConfig(max_tool_calls=2))
        tracker.record_tool_call()
        tracker.record_tool_call()
        tracker.record_tool_call()
        result = tracker.check_budget()
        assert result is not None
        assert "tool call limit" in result
        assert "2" in result
        assert "3" in result

    def test_check_budget_within_limits(self) -> None:
        tracker = BudgetTracker(
            model="gpt-4o-mini",
            config=BudgetConfig(max_tokens=1000, max_cost_usd=1.0, max_tool_calls=10),
        )

        class Usage:
            input_tokens = 100
            output_tokens = 50
            total_tokens = 150

        tracker.record_usage(Usage())
        tracker.record_tool_call()
        assert tracker.check_budget() is None

    def test_no_config_always_within_budget(self) -> None:
        """No config (all defaults 0 = unlimited) should never exceed budget."""
        tracker = BudgetTracker(model="gpt-4o")

        class Usage:
            input_tokens = 1_000_000
            output_tokens = 1_000_000
            total_tokens = 2_000_000

        tracker.record_usage(Usage())
        for _ in range(100):
            tracker.record_tool_call()
        assert tracker.check_budget() is None

    def test_record_tool_call(self) -> None:
        tracker = BudgetTracker(model="gpt-4o-mini")
        assert tracker.summary["tool_calls"] == 0
        tracker.record_tool_call()
        assert tracker.summary["tool_calls"] == 1
        tracker.record_tool_call()
        tracker.record_tool_call()
        assert tracker.summary["tool_calls"] == 3

    def test_none_usage_ignored(self) -> None:
        tracker = BudgetTracker(model="gpt-4o-mini")
        tracker.record_usage(None)
        summary = tracker.summary
        assert summary["input_tokens"] == 0
        assert summary["output_tokens"] == 0
        assert summary["total_tokens"] == 0
        assert summary["cost_usd"] == 0.0

    def test_summary_contains_duration(self) -> None:
        tracker = BudgetTracker(model="gpt-4o-mini")
        summary = tracker.summary
        assert "duration_s" in summary
        assert summary["duration_s"] >= 0.0

    def test_model_pricing_has_expected_models(self) -> None:
        """Verify pricing table contains expected models."""
        expected_models = ["gpt-4o", "gpt-4o-mini", "gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "o3", "o3-mini", "o4-mini"]
        for model in expected_models:
            assert model in MODEL_PRICING, f"Missing pricing for {model}"
            inp, out = MODEL_PRICING[model]
            assert inp > 0
            assert out > 0
