"""Tests for budget tracking and enforcement.

Phase 5, step 3: Token budget precedence and exhaustion handling.
"""

import pytest

from anima.core.budget import (
    resolve_budget,
    grace_system_message,
    format_budget_status,
)
from anima.core.tracker import BudgetTracker
from anima.core.schemas import JobSpec


# =============================================================================
# Test Resolve Budget
# =============================================================================


class TestResolveBudget:
    """Tests for budget source precedence (CLI > Job)."""

    def test_cli_overrides_all(self) -> None:
        """CLI max_tokens takes precedence over job."""
        job = JobSpec(
            persona_ref="test@v1",
            instance_id="inst_1",
            max_tokens=500,
        )
        # CLI=1000 overrides job=500
        result = resolve_budget(1000, job)
        assert result == 1000

    def test_job_overrides_persona(self) -> None:
        """Job max_tokens is used when CLI is not set."""
        job = JobSpec(
            persona_ref="test@v1",
            instance_id="inst_1",
            max_tokens=500,
        )
        # CLI=None, job=500 -> 500
        result = resolve_budget(None, job)
        assert result == 500

    def test_persona_fallback(self) -> None:
        """Returns None when CLI and job are not set."""
        job = JobSpec(
            persona_ref="test@v1",
            instance_id="inst_1",
            # max_tokens not set
        )
        # CLI=None, job.max_tokens=None -> None (unlimited)
        result = resolve_budget(None, job)
        assert result is None

    def test_unlimited_when_all_none(self) -> None:
        """Returns None (unlimited) when all sources are unset."""
        job = JobSpec(
            persona_ref="test@v1",
            instance_id="inst_1",
            # max_tokens not set
        )
        result = resolve_budget(None, job)
        assert result is None

    def test_cli_zero_invalid(self) -> None:
        """CLI max_tokens=0 is rejected by contract (must be >= 1 or None)."""
        job = JobSpec(persona_ref="test@v1", instance_id="inst_1")
        # Contract requires cli_max_tokens >= 1 if not None
        # This is enforced by @pre decorator
        with pytest.raises(Exception):  # PreContractError
            resolve_budget(0, job)

    def test_cli_negative_invalid(self) -> None:
        """CLI max_tokens=-1 is rejected by contract."""
        job = JobSpec(persona_ref="test@v1", instance_id="inst_1")
        with pytest.raises(Exception):  # PreContractError
            resolve_budget(-1, job)


# =============================================================================
# Test BudgetTracker
# =============================================================================


class TestBudgetTracker:
    """Tests for BudgetTracker class."""

    # -------------------------------------------------------------------------
    # Normal tracking tests
    # -------------------------------------------------------------------------

    def test_normal_tracking(self) -> None:
        """Track usage and calculate remaining."""
        tracker = BudgetTracker(1000)
        assert tracker.used == 0
        assert tracker.remaining() == 1000
        assert tracker.is_exhausted() is False

        tracker.record_usage(500)
        assert tracker.used == 500
        assert tracker.remaining() == 500
        assert tracker.is_exhausted() is False

        tracker.record_usage(300)
        assert tracker.used == 800
        assert tracker.remaining() == 200

    def test_exhaustion(self) -> None:
        """Tracker is exhausted when used >= budget."""
        tracker = BudgetTracker(1000)
        tracker.record_usage(1000)
        assert tracker.used == 1000
        assert tracker.remaining() == 0
        assert tracker.is_exhausted() is True

    def test_over_budget(self) -> None:
        """Usage can exceed budget (remaining = 0)."""
        tracker = BudgetTracker(100)
        tracker.record_usage(150)
        assert tracker.used == 150
        # Remaining is capped at 0
        assert tracker.remaining() == 0
        assert tracker.is_exhausted() is True

    # -------------------------------------------------------------------------
    # Unlimited budget tests
    # -------------------------------------------------------------------------

    def test_unlimited_never_exhausted(self) -> None:
        """Unlimited budget (None) is never exhausted."""
        tracker = BudgetTracker(None)
        tracker.record_usage(999999)
        assert tracker.used == 999999
        assert tracker.remaining() is None
        assert tracker.is_exhausted() is False
        assert tracker.should_enter_grace() is False

    def test_unlimited_no_grace(self) -> None:
        """Unlimited budget cannot enter grace."""
        tracker = BudgetTracker(None, grace_budget=500)
        tracker.record_usage(1000000)
        tracker.enter_grace()
        assert tracker.in_grace is False

    # -------------------------------------------------------------------------
    # Grace entry tests
    # -------------------------------------------------------------------------

    def test_grace_entry(self) -> None:
        """Enter grace when exhausted."""
        tracker = BudgetTracker(100, grace_budget=50)
        tracker.record_usage(100)
        assert tracker.is_exhausted() is True
        assert tracker.should_enter_grace() is True

        tracker.enter_grace()
        assert tracker.in_grace is True
        assert tracker.grace_start_usage == 100
        assert tracker.remaining() == 50  # Grace budget available
        assert tracker.is_exhausted() is False  # Grace resets exhaustion

    def test_grace_exhaustion(self) -> None:
        """Grace budget can be exhausted."""
        tracker = BudgetTracker(100, grace_budget=50)
        tracker.record_usage(100)  # Exhaust normal budget
        tracker.enter_grace()
        assert tracker.in_grace is True
        assert tracker.is_exhausted() is False

        tracker.record_usage(25)  # Use grace budget
        assert tracker.remaining() == 25
        assert tracker.is_grace_exhausted() is False

        tracker.record_usage(25)  # Use rest of grace
        assert tracker.remaining() == 0
        assert tracker.is_grace_exhausted() is True
        assert tracker.is_exhausted() is True

    def test_grace_not_entered_when_not_exhausted(self) -> None:
        """Cannot enter grace if not exhausted."""
        tracker = BudgetTracker(100, grace_budget=50)
        tracker.record_usage(50)  # Not exhausted yet
        assert tracker.is_exhausted() is False
        assert tracker.should_enter_grace() is False

        tracker.enter_grace()
        assert tracker.in_grace is False

    def test_grace_not_entered_with_zero_grace_budget(self) -> None:
        """Cannot enter grace if grace_budget is 0."""
        tracker = BudgetTracker(100, grace_budget=0)
        tracker.record_usage(100)
        assert tracker.is_exhausted() is True
        assert tracker.should_enter_grace() is False

        tracker.enter_grace()
        assert tracker.in_grace is False

    def test_grace_idempotent(self) -> None:
        """enter_grace is idempotent."""
        tracker = BudgetTracker(100, grace_budget=50)
        tracker.record_usage(100)
        tracker.enter_grace()
        assert tracker.in_grace is True
        grace_start = tracker.grace_start_usage

        tracker.enter_grace()  # Call again
        assert tracker.in_grace is True
        assert tracker.grace_start_usage == grace_start  # Unchanged

    def test_grace_remaining_tracking(self) -> None:
        """Grace remaining is tracked correctly."""
        tracker = BudgetTracker(100, grace_budget=50)
        tracker.record_usage(100)
        tracker.enter_grace()

        assert tracker.remaining() == 50
        assert tracker.is_grace_exhausted() is False

        tracker.record_usage(10)
        assert tracker.remaining() == 40
        assert tracker.is_grace_exhausted() is False

        tracker.record_usage(30)
        assert tracker.remaining() == 10
        assert tracker.is_grace_exhausted() is False

        tracker.record_usage(10)
        assert tracker.remaining() == 0
        assert tracker.is_grace_exhausted() is True

    # -------------------------------------------------------------------------
    # Contract tests
    # -------------------------------------------------------------------------

    def test_record_usage_negative_rejected(self) -> None:
        """record_usage rejects negative tokens."""
        tracker = BudgetTracker(100)
        with pytest.raises(Exception):  # PreContractError
            tracker.record_usage(-1)

    def test_budget_limit_zero_rejected(self) -> None:
        """BudgetTracker rejects budget_limit=0."""
        with pytest.raises(Exception):  # PreContractError
            BudgetTracker(0)

    def test_grace_budget_negative_rejected(self) -> None:
        """BudgetTracker rejects negative grace_budget."""
        with pytest.raises(Exception):  # PreContractError
            BudgetTracker(100, grace_budget=-1)


# =============================================================================
# Test Grace Helpers
# =============================================================================


class TestGraceHelpers:
    """Tests for grace_system_message and format_budget_status."""

    # -------------------------------------------------------------------------
    # grace_system_message tests
    # -------------------------------------------------------------------------

    def test_grace_message_content(self) -> None:
        """Grace message contains ctl_finish instruction."""
        msg = grace_system_message(500)
        assert "ctl_finish" in msg
        assert "grace" in msg.lower()
        assert "500" in msg

    def test_grace_message_zero_budget(self) -> None:
        """Grace message with zero budget still valid."""
        msg = grace_system_message(0)
        assert "ctl_finish" in msg
        assert len(msg) > 0

    def test_grace_message_token_mention(self) -> None:
        """Grace message mentions tokens."""
        msg = grace_system_message(1000)
        assert "token" in msg.lower()
        assert "1000" in msg

    def test_grace_message_negative_rejected(self) -> None:
        """grace_system_message rejects negative budget."""
        with pytest.raises(Exception):  # PreContractError
            grace_system_message(-1)

    # -------------------------------------------------------------------------
    # format_budget_status tests
    # -------------------------------------------------------------------------

    def test_format_budget_normal(self) -> None:
        """Normal budget shows used/remaining."""
        tracker = BudgetTracker(1000)
        tracker.record_usage(300)
        status = format_budget_status(tracker)
        assert "Budget:" in status
        assert "300" in status  # used
        assert "700" in status  # remaining
        assert "grace" not in status.lower()

    def test_format_budget_exhausted(self) -> None:
        """Exhausted budget shows exhausted status."""
        tracker = BudgetTracker(100)
        tracker.record_usage(100)
        status = format_budget_status(tracker)
        assert "Budget:" in status
        assert "exhausted" in status.lower() or "0" in status

    def test_format_budget_unlimited(self) -> None:
        """Unlimited budget shows unlimited status."""
        tracker = BudgetTracker(None)
        status = format_budget_status(tracker)
        assert "unlimited" in status.lower()

        tracker.record_usage(10000)
        status = format_budget_status(tracker)
        assert "unlimited" in status.lower()

    def test_format_budget_grace(self) -> None:
        """Grace mode shows grace budget status."""
        tracker = BudgetTracker(100, grace_budget=50)
        tracker.record_usage(100)  # Exhaust normal
        tracker.enter_grace()
        tracker.record_usage(20)  # Use some grace
        status = format_budget_status(tracker)
        assert "grace" in status.lower()
        assert "20" in status  # used in grace
        assert "30" in status  # remaining in grace

    def test_format_budget_grace_just_entered(self) -> None:
        """Just entered grace shows full grace budget."""
        tracker = BudgetTracker(200, grace_budget=100)
        tracker.record_usage(200)
        tracker.enter_grace()
        status = format_budget_status(tracker)
        assert "grace" in status.lower()
        assert "100" in status  # full grace available

    def test_format_budget_grace_exhausted(self) -> None:
        """Grace exhausted shows exhausted status."""
        tracker = BudgetTracker(100, grace_budget=30)
        tracker.record_usage(100)
        tracker.enter_grace()
        tracker.record_usage(30)  # Exhaust grace
        status = format_budget_status(tracker)
        assert "grace" in status.lower()
        # 30/30 used
        assert "30" in status


# =============================================================================
# Additional Coverage Tests
# =============================================================================


class TestBudgetCoverage:
    """Additional tests for edge cases and coverage."""

    def test_record_usage_zero(self) -> None:
        """record_usage(0) is a no-op."""
        tracker = BudgetTracker(100)
        tracker.record_usage(0)
        assert tracker.used == 0

    def test_multiple_record_usage(self) -> None:
        """Multiple record_usage calls accumulate."""
        tracker = BudgetTracker(1000)
        tracker.record_usage(100)
        tracker.record_usage(200)
        tracker.record_usage(50)
        assert tracker.used == 350

    def test_grace_budget_default(self) -> None:
        """BudgetTracker has default grace_budget of 500."""
        tracker = BudgetTracker(100)
        assert tracker.grace_budget == 500

    def test_custom_grace_budget(self) -> None:
        """BudgetTracker accepts custom grace_budget."""
        tracker = BudgetTracker(100, grace_budget=200)
        assert tracker.grace_budget == 200

    def test_is_grace_exhausted_not_in_grace(self) -> None:
        """is_grace_exhausted returns False when not in grace."""
        tracker = BudgetTracker(100)
        tracker.record_usage(100)
        # Exhausted but not in grace
        assert tracker.is_grace_exhausted() is False

    def test_remaining_caps_at_zero(self) -> None:
        """remaining() caps at zero, never negative."""
        tracker = BudgetTracker(50)
        tracker.record_usage(100)
        assert tracker.remaining() == 0

    def test_grace_remaining_caps_at_zero(self) -> None:
        """Grace remaining caps at zero."""
        tracker = BudgetTracker(50, grace_budget=30)
        tracker.record_usage(50)
        tracker.enter_grace()
        tracker.record_usage(100)  # Use way more than grace
        assert tracker.remaining() == 0
