"""Tests for SageAttention layer.

Tests cover:
- Basic attention computation
- Integration with AttentionBackend
- Integration with KVCacheManager
- Layer replacement functionality
"""

from __future__ import annotations

import pytest
import torch
import torch.nn as nn

from sagellm_backend.attention import (
    AttentionMetadata,
    AttentionType,
    CPUAttentionBackend,
)
from sagellm_core.attention import (
    SageAttention,
    SageAttentionConfig,
    replace_attention_layers,
)


class TestSageAttentionConfig:
    """Test SageAttentionConfig."""

    def test_default_config(self) -> None:
        """Test default configuration."""
        config = SageAttentionConfig(num_heads=8)

        assert config.num_heads == 8
        assert config.num_kv_heads == 8  # Should default to num_heads
        assert config.head_dim == 64
        assert config.hidden_size == 8 * 64
        assert config.scale is not None

    def test_gqa_config(self) -> None:
        """Test GQA configuration."""
        config = SageAttentionConfig(
            num_heads=32,
            num_kv_heads=8,
            head_dim=128,
        )

        assert config.num_heads == 32
        assert config.num_kv_heads == 8
        assert config.head_dim == 128
        assert config.hidden_size == 32 * 128

    def test_custom_scale(self) -> None:
        """Test custom attention scale."""
        config = SageAttentionConfig(num_heads=8, scale=0.5)
        assert config.scale == 0.5


class TestSageAttention:
    """Test SageAttention layer."""

    @pytest.fixture
    def attention_backend(self) -> CPUAttentionBackend:
        """Create CPU attention backend for testing."""
        return CPUAttentionBackend()

    @pytest.fixture
    def attention_config(self) -> SageAttentionConfig:
        """Create test attention config."""
        return SageAttentionConfig(
            num_heads=4,
            head_dim=64,
            hidden_size=256,
        )

    @pytest.fixture
    def sage_attention(
        self,
        attention_config: SageAttentionConfig,
        attention_backend: CPUAttentionBackend,
    ) -> SageAttention:
        """Create SageAttention layer."""
        return SageAttention(attention_config, attention_backend)

    def test_initialization(
        self,
        sage_attention: SageAttention,
        attention_config: SageAttentionConfig,
    ) -> None:
        """Test layer initialization."""
        assert sage_attention.num_heads == attention_config.num_heads
        assert sage_attention.head_dim == attention_config.head_dim
        assert sage_attention.hidden_size == attention_config.hidden_size

    def test_forward_shape(
        self,
        sage_attention: SageAttention,
    ) -> None:
        """Test forward pass output shape."""
        batch_size = 2
        seq_len = 10
        hidden_size = 256

        # Create input
        hidden_states = torch.randn(batch_size, seq_len, hidden_size)

        # Create attention metadata
        attn_metadata = AttentionMetadata(
            seq_lens=[seq_len] * batch_size,
            is_prefill=True,
            attn_type=AttentionType.PREFILL,
            max_seq_len=seq_len,
            num_tokens=batch_size * seq_len,
        )

        # Forward pass
        output = sage_attention(hidden_states, attn_metadata)

        # Check output shape
        assert output.hidden_states.shape == (batch_size, seq_len, hidden_size)

    def test_forward_with_kv_cache(
        self,
        sage_attention: SageAttention,
    ) -> None:
        """Test forward pass with KV cache."""
        batch_size = 1
        seq_len = 5
        hidden_size = 256
        num_kv_heads = 4
        head_dim = 64

        # Create input
        hidden_states = torch.randn(batch_size, seq_len, hidden_size)

        # Create KV cache (for testing, just empty tensors)
        num_blocks = 2
        block_size = 16
        key_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_dim)
        value_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_dim)
        kv_cache = (key_cache, value_cache)

        # Create attention metadata
        attn_metadata = AttentionMetadata(
            seq_lens=[seq_len],
            is_prefill=True,
            attn_type=AttentionType.PREFILL,
            max_seq_len=seq_len,
            num_tokens=seq_len,
            block_size=block_size,
        )

        # Forward pass
        output = sage_attention(
            hidden_states,
            attn_metadata,
            kv_cache=kv_cache,
        )

        assert output.hidden_states.shape == (batch_size, seq_len, hidden_size)

    def test_gqa_attention(
        self,
        attention_backend: CPUAttentionBackend,
    ) -> None:
        """Test GQA (Grouped Query Attention)."""
        config = SageAttentionConfig(
            num_heads=8,
            num_kv_heads=2,  # GQA: fewer KV heads
            head_dim=64,
        )

        sage_attn = SageAttention(config, attention_backend)

        assert sage_attn.num_heads == 8
        assert sage_attn.num_kv_heads == 2
        assert sage_attn.num_queries_per_kv == 4

    def test_projection_shapes(
        self,
        sage_attention: SageAttention,
    ) -> None:
        """Test QKV projection layer shapes."""
        hidden_size = sage_attention.hidden_size
        num_heads = sage_attention.num_heads
        num_kv_heads = sage_attention.num_kv_heads
        head_dim = sage_attention.head_dim

        # Check Q projection
        assert sage_attention.q_proj.in_features == hidden_size
        assert sage_attention.q_proj.out_features == num_heads * head_dim

        # Check K projection
        assert sage_attention.k_proj.in_features == hidden_size
        assert sage_attention.k_proj.out_features == num_kv_heads * head_dim

        # Check V projection
        assert sage_attention.v_proj.in_features == hidden_size
        assert sage_attention.v_proj.out_features == num_kv_heads * head_dim

        # Check output projection
        assert sage_attention.o_proj.in_features == num_heads * head_dim
        assert sage_attention.o_proj.out_features == hidden_size


class TestLayerReplacement:
    """Test layer replacement functionality."""

    def test_replace_simple_model(self) -> None:
        """Test replacing attention in a simple model."""

        # Create a simple model with attention-like layer
        class SimpleAttention(nn.Module):
            def __init__(self, hidden_size: int = 256, num_heads: int = 4):
                super().__init__()
                self.hidden_size = hidden_size
                self.num_heads = num_heads
                self.q_proj = nn.Linear(hidden_size, hidden_size)
                self.k_proj = nn.Linear(hidden_size, hidden_size)
                self.v_proj = nn.Linear(hidden_size, hidden_size)
                self.o_proj = nn.Linear(hidden_size, hidden_size)

                # Add config for factory
                self.config = type(
                    "Config",
                    (),
                    {
                        "num_attention_heads": num_heads,
                        "hidden_size": hidden_size,
                    },
                )()

        class SimpleModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.attention = SimpleAttention()

        model = SimpleModel()
        backend = CPUAttentionBackend()

        # Create custom factory for SimpleAttention
        from sagellm_core.attention.layer_registry import (
            AttentionLayerFactory,
            AttentionRegistry,
        )

        class SimpleAttentionFactory(AttentionLayerFactory):
            def detect_attention_layer(self, module: nn.Module) -> bool:
                return isinstance(module, SimpleAttention)

            def extract_config(self, hf_attention: nn.Module) -> SageAttentionConfig:
                return SageAttentionConfig(
                    num_heads=hf_attention.num_heads,
                    hidden_size=hf_attention.hidden_size,
                )

            def create_sage_attention(
                self,
                hf_attention: nn.Module,
                attention_backend,
                kv_cache_manager=None,
            ):
                config = self.extract_config(hf_attention)
                sage_attn = SageAttention(config, attention_backend, kv_cache_manager)
                self.copy_weights(hf_attention, sage_attn)
                return sage_attn

            def copy_weights(self, hf_attention: nn.Module, sage_attention: SageAttention):
                sage_attention.q_proj.weight.data.copy_(hf_attention.q_proj.weight.data)
                sage_attention.k_proj.weight.data.copy_(hf_attention.k_proj.weight.data)
                sage_attention.v_proj.weight.data.copy_(hf_attention.v_proj.weight.data)
                sage_attention.o_proj.weight.data.copy_(hf_attention.o_proj.weight.data)

        # Register and replace
        registry = AttentionRegistry()
        registry.register("simple", SimpleAttentionFactory())

        model = replace_attention_layers(
            model,
            backend,
            model_type="simple",
            registry=registry,
        )

        # Check replacement
        assert isinstance(model.attention, SageAttention)
        assert model.attention.num_heads == 4
        assert model.attention.hidden_size == 256


@pytest.mark.integration
class TestIntegration:
    """Integration tests with real components."""

    @pytest.mark.skipif(
        not torch.cuda.is_available(),
        reason="CUDA not available",
    )
    def test_with_paged_attention_backend(self) -> None:
        """Test with PagedAttention backend (requires CUDA)."""
        from sagellm_backend.attention import PagedAttentionBackend

        backend = PagedAttentionBackend(block_size=16)
        config = SageAttentionConfig(num_heads=8)
        attention = SageAttention(config, backend)

        batch_size = 2
        seq_len = 10
        hidden_size = 512

        hidden_states = torch.randn(batch_size, seq_len, hidden_size, device="cuda")

        attn_metadata = AttentionMetadata(
            seq_lens=[seq_len] * batch_size,
            is_prefill=True,
            attn_type=AttentionType.PREFILL,
            max_seq_len=seq_len,
            num_tokens=batch_size * seq_len,
        )

        output = attention(hidden_states, attn_metadata)
        assert output.hidden_states.shape == (batch_size, seq_len, hidden_size)

    def test_with_kv_cache_manager(self) -> None:
        """Test with KVCacheManager."""
        from sagellm_kv_cache import KVCacheManager

        backend = CPUAttentionBackend()
        kv_manager = KVCacheManager(max_tokens=1024, block_size=16)

        config = SageAttentionConfig(num_heads=4)
        attention = SageAttention(config, backend, kv_manager)

        batch_size = 1
        seq_len = 8
        hidden_size = 256

        hidden_states = torch.randn(batch_size, seq_len, hidden_size)

        attn_metadata = AttentionMetadata(
            seq_lens=[seq_len],
            is_prefill=True,
            attn_type=AttentionType.PREFILL,
            max_seq_len=seq_len,
            num_tokens=seq_len,
        )

        output = attention(hidden_states, attn_metadata)

        # Should allocate KV cache
        assert output.kv_handle is not None
        assert output.kv_handle.num_tokens == seq_len
