"""
Tests for ARES Context Compression Layer.
"""

import pytest
import torch
from ares_compression import (
    ContextCompressor,
    CompressionConfig,
    CompressionMode,
)


class TestContextCompressor:
    """Test suite for ContextCompressor."""

    @pytest.fixture
    def sample_context(self):
        """Create sample context for testing."""
        return torch.randn(512, 768)

    @pytest.fixture
    def compressor(self):
        """Create compressor instance."""
        return ContextCompressor.from_preset(CompressionMode.BALANCED, device="cpu")

    def test_initialization(self, compressor):
        """Test compressor initialization."""
        assert compressor is not None
        assert compressor.config is not None
        assert compressor.config.mode == CompressionMode.BALANCED

    def test_compress_returns_result(self, compressor, sample_context):
        """Test that compress returns CompressionResult."""
        result = compressor.compress(sample_context)
        assert result is not None
        assert hasattr(result, 'compressed_context')
        assert hasattr(result, 'compression_ratio')
        assert hasattr(result, 'memory_saved')
        assert hasattr(result, 'latency_ms')

    def test_compression_reduces_tokens(self, compressor, sample_context):
        """Test that compression reduces token count."""
        original_length = sample_context.shape[0]
        result = compressor.compress(sample_context)
        assert result.compressed_length <= original_length

    def test_compression_saves_memory(self, compressor, sample_context):
        """Test that compression saves memory."""
        result = compressor.compress(sample_context)
        assert result.memory_saved >= 0

    def test_pruned_attention_mask_is_returned(self, compressor, sample_context):
        """Test that the pruned attention mask is returned with matching length."""
        attention_mask = torch.ones(sample_context.shape[0])
        attention_mask[-32:] = 0
        result = compressor.compress(sample_context, attention_mask=attention_mask)
        assert result.attention_mask is not None
        assert result.attention_mask.shape[0] == result.compressed_length
        assert torch.all(result.attention_mask >= 0)

    def test_stats_tracking(self, compressor, sample_context):
        """Test that compressor tracks statistics."""
        initial_count = compressor.get_stats()['total_compressions']
        compressor.compress(sample_context)
        stats = compressor.get_stats()
        assert stats['total_compressions'] == initial_count + 1

    def test_different_modes(self, sample_context):
        """Test different compression modes."""
        modes = [
            CompressionMode.CONSERVATIVE,
            CompressionMode.BALANCED,
            CompressionMode.AGGRESSIVE,
        ]
        for mode in modes:
            compressor = ContextCompressor.from_preset(mode, device="cpu")
            result = compressor.compress(sample_context)
            assert result is not None
            assert result.compressed_length > 0

    def test_custom_config(self, sample_context):
        """Test custom configuration."""
        config = CompressionConfig(
            mode=CompressionMode.CUSTOM,
            semantic_pruning={
                'enabled': True,
                'target_ratio': 0.5,
                'min_tokens': 100
            }
        )
        compressor = ContextCompressor.from_config(config)
        result = compressor.compress(sample_context)
        assert result is not None


class TestCompressionConfig:
    """Test suite for CompressionConfig."""

    def test_preset_modes(self):
        """Test that preset modes create valid configs."""
        modes = [
            CompressionMode.CONSERVATIVE,
            CompressionMode.BALANCED,
            CompressionMode.AGGRESSIVE,
        ]
        for mode in modes:
            config = CompressionConfig(mode=mode)
            assert config.mode == mode
            assert config.semantic_pruning is not None
            assert config.dynamic_precision is not None

    def test_custom_config_validation(self):
        """Test custom configuration validation."""
        config = CompressionConfig(
            mode=CompressionMode.CUSTOM,
            max_memory_reduction=0.7,
            max_latency_overhead=0.15
        )
        assert config.max_memory_reduction == 0.7
        assert config.max_latency_overhead == 0.15


@pytest.mark.integration
class TestIntegration:
    """Integration tests."""

    def test_full_pipeline(self):
        """Test full compression pipeline."""
        compressor = ContextCompressor.from_preset(CompressionMode.BALANCED)
        compressor = compressor.to("cpu")

        # Simulate multiple compressions
        contexts = [
            torch.randn(256, 768),
            torch.randn(512, 768),
            torch.randn(1024, 768),
        ]

        results = []
        for context in contexts:
            result = compressor.compress(context)
            results.append(result)
            assert result is not None
            assert result.compressed_length > 0

        # Check statistics
        stats = compressor.get_stats()
        assert stats['total_compressions'] == 3
        assert stats['avg_compression_ratio'] > 0
