"""Tests for context budget management."""

import pytest

from oclawma.context import BudgetExceededError, BudgetThreshold, ContextBudget


class TestContextBudget:
    """Test suite for ContextBudget class."""

    def test_default_budget_initialization(self):
        """Test budget initializes with correct defaults."""
        budget = ContextBudget()
        assert budget.total_budget == 8192
        assert budget.warning_threshold == 6144  # 75%
        assert budget.critical_threshold == 7168  # 87%
        assert budget.used_tokens == 0
        assert budget.remaining_tokens == 8192
        assert budget.usage_percent == 0.0

    def test_custom_budget_initialization(self):
        """Test budget can be initialized with custom values."""
        budget = ContextBudget(
            total_budget=4096,
            warning_threshold=3000,
            critical_threshold=3500,
        )
        assert budget.total_budget == 4096
        assert budget.warning_threshold == 3000
        assert budget.critical_threshold == 3500

    def test_allocate_tokens(self):
        """Test token allocation."""
        budget = ContextBudget()

        allocated = budget.allocate(1000)
        assert allocated == 1000
        assert budget.used_tokens == 1000
        assert budget.remaining_tokens == 7192

    def test_allocate_multiple_times(self):
        """Test multiple allocations."""
        budget = ContextBudget()

        budget.allocate(1000)
        budget.allocate(2000)
        budget.allocate(500)

        assert budget.used_tokens == 3500
        assert budget.remaining_tokens == 8192 - 3500

    def test_allocate_negative_raises_error(self):
        """Test that negative allocation raises ValueError."""
        budget = ContextBudget()

        with pytest.raises(ValueError, match="Cannot allocate negative tokens"):
            budget.allocate(-100)

    def test_allocate_zero_returns_zero(self):
        """Test that allocating zero returns zero."""
        budget = ContextBudget()
        assert budget.allocate(0) == 0
        assert budget.used_tokens == 0

    def test_release_tokens(self):
        """Test releasing tokens back to budget."""
        budget = ContextBudget()

        budget.allocate(5000)
        released = budget.release(2000)

        assert released == 2000
        assert budget.used_tokens == 3000

    def test_release_more_than_used(self):
        """Test releasing more tokens than used only releases what's available."""
        budget = ContextBudget()

        budget.allocate(1000)
        released = budget.release(2000)

        assert released == 1000
        assert budget.used_tokens == 0

    def test_release_negative_raises_error(self):
        """Test that releasing negative tokens raises ValueError."""
        budget = ContextBudget()

        with pytest.raises(ValueError, match="Cannot release negative tokens"):
            budget.release(-100)

    def test_budget_exceeded_strict(self):
        """Test that budget exceeded raises error in strict mode."""
        budget = ContextBudget(total_budget=1000)
        budget.allocate(800)

        with pytest.raises(BudgetExceededError) as exc_info:
            budget.allocate(300)  # Would exceed 1000

        assert "CONTEXT BUDGET EXHAUSTED" in str(exc_info.value)
        assert exc_info.value.used == 1100
        assert exc_info.value.limit == 1000

    def test_budget_exceeded_non_strict(self):
        """Test that budget exceeded in non-strict mode allocates what it can."""
        budget = ContextBudget(total_budget=1000)
        budget.allocate(800)

        allocated = budget.allocate(300, strict=False)

        assert allocated == 200  # Only allocated what was available
        assert budget.used_tokens == 1000  # Budget is now full

    def test_can_allocate(self):
        """Test can_allocate method."""
        budget = ContextBudget(total_budget=1000)
        budget.allocate(500)

        assert budget.can_allocate(400) is True
        assert budget.can_allocate(600) is False

    def test_reset(self):
        """Test reset clears all tokens."""
        budget = ContextBudget()

        budget.allocate(5000)
        budget.reset()

        assert budget.used_tokens == 0
        assert budget.remaining_tokens == 8192

    def test_get_status_normal(self):
        """Test status when budget is normal."""
        budget = ContextBudget()
        budget.allocate(1000)  # Well below warning

        status = budget.get_status()

        assert status.used_tokens == 1000
        assert status.total_tokens == 8192
        assert status.remaining_tokens == 7192
        assert status.threshold == BudgetThreshold.NORMAL
        assert status.warning_message is None

    def test_get_status_warning(self):
        """Test status at warning threshold (75%)."""
        budget = ContextBudget()
        budget.allocate(budget.warning_threshold)

        status = budget.get_status()

        assert status.threshold == BudgetThreshold.WARNING
        assert status.warning_message is not None
        assert "75%" in status.warning_message

    def test_get_status_critical(self):
        """Test status at critical threshold (87%)."""
        budget = ContextBudget()
        budget.allocate(budget.critical_threshold)

        status = budget.get_status()

        assert status.threshold == BudgetThreshold.CRITICAL
        assert status.warning_message is not None
        assert "87%" in status.warning_message

    def test_get_status_exhausted(self):
        """Test status when budget is exhausted."""
        budget = ContextBudget(total_budget=1000)
        budget.allocate(1000)

        status = budget.get_status()

        assert status.threshold == BudgetThreshold.EXHAUSTED
        assert status.warning_message is not None
        assert "exhausted" in status.warning_message

    def test_warning_callback_triggered(self):
        """Test warning callback is triggered at threshold."""
        warnings = []

        def on_warning(msg):
            warnings.append(msg)

        budget = ContextBudget(on_warning=on_warning)
        budget.allocate(budget.warning_threshold)

        assert len(warnings) == 1

    def test_critical_callback_triggered(self):
        """Test critical callback is triggered at threshold."""
        criticals = []

        def on_critical(msg):
            criticals.append(msg)

        budget = ContextBudget(on_critical=on_critical)
        budget.allocate(budget.critical_threshold)

        assert len(criticals) == 1

    def test_warning_not_triggered_twice(self):
        """Test warning callback only triggers once per threshold crossing."""
        warnings = []

        def on_warning(msg):
            warnings.append(msg)

        budget = ContextBudget(on_warning=on_warning)

        budget.allocate(budget.warning_threshold)
        budget.allocate(100)  # Still above threshold

        assert len(warnings) == 1

    def test_warning_resets_after_release(self):
        """Test warning flag resets when tokens are released."""
        warnings = []

        def on_warning(msg):
            warnings.append(msg)

        budget = ContextBudget(on_warning=on_warning)

        # First, allocate to just trigger warning
        budget.allocate(budget.warning_threshold)
        assert len(warnings) == 1

        budget.release(1000)  # Back below threshold
        # Calculate how many more we can allocate to re-cross threshold
        needed = budget.warning_threshold - budget.used_tokens
        budget.allocate(needed)  # Re-cross threshold

        assert len(warnings) == 2  # Triggered again

    def test_estimate_tokens(self):
        """Test token estimation."""
        budget = ContextBudget()

        # Roughly 4 chars per token
        assert budget.estimate_tokens("hello") == 1
        assert budget.estimate_tokens("a" * 100) == 25
        assert budget.estimate_tokens("") == 1  # Minimum 1

    def test_visualize_normal(self):
        """Test visualization in normal state."""
        budget = ContextBudget()
        budget.allocate(1000)

        viz = budget.visualize(width=40)

        assert "✅" in viz
        assert "Context Budget:" in viz
        assert "1,000 / 8,192" in viz
        assert "12.2%" in viz or "12.3%" in viz

    def test_visualize_warning(self):
        """Test visualization at warning level."""
        budget = ContextBudget()
        budget.allocate(budget.warning_threshold)

        viz = budget.visualize(width=40)

        assert "⚠️" in viz
        assert "75.0%" in viz

    def test_visualize_critical(self):
        """Test visualization at critical level."""
        budget = ContextBudget()
        budget.allocate(budget.critical_threshold)

        viz = budget.visualize(width=40)

        assert "🚨" in viz
        assert "87.5%" in viz

    def test_visualize_exhausted(self):
        """Test visualization when exhausted."""
        budget = ContextBudget(total_budget=1000)
        budget.allocate(1000)

        viz = budget.visualize(width=40)

        assert "🛑" in viz
        assert "100.0%" in viz

    def test_to_dict(self):
        """Test dictionary conversion."""
        budget = ContextBudget()
        budget.allocate(1000)

        data = budget.to_dict()

        assert data["total_budget"] == 8192
        assert data["used_tokens"] == 1000
        assert data["remaining_tokens"] == 7192
        assert "usage_percent" in data
        assert data["threshold"] == "normal"


class TestBudgetExceededError:
    """Test suite for BudgetExceededError."""

    def test_error_message_default(self):
        """Test default error message."""
        error = BudgetExceededError(used=8500, limit=8192)

        assert "CONTEXT BUDGET EXHAUSTED" in str(error)
        assert error.used == 8500
        assert error.limit == 8192

    def test_error_message_custom(self):
        """Test custom error message."""
        error = BudgetExceededError(used=100, limit=100, message="Custom message")

        assert str(error) == "Custom message"

    def test_error_includes_helpful_tips(self):
        """Test error message includes helpful tips."""
        error = BudgetExceededError(used=8500, limit=8192)

        message = str(error)
        assert "/compact" in message
        assert "new conversation" in message
        assert "/clear" in message
