"""
tests/test_voice_budget.py

Run: pytest tests/ -v
"""

import asyncio
import warnings
from typing import Dict, List
import pytest

from voice_budget.compressors import (
    BudgetCompressor,
    DEFAULT_SUMMARY_PROMPT,
    SemanticTrimCompressor,
    SlidingWindowCompressor,
    SummariseTailCompressor,
)
from voice_budget.core import TTFTMeasurer
from voice_budget.wrapper import VoiceBudget, wrap



def make_messages(n_turns: int = 10, words_per_turn: int = 50) -> List[Dict]:
    """Generate a fake conversation of n_turns."""
    msgs = [{"role": "system", "content": "You are a helpful voice assistant."}]
    for i in range(n_turns):
        msgs.append({
            "role": "user",
            "content": f"Turn {i}: " + " ".join([f"word{j}" for j in range(words_per_turn)])
        })
        msgs.append({
            "role": "assistant",
            "content": f"Response {i}: " + " ".join([f"resp{j}" for j in range(words_per_turn)])
        })
    return msgs


async def fast_llm(messages, **kwargs) -> str:
    """Simulates a fast LLM (100ms TTFT)."""
    await asyncio.sleep(0.1)
    return "Fast response."


async def slow_llm(messages, **kwargs) -> str:
    """Simulates a slow LLM (1500ms TTFT) — like context accumulation."""
    # Simulate latency growing with message count
    delay = 0.1 + (len(messages) * 0.05)
    await asyncio.sleep(delay)
    return "Slow response."


async def variable_llm(messages, **kwargs) -> str:
    """High variance — simulates real production LLM."""
    import random
    delay = random.uniform(0.2, 2.0)
    await asyncio.sleep(delay)
    return "Variable response."


class TestTTFTMeasurer:

    def test_deprecated_params_warn(self):
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            TTFTMeasurer(target_ms=800, p95_trigger=0.9)
        assert any(issubclass(x.category, DeprecationWarning) for x in w)

    def test_no_stats_below_3_samples(self):
        m = TTFTMeasurer(target_ms=800)
        assert m.stats() is None
        m.record_sample(500, 100)
        m.record_sample(600, 120)
        assert m.stats() is None
        m.record_sample(700, 140)
        assert m.stats() is not None

    def test_p95_calculation(self):
        m = TTFTMeasurer(target_ms=800, window_size=100)
        for _ in range(90):
            m.record_sample(500.0, 100)
        for _ in range(10):
            m.record_sample(2000.0, 100)
        s = m.stats()
        assert s.p95_ms > 1000, "P95 should reflect the high outliers"
        assert s.p50_ms < 600, "P50 should reflect the majority"

    def test_should_compress_triggers_correctly(self):
        m = TTFTMeasurer(target_ms=800, window_size=10)
        # Below budget
        for _ in range(10):
            m.record_sample(400.0, 100)
        assert not m.should_compress()

        # Above budget
        m2 = TTFTMeasurer(target_ms=800, window_size=10)
        for _ in range(10):
            m2.record_sample(1200.0, 100)
        assert m2.should_compress()

    def test_token_counting(self):
        m = TTFTMeasurer(model="gpt-4o")
        msgs = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello, how are you?"},
            {"role": "assistant", "content": "I'm doing well, thank you!"},
        ]
        count = m.count_tokens(msgs)
        assert count > 0
        assert count < 200  # sanity check

    def test_compression_event_tracking(self):
        m = TTFTMeasurer(target_ms=800)
        for _ in range(5):
            m.record_sample(1500.0, 500)
        ev = m.record_compression("sliding_window", 500, 300, 1500.0)
        m.record_sample(700.0, 300, compressed=True, strategy="sliding_window")
        assert ev.ttft_after_ms == 700.0
        assert ev.delta_ms == 800.0  # 1500 - 700

    def test_rollback_detection(self):
        m = TTFTMeasurer(target_ms=800)
        for _ in range(5):
            m.record_sample(1200.0, 500)
        ev = m.record_compression("sliding_window", 500, 300, 1200.0)
        # Simulate compression making things WORSE
        m.record_sample(1500.0, 300, compressed=True)
        assert ev.rolled_back is True

    def test_snapshot_report_structure(self):
        m = TTFTMeasurer(target_ms=800)
        for _ in range(10):
            m.record_sample(600.0, 200)
        r = m.snapshot_report()
        assert "total_turns" in r
        assert "current_p95_ms" in r
        assert "total_tokens_saved" in r
        assert r["total_turns"] == 10


class TestSlidingWindowCompressor:

    def test_removes_oldest_non_system(self):
        c = SlidingWindowCompressor()
        msgs = make_messages(n_turns=20)
        original_len = len(msgs)
        compressed, removed = c.compress(msgs, target_tokens=500, current_tokens=2000)
        assert len(compressed) < original_len
        # System message should still be first
        assert compressed[0]["role"] == "system"
        assert removed > 0

    def test_never_removes_system_message(self):
        c = SlidingWindowCompressor()
        msgs = make_messages(n_turns=5)
        system_content = msgs[0]["content"]
        compressed, _ = c.compress(msgs, target_tokens=100, current_tokens=1000)
        assert compressed[0]["content"] == system_content

    def test_preserves_minimum_messages(self):
        c = SlidingWindowCompressor()
        msgs = make_messages(n_turns=2)
        compressed, _ = c.compress(msgs, target_tokens=1, current_tokens=1000)
        # Should always keep system + at least 1 user + 1 assistant
        assert len(compressed) >= 2


class TestSemanticTrimCompressor:

    def test_falls_back_gracefully_without_sentence_transformers(self):
        c = SemanticTrimCompressor()
        c._available = False  # force fallback
        msgs = make_messages(n_turns=10)
        compressed, removed = c.compress(msgs, target_tokens=500, current_tokens=2000)
        assert len(compressed) < len(msgs) or removed >= 0

    def test_preserves_system_message(self):
        c = SemanticTrimCompressor()
        c._available = False
        msgs = make_messages(n_turns=5)
        compressed, _ = c.compress(msgs, target_tokens=100, current_tokens=1000)
        assert compressed[0]["role"] == "system"


class TestBudgetCompressor:

    @pytest.mark.asyncio
    async def test_compress_reduces_tokens(self):
        c = BudgetCompressor(target_tokens=300, use_semantic=False, use_summarise=False)
        msgs = make_messages(n_turns=20)

        def counter(m):
            return sum(len(str(msg.get("content", "")).split()) * 4 // 3 for msg in m)

        current = counter(msgs)
        compressed, strategy, removed = await c.compress(msgs, current, counter)
        assert counter(compressed) < current
        assert strategy == "sliding_window"
        assert removed > 0


class TestVoiceBudget:

    @pytest.mark.asyncio
    async def test_basic_call_passthrough(self):
        budget = VoiceBudget(fast_llm, target_ms=800)
        msgs = make_messages(n_turns=3)
        result = await budget(msgs)
        assert result == "Fast response."

    @pytest.mark.asyncio
    async def test_stats_accumulate(self):
        budget = VoiceBudget(fast_llm, target_ms=800, window_size=10)
        msgs = make_messages(n_turns=3)
        for _ in range(5):
            await budget(msgs)
        s = budget.stats()
        assert s is not None
        assert s.turn == 5
        assert s.p50_ms > 0

    @pytest.mark.asyncio
    async def test_compression_triggers_on_slow_llm(self):
        """Simulate context accumulation: latency grows with message count."""
        budget = VoiceBudget(
            slow_llm,
            target_ms=500,      # tight budget
            window_size=5,
            token_budget=200,
            use_semantic=False,
            verbose=False,
        )
        # Run many turns to trigger compression
        msgs = make_messages(n_turns=30)
        for _ in range(10):
            await budget(msgs)

        # Should have triggered at least one compression
        history = budget.compression_history()
        s = budget.stats()
        assert s is not None
        assert s.turn == 10
        assert len(history) > 0, "Expected at least one compression event"
        assert history[0].strategy == "sliding_window"
        assert history[0].tokens_before > history[0].tokens_after

    @pytest.mark.asyncio
    async def test_wrap_convenience_function(self):
        managed = wrap(fast_llm, target_ms=800, verbose=False)
        msgs = make_messages(n_turns=2)
        result = await managed(msgs)
        assert isinstance(result, str)

    @pytest.mark.asyncio
    async def test_on_compression_callback_fires(self):
        fired = []

        def on_comp(ev):
            fired.append(ev)

        budget = VoiceBudget(
            slow_llm,
            target_ms=200,      # very tight — will trigger quickly
            window_size=4,
            token_budget=100,
            use_semantic=False,
            on_compression=on_comp,
            verbose=False,
        )
        msgs = make_messages(n_turns=20)
        for _ in range(8):
            await budget(msgs)

        # Callback may or may not have fired depending on timing
        # — just check it's callable and didn't crash
        assert True  # test passes if no exception

    def test_print_report_no_crash(self):
        budget = VoiceBudget(fast_llm)
        # Should not crash even with no data
        budget.print_report()

    @pytest.mark.asyncio
    async def test_report_structure(self):
        budget = VoiceBudget(fast_llm, target_ms=800)
        msgs = make_messages(n_turns=3)
        for _ in range(10):
            await budget(msgs)
        r = budget.report()
        assert "total_turns" in r
        assert r["total_turns"] == 10
        assert "strategies_used" in r
        assert "total_tokens_saved" in r


class TestSummariseTailCompressor:

    @pytest.mark.asyncio
    async def test_default_summary_prompt_used(self):
        """When no custom prompt given, DEFAULT_SUMMARY_PROMPT is used."""
        captured_prompts = []

        async def fake_llm(messages, **kwargs):
            captured_prompts.append(messages[0]["content"])
            return "Summary of conversation."

        c = SummariseTailCompressor(llm_fn=fake_llm)
        msgs = make_messages(n_turns=5)
        await c.async_compress(msgs, target_tokens=100, current_tokens=1000)
        assert len(captured_prompts) == 1
        assert DEFAULT_SUMMARY_PROMPT in captured_prompts[0]

    @pytest.mark.asyncio
    async def test_custom_summary_prompt_used(self):
        """When a custom prompt is given, it replaces the default."""
        captured_prompts = []
        custom = "Condense this chat into bullet points."

        async def fake_llm(messages, **kwargs):
            captured_prompts.append(messages[0]["content"])
            return "Bullet summary."

        c = SummariseTailCompressor(llm_fn=fake_llm, summary_prompt=custom)
        msgs = make_messages(n_turns=5)
        await c.async_compress(msgs, target_tokens=100, current_tokens=1000)
        assert len(captured_prompts) == 1
        assert custom in captured_prompts[0]
        assert DEFAULT_SUMMARY_PROMPT not in captured_prompts[0]

    @pytest.mark.asyncio
    async def test_tail_turns_preserves_last_user_message(self):
        """Regression: short chats ending with a user turn keep that turn after compression."""
        async def fake_llm(messages, **kwargs):
            return "Some summary."

        c = SummariseTailCompressor(llm_fn=fake_llm, tail_turns=100)
        msgs = [
            {"role": "system", "content": "You are a helpful voice assistant."},
            {"role": "user", "content": "First user message."},
            {"role": "assistant", "content": "Assistant reply."},
            {"role": "user", "content": "Final user message should be preserved."},
        ]
        compressed, _ = await c.async_compress(
            msgs, target_tokens=100, current_tokens=1000
        )
        # The last user message must still be present.
        assert any(
            m["role"] == "user" and m["content"] == "Final user message should be preserved."
            for m in compressed
        ), "Last user message was dropped by compression"
        # It should be the last non-summary message (or at least present at end).
        user_msgs = [m for m in compressed if m["role"] == "user" and "summary" not in m.get("content", "").lower()]
        assert user_msgs[-1]["content"] == "Final user message should be preserved."


class TestEffectiveTarget:

    @pytest.mark.asyncio
    async def test_effective_target_overrides_default(self):
        """BudgetCompressor.compress with effective_target lower than target_tokens."""
        c = BudgetCompressor(target_tokens=5000, use_semantic=False, use_summarise=False)
        msgs = make_messages(n_turns=20)

        def counter(m):
            return sum(len(str(msg.get("content", "")).split()) * 4 // 3 for msg in m)

        current = counter(msgs)
        # effective_target much lower than target_tokens — should compress
        compressed, strategy, removed = await c.compress(
            msgs, current, counter, effective_target=300
        )
        assert counter(compressed) < current
        assert removed > 0

    @pytest.mark.asyncio
    async def test_dynamic_target_when_tokens_under_budget(self):
        """When tokens < token_budget but TTFT is high, wrapper uses proportional target."""
        budget = VoiceBudget(
            slow_llm,
            target_ms=300,
            window_size=5,
            token_budget=50000,  # very high — tokens will always be "under budget"
            use_semantic=False,
            verbose=False,
        )
        msgs = make_messages(n_turns=10)
        for _ in range(8):
            await budget(msgs)
        history = budget.compression_history()
        # Compression should have been triggered despite tokens < token_budget
        assert history, "Expected at least one compression event"
        assert history[0].tokens_after < history[0].tokens_before


class TestRollback:

    @pytest.mark.asyncio
    async def test_rollback_restores_context(self):
        """After harmful compression, next call uses last known good messages."""
        call_count = [0]
        received_msgs = []

        async def tracking_llm(messages, **kwargs):
            call_count[0] += 1
            received_msgs.append(len(messages))
            # First 5 calls: fast. Then one slow call to trigger compression,
            # then make post-compression call even slower to trigger rollback.
            delay = 0.05
            if call_count[0] > 3:
                delay = 0.05 + len(messages) * 0.03
            await asyncio.sleep(delay)
            return "OK"

        budget = VoiceBudget(
            tracking_llm,
            target_ms=100,
            window_size=4,
            token_budget=50000,
            use_semantic=False,
            verbose=False,
        )
        msgs = make_messages(n_turns=15)
        for _ in range(10):
            await budget(msgs)

        # Verify rollback was detected and didn't crash
        history = budget.compression_history()
        assert budget.stats() is not None
        # If a rollback happened, verify it was flagged
        rolled_back = [ev for ev in history if ev.rolled_back]
        if rolled_back:
            # After rollback, next call should have received more messages
            # (the original uncompressed context)
            assert len(received_msgs) >= 2, "Expected multiple LLM calls"


class TestStreaming:

    @pytest.mark.asyncio
    async def test_streaming_ttft_recorded(self):
        """Streaming LLM (async generator) should still record TTFT samples."""
        call_count = [0]

        async def streaming_llm(messages, **kwargs):
            call_count[0] += 1
            # Simulate a streaming response that yields chunks
            async def _gen():
                await asyncio.sleep(0.05)
                yield "Hello"
                yield " world"
            return _gen()

        budget = VoiceBudget(
            streaming_llm,
            target_ms=5000,
            window_size=5,
            token_budget=50000,
            use_semantic=False,
            verbose=False,
        )
        msgs = make_messages(n_turns=3)

        for _ in range(4):
            result = await budget(msgs)
            # Consume the stream
            chunks = []
            async for chunk in result:
                chunks.append(chunk)
            assert chunks == ["Hello", " world"]

        stats = budget.stats()
        assert stats is not None
        assert stats.turn == 4
        assert stats.p50_ms > 0