"""
Tests for cortex/hippo/budget.py — cost circuit breaker.

Issue: #1178
"""

from __future__ import annotations

import time
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from cortex.hippo.budget import (
    BudgetDecision,
    BudgetEvent,
    BudgetStats,
    CompactionBudget,
    _current_bucket,
)


# ── Helpers ──────────────────────────────────────────────────────────

def _make_budget(**kwargs) -> CompactionBudget:
    defaults = dict(
        max_summarization_tokens_per_hour=100_000,
        global_max_summarization_tokens_per_hour=500_000,
        circuit_breaker_threshold=3,
        circuit_breaker_cooldown_seconds=300,
        min_compaction_reduction_ratio=0.10,
        budget_warning_threshold=0.80,
    )
    defaults.update(kwargs)
    return CompactionBudget(**defaults)


# ── test_budget_tracks_tokens ────────────────────────────────────────

def test_budget_tracks_tokens():
    budget = _make_budget()
    budget.record_summarization("s1", 1000, 500)
    budget.record_summarization("s1", 2000, 800)

    decision = budget.check_budget("s1")
    assert decision.allowed is True
    assert decision.reason == "ok"
    # 100_000 - (1000+500+2000+800) = 95_700
    assert decision.remaining_tokens == 95_700


# ── test_budget_blocks_when_exceeded ─────────────────────────────────

def test_budget_blocks_when_exceeded():
    budget = _make_budget(max_summarization_tokens_per_hour=5000)
    budget.record_summarization("s1", 3000, 2500)  # 5500 > 5000

    decision = budget.check_budget("s1")
    assert decision.allowed is False
    assert decision.reason == "budget_exceeded"
    assert decision.remaining_tokens == 0


# ── test_budget_sliding_window ───────────────────────────────────────

def test_budget_sliding_window():
    budget = _make_budget(max_summarization_tokens_per_hour=5000)

    # Inject tokens into an old bucket (> 1 hour ago)
    with budget._lock:
        sess = budget._get_session("s1")
        old_bucket = _current_bucket() - 61  # 61 minutes ago
        sess.token_buckets[old_bucket] = 4000
        budget._global_buckets[old_bucket] = 4000

    # Old usage should be expired
    decision = budget.check_budget("s1")
    assert decision.allowed is True
    assert decision.remaining_tokens == 5000

    # Add current usage
    budget.record_summarization("s1", 3000, 0)
    decision = budget.check_budget("s1")
    assert decision.allowed is True
    assert decision.remaining_tokens == 2000


# ── test_circuit_breaker_trips ───────────────────────────────────────

def test_circuit_breaker_trips():
    budget = _make_budget(circuit_breaker_threshold=3)

    # 3 consecutive low-reduction compactions
    budget.record_compaction_result("s1", 0.05)
    budget.record_compaction_result("s1", 0.03)

    # Not yet tripped
    decision = budget.check_budget("s1")
    assert decision.allowed is True

    budget.record_compaction_result("s1", 0.02)

    # Now tripped
    decision = budget.check_budget("s1")
    assert decision.allowed is False
    assert decision.reason == "circuit_breaker_tripped"
    assert decision.cooldown_until is not None


# ── test_circuit_breaker_cooldown ────────────────────────────────────

def test_circuit_breaker_cooldown():
    budget = _make_budget(
        circuit_breaker_threshold=3,
        circuit_breaker_cooldown_seconds=10,
    )

    # Trip the breaker
    for _ in range(3):
        budget.record_compaction_result("s1", 0.01)

    decision = budget.check_budget("s1")
    assert decision.allowed is False

    # Simulate cooldown expiry
    with budget._lock:
        sess = budget._get_session("s1")
        sess.cooldown_until = time.time() - 1  # expired

    # Should allow retry
    decision = budget.check_budget("s1")
    assert decision.allowed is True


# ── test_circuit_breaker_exponential_backoff ─────────────────────────

def test_circuit_breaker_exponential_backoff():
    budget = _make_budget(
        circuit_breaker_threshold=2,
        circuit_breaker_cooldown_seconds=100,
    )

    # First trip
    budget.record_compaction_result("s1", 0.01)
    budget.record_compaction_result("s1", 0.01)

    with budget._lock:
        sess = budget._get_session("s1")
        assert sess.trip_count == 1
        # First trip = 100s cooldown (2^0 * 100)
        expected_cooldown_1 = 100
        cooldown_1 = sess.cooldown_until - (time.time())
        assert abs(cooldown_1 - expected_cooldown_1) < 2

    # Expire cooldown, allow retry
    with budget._lock:
        sess.cooldown_until = time.time() - 1
        # breaker is still tripped but cooldown expired → allows retry

    # Retry fails again → re-trip with 2x cooldown
    # First, a good compaction would clear it, but a bad one won't.
    # After cooldown, consecutive_low is still >=threshold, but breaker_tripped is True.
    # record_compaction_result with low ratio when breaker already tripped:
    # consecutive_low goes up but no re-trip because already tripped.
    # We need to reset breaker_tripped=False to allow re-trip on next consecutive failures.
    # Actually the logic: after cooldown expires, check_budget allows,
    # then if compaction is bad, record_compaction_result will see consecutive_low still >= threshold
    # but breaker is already tripped. Let's reset consecutive_low and re-trip.
    
    # Reset the breaker so it can re-trip
    with budget._lock:
        sess.breaker_tripped = False
        sess.consecutive_low = 0

    # Second trip
    budget.record_compaction_result("s1", 0.01)
    budget.record_compaction_result("s1", 0.01)

    with budget._lock:
        assert sess.trip_count == 2
        # Second trip = 200s cooldown (2^1 * 100)
        cooldown_2 = sess.cooldown_until - time.time()
        assert abs(cooldown_2 - 200) < 2


# ── test_circuit_breaker_reset ───────────────────────────────────────

def test_circuit_breaker_reset():
    budget = _make_budget(circuit_breaker_threshold=2)

    # Trip it
    budget.record_compaction_result("s1", 0.01)
    budget.record_compaction_result("s1", 0.01)

    decision = budget.check_budget("s1")
    assert decision.allowed is False

    # Manual reset
    budget.reset_session("s1")

    decision = budget.check_budget("s1")
    assert decision.allowed is True


# ── test_global_budget ───────────────────────────────────────────────

def test_global_budget():
    budget = _make_budget(
        max_summarization_tokens_per_hour=1_000_000,  # high per-session
        global_max_summarization_tokens_per_hour=10_000,  # low global
    )

    # Multiple sessions eat into global budget
    budget.record_summarization("s1", 4000, 0)
    budget.record_summarization("s2", 4000, 0)
    budget.record_summarization("s3", 3000, 0)  # total 11_000 > 10_000

    decision = budget.check_budget("s4")
    assert decision.allowed is False
    assert decision.reason == "global_budget_exceeded"


# ── test_budget_callback ────────────────────────────────────────────

def test_budget_callback():
    events: list[BudgetEvent] = []

    def on_event(session_id: str, event: BudgetEvent):
        events.append(event)

    budget = _make_budget(
        max_summarization_tokens_per_hour=10_000,
        budget_warning_threshold=0.50,
        circuit_breaker_threshold=2,
        on_budget_event=on_event,
    )

    # Trigger warning at 50%
    budget.record_summarization("s1", 5500, 0)  # 55% > 50%
    assert any(e.type == "warning" for e in events)

    # Trigger exceeded
    events.clear()
    budget.record_summarization("s1", 5000, 0)  # total 10500 > 10000
    budget.check_budget("s1")
    assert any(e.type == "exceeded" for e in events)

    # Trigger breaker
    events.clear()
    budget.record_compaction_result("s1", 0.01)
    budget.record_compaction_result("s1", 0.01)
    assert any(e.type == "breaker_tripped" for e in events)

    # Trigger reset
    events.clear()
    budget.reset_session("s1")
    assert any(e.type == "breaker_reset" for e in events)


# ── test_engine_respects_budget ──────────────────────────────────────

@pytest.mark.asyncio
async def test_engine_respects_budget():
    """Integration: after_turn skips compaction when breaker tripped."""
    from cortex.hippo.config import HippoConfig
    from cortex.hippo.engine import HippoEngine

    config = HippoConfig(
        enabled=True,
        circuit_breaker_threshold=1,
    )
    engine = HippoEngine(config=config)
    await engine.initialize(db_path=":memory:", owner_id="test")

    # Trip the breaker
    assert engine._budget is not None
    conv_id = engine._conversation.conversation_id
    engine._budget.record_compaction_result(str(conv_id), 0.01)

    # Verify breaker is tripped
    decision = engine._budget.check_budget(str(conv_id))
    assert decision.allowed is False

    # after_turn should skip compaction silently (no error)
    messages = [
        {"role": "user", "content": "hello"},
        {"role": "assistant", "content": "hi"},
    ]
    # This should not raise, should just log warning and skip compaction
    await engine.after_turn(
        conversation_id=conv_id,
        messages=messages,
        pre_prompt_count=0,
        token_budget=100_000,
    )

    await engine.shutdown()


# ── test_compaction_records_reduction_ratio ───────────────────────────

def test_compaction_records_reduction_ratio():
    budget = _make_budget()

    budget.record_compaction_result("s1", 0.30)
    budget.record_compaction_result("s1", 0.25)
    budget.record_compaction_result("s1", 0.05)

    stats = budget.get_stats()
    s1 = stats.sessions["s1"]
    assert s1.last_reduction_ratios == [0.30, 0.25, 0.05]
    assert s1.consecutive_low_reductions == 1  # only last one was low


# ── test_good_compaction_clears_breaker ──────────────────────────────

def test_good_compaction_clears_breaker():
    budget = _make_budget(circuit_breaker_threshold=2)

    # Trip
    budget.record_compaction_result("s1", 0.01)
    budget.record_compaction_result("s1", 0.01)

    decision = budget.check_budget("s1")
    assert decision.allowed is False

    # Expire cooldown
    with budget._lock:
        sess = budget._get_session("s1")
        sess.cooldown_until = time.time() - 1

    # Now allowed (cooldown expired)
    decision = budget.check_budget("s1")
    assert decision.allowed is True

    # Good compaction clears the breaker
    budget.record_compaction_result("s1", 0.30)

    # Should be fully clear now
    stats = budget.get_stats()
    assert stats.sessions["s1"].breaker_tripped is False
    assert stats.sessions["s1"].breaker_trip_count == 0


# ── test_get_stats ───────────────────────────────────────────────────

def test_get_stats():
    budget = _make_budget()
    budget.record_summarization("s1", 1000, 500)
    budget.record_summarization("s2", 2000, 0)

    stats = budget.get_stats()
    assert isinstance(stats, BudgetStats)
    assert stats.global_tokens_used_last_hour == 3500
    assert "s1" in stats.sessions
    assert "s2" in stats.sessions
    assert stats.sessions["s1"].tokens_used_last_hour == 1500
    assert stats.sessions["s2"].tokens_used_last_hour == 2000
