"""Tests for the Context Distillation (Distill) module."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
from context_manager.distill import LLMLinguaCompressor
from context_manager.engine import ContextEngine
from context_manager.models import ContextBlock


class TestLLMLinguaCompressor:
    """Unit tests for the LLMLinguaCompressor wrapper."""

    def test_lazy_initialization(self):
        """Test that the heavy model isn't loaded until used."""
        compressor = LLMLinguaCompressor(model_name="test-model")
        assert compressor._compressor is None

    @patch("llmlingua.PromptCompressor")
    def test_compress_call(self, mock_cls):
        """Test that compress() delegates to the underlying library."""
        # Setup mock
        mock_instance = MagicMock()
        mock_instance.compress_prompt.return_value = {
            "compressed_prompt": "compressed text",
            "origin_tokens": 100,
            "compressed_tokens": 50,
        }
        mock_cls.return_value = mock_instance

        compressor = LLMLinguaCompressor()
        result = compressor.compress("original text", target_token_count=50)

        assert result == "compressed text"
        mock_cls.assert_called_once()  # initialized on first use
        mock_instance.compress_prompt.assert_called_once_with(
            prompt="original text",
            target_token=50,
            rate=0.5,
        )

    @patch("llmlingua.PromptCompressor")
    def test_compress_rate(self, mock_cls):
        """Test compression using rate."""
        mock_instance = MagicMock()
        mock_instance.compress_prompt.return_value = {
            "compressed_prompt": "short",
        }
        mock_cls.return_value = mock_instance

        compressor = LLMLinguaCompressor()
        compressor.compress("long text", rate=0.3)

        mock_instance.compress_prompt.assert_called_once_with(
            prompt="long text",
            target_token=-1,
            rate=0.3,
        )


class TestEngineIntegration:
    """Test that ContextEngine integrates the compressor correctly."""

    def test_compress_flagged_blocks(self):
        """Only blocks with can_compress=True should be compressed."""
        mock_compressor = MagicMock()
        mock_compressor.compress.side_effect = lambda t: f"compressed({t})"

        engine = ContextEngine(compressor=mock_compressor)

        b1 = ContextBlock(content="A", can_compress=True)
        b2 = ContextBlock(content="B", can_compress=False)  # default
        b3 = ContextBlock(content="C", can_compress=True)

        engine.add(b1).add(b2).add(b3)
        result = engine.compile()

        # Check content in compiled result
        contents = [m["content"] for m in result]
        assert "compressed(A)" in contents
        assert "B" in contents
        assert "compressed(C)" in contents

        # Verify compressor call count
        assert mock_compressor.compress.call_count == 2
        mock_compressor.compress.assert_any_call("A")
        mock_compressor.compress.assert_any_call("C")

    def test_no_compressor_configured(self):
        """If no compressor is provided, can_compress flag is ignored."""
        engine = ContextEngine(compressor=None)
        b1 = ContextBlock(content="A", can_compress=True)
        engine.add(b1)
        result = engine.compile()
        assert result[0]["content"] == "A"
