"""Tests for BudgetController."""

import time

from dpe.engine.budget import BudgetController
from dpe.models.config import BudgetLimits, ThrottleBehavior


class TestPreCallCheck:
    def test_allows_within_budget(self) -> None:
        bc = BudgetController(limits=BudgetLimits())
        decision = bc.pre_call_check(is_premium=False, stage="TEST", adapter="test")
        assert decision.allowed is True

    def test_blocks_total_calls_exceeded(self) -> None:
        limits = BudgetLimits(max_total_calls=2)
        bc = BudgetController(limits=limits)
        bc.total_calls_used = 2
        decision = bc.pre_call_check(is_premium=False, stage="TEST", adapter="test")
        assert decision.allowed is False
        assert "Total calls exceeded" in decision.reason

    def test_blocks_premium_calls_exceeded(self) -> None:
        limits = BudgetLimits(max_premium_calls=1)
        bc = BudgetController(limits=limits)
        bc.premium_calls_used = 1
        decision = bc.pre_call_check(is_premium=True, stage="TEST", adapter="test")
        assert decision.allowed is False
        assert "Premium calls exceeded" in decision.reason

    def test_allows_non_premium_when_premium_exhausted(self) -> None:
        limits = BudgetLimits(max_premium_calls=1)
        bc = BudgetController(limits=limits)
        bc.premium_calls_used = 1
        decision = bc.pre_call_check(is_premium=False, stage="TEST", adapter="local")
        assert decision.allowed is True

    def test_blocks_wall_time_exceeded(self) -> None:
        limits = BudgetLimits(max_wall_time_seconds=1)
        bc = BudgetController(limits=limits)
        bc.start_time = time.monotonic() - 2  # 2 seconds ago
        decision = bc.pre_call_check(is_premium=False, stage="TEST", adapter="test")
        assert decision.allowed is False
        assert "Wall time exceeded" in decision.reason


class TestRecordCall:
    def test_increments_total(self) -> None:
        bc = BudgetController(limits=BudgetLimits())
        bc.record_call(is_premium=False)
        assert bc.total_calls_used == 1
        assert bc.premium_calls_used == 0

    def test_increments_premium(self) -> None:
        bc = BudgetController(limits=BudgetLimits())
        bc.record_call(is_premium=True)
        assert bc.total_calls_used == 1
        assert bc.premium_calls_used == 1

    def test_monotonic(self) -> None:
        bc = BudgetController(limits=BudgetLimits())
        for _ in range(5):
            bc.record_call(is_premium=True)
        assert bc.total_calls_used == 5
        assert bc.premium_calls_used == 5


class TestCritiqueRounds:
    def test_can_critique_initially(self) -> None:
        bc = BudgetController(limits=BudgetLimits(max_critique_rounds=2))
        assert bc.can_critique() is True

    def test_cannot_critique_after_max(self) -> None:
        bc = BudgetController(limits=BudgetLimits(max_critique_rounds=2))
        bc.record_critique_round()
        bc.record_critique_round()
        assert bc.can_critique() is False


class TestThrottleBehavior:
    def test_abort_denies(self) -> None:
        limits = BudgetLimits(max_total_calls=1, throttle_behavior=ThrottleBehavior.ABORT)
        bc = BudgetController(limits=limits)
        bc.total_calls_used = 1
        decision = bc.pre_call_check(is_premium=False, stage="TEST", adapter="test")
        assert decision.allowed is False

    def test_degrade_to_local_allows(self) -> None:
        limits = BudgetLimits(
            max_premium_calls=1,
            throttle_behavior=ThrottleBehavior.DEGRADE_TO_LOCAL,
        )
        bc = BudgetController(limits=limits)
        bc.premium_calls_used = 1
        decision = bc.pre_call_check(is_premium=True, stage="TEST", adapter="premium")
        assert decision.allowed is True
        assert decision.degraded is True
        assert bc.is_degraded is True

    def test_retry_then_degrade(self) -> None:
        limits = BudgetLimits(
            max_premium_calls=1,
            throttle_behavior=ThrottleBehavior.RETRY_THEN_DEGRADE,
        )
        bc = BudgetController(limits=limits)
        bc.premium_calls_used = 1

        # First breach: retry (not degraded yet)
        d1 = bc.pre_call_check(is_premium=True, stage="TEST", adapter="premium")
        assert d1.allowed is True
        assert d1.degraded is False

        # Second breach: degrade
        d2 = bc.pre_call_check(is_premium=True, stage="TEST", adapter="premium")
        assert d2.allowed is True
        assert d2.degraded is True

    def test_throttle_events_recorded(self) -> None:
        limits = BudgetLimits(max_total_calls=1, throttle_behavior=ThrottleBehavior.ABORT)
        bc = BudgetController(limits=limits)
        bc.total_calls_used = 1
        bc.pre_call_check(is_premium=False, stage="TEST", adapter="test")
        assert len(bc.throttle_events) == 1
        assert bc.throttle_events[0].stage == "TEST"


class TestSnapshot:
    def test_snapshot_fields(self) -> None:
        bc = BudgetController(limits=BudgetLimits())
        bc.record_call(is_premium=True)
        snap = bc.snapshot()
        assert snap["total_calls_used"] == 1
        assert snap["premium_calls_used"] == 1
        assert "wall_time_seconds" in snap
