"""Tests for attention backend implementations.

Tests:
- CPUAttentionBackend.forward() output shape correctness
- AttentionMetadata.for_prefill() / for_decode() construction
- FlashAttentionBackend optional import (skip without flash-attn)
- PagedAttentionBackend block_tables handling
"""

import pytest
import torch

from sagellm_backend.attention import (
    AttentionMetadata,
    AttentionType,
    get_attention_backend,
    is_attention_backend_available,
)


class TestCPUAttentionBackend:
    """Test suite for CPU attention backend."""

    @pytest.fixture
    def cpu_backend(self):
        """Get CPU attention backend."""
        return get_attention_backend("cpu")

    def test_backend_basic(self, cpu_backend):
        """Test that CPU backend exists and has correct name."""
        assert cpu_backend.name == "cpu"
        assert cpu_backend.supports_paged_attention is False

    def test_prefill_forward_shape(self, cpu_backend):
        """Test that prefill forward produces correct output shape."""
        total_tokens = 128
        num_heads = 8
        head_dim = 64

        query = torch.randn(total_tokens, num_heads, head_dim)
        key = torch.randn(total_tokens, num_heads, head_dim)
        value = torch.randn(total_tokens, num_heads, head_dim)

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[32, 48, 48],
            device="cpu",
        )

        output = cpu_backend.forward(query, key, value, metadata)

        # Output should match query shape
        assert output.shape == query.shape
        assert output.dtype == query.dtype

    def test_custom_scale(self, cpu_backend):
        """Test attention with custom scale factor."""
        batch_size = 16
        num_heads = 4
        head_dim = 32

        query = torch.randn(batch_size, num_heads, head_dim)
        key = torch.randn(batch_size, num_heads, head_dim)
        value = torch.randn(batch_size, num_heads, head_dim)

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[8, 8],
            device="cpu",
        )

        # Forward with custom scale
        output = cpu_backend.forward(query, key, value, metadata, scale=0.5)
        assert output.shape == query.shape

    def test_different_shapes(self, cpu_backend):
        """Test attention with different tensor shapes."""
        # Small shape
        query = torch.randn(16, 4, 32)
        key = torch.randn(16, 4, 32)
        value = torch.randn(16, 4, 32)

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[8, 8],
            device="cpu",
        )

        output = cpu_backend.forward(query, key, value, metadata)
        assert output.shape == query.shape

        # Larger shape
        query2 = torch.randn(256, 16, 64)
        key2 = torch.randn(256, 16, 64)
        value2 = torch.randn(256, 16, 64)

        metadata2 = AttentionMetadata.for_prefill(
            seq_lens=[128, 128],
            device="cpu",
        )

        output2 = cpu_backend.forward(query2, key2, value2, metadata2)
        assert output2.shape == query2.shape


class TestAttentionMetadata:
    """Test suite for AttentionMetadata construction."""

    def test_for_prefill_basic(self):
        """Test basic prefill metadata construction."""
        metadata = AttentionMetadata.for_prefill(
            seq_lens=[16, 32, 48],
            device="cpu",
        )

        assert metadata.attn_type == AttentionType.PREFILL
        assert metadata.num_tokens == 16 + 32 + 48
        assert metadata.max_seq_len == 48
        assert len(metadata.seq_lens) == 3
        assert metadata.is_prefill is True

    def test_for_prefill_single_seq(self):
        """Test prefill metadata with single sequence."""
        metadata = AttentionMetadata.for_prefill(
            seq_lens=[128],
            device="cpu",
        )

        assert metadata.num_tokens == 128
        assert metadata.max_seq_len == 128
        assert len(metadata.seq_lens) == 1

    def test_for_decode_basic(self):
        """Test basic decode metadata construction."""
        block_tables = torch.randint(0, 100, (3, 32), dtype=torch.long)

        metadata = AttentionMetadata.for_decode(
            context_lens=[128, 256, 512],
            block_tables=block_tables,
            device="cpu",
        )

        assert metadata.attn_type == AttentionType.DECODE
        assert len(metadata.context_lens) == 3
        assert metadata.max_seq_len == 512
        assert metadata.is_prefill is False
        assert metadata.block_tables is not None

    def test_for_decode_with_block_tables(self):
        """Test decode metadata with block tables for paging."""
        block_tables = torch.randint(0, 100, (3, 32), dtype=torch.long)

        metadata = AttentionMetadata.for_decode(
            context_lens=[128, 256, 384],
            block_tables=block_tables,
            device="cpu",
        )

        assert metadata.block_tables is not None
        assert metadata.block_tables.shape == (3, 32)
        assert metadata.num_decode_tokens == 3

    def test_metadata_device_placement(self):
        """Test that metadata tensors are on correct device."""
        block_tables_cpu = torch.randint(0, 100, (2, 16), dtype=torch.long)

        metadata = AttentionMetadata.for_decode(
            context_lens=[64, 128],
            block_tables=block_tables_cpu,
            device="cpu",
        )

        if metadata.seq_start_loc is not None:
            assert metadata.seq_start_loc.device.type == "cpu"

        if torch.cuda.is_available():
            block_tables_cuda = torch.randint(0, 100, (2, 16), dtype=torch.long, device="cuda")
            metadata_cuda = AttentionMetadata.for_decode(
                context_lens=[64, 128],
                block_tables=block_tables_cuda,
                device="cuda",
            )
            if metadata_cuda.seq_start_loc is not None:
                assert metadata_cuda.seq_start_loc.device.type == "cuda"

    def test_seq_start_loc_computation(self):
        """Test that seq_start_loc is computed correctly for prefill."""
        seq_lens = [10, 20, 30]
        metadata = AttentionMetadata.for_prefill(
            seq_lens=seq_lens,
            device="cpu",
        )

        if metadata.seq_start_loc is not None:
            expected = torch.tensor([0, 10, 30, 60], dtype=torch.long)
            assert torch.all(metadata.seq_start_loc == expected)


class TestFlashAttentionBackend:
    """Test suite for Flash Attention backend (optional)."""

    @pytest.fixture
    def flash_backend(self):
        """Get Flash attention backend."""
        # Check if flash-attn is actually available (not just registered)
        try:
            import flash_attn  # noqa: F401
        except ImportError:
            pytest.skip("flash-attn package not installed")
        return get_attention_backend("flash")

    def test_backend_basic(self, flash_backend):
        """Test that Flash backend has correct properties."""
        assert flash_backend.name == "flash"

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
    def test_prefill_forward(self, flash_backend):
        """Test Flash attention prefill forward."""
        total_tokens = 128
        num_heads = 8
        head_dim = 64

        query = torch.randn(total_tokens, num_heads, head_dim, device="cuda")
        key = torch.randn(total_tokens, num_heads, head_dim, device="cuda")
        value = torch.randn(total_tokens, num_heads, head_dim, device="cuda")

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[32, 48, 48],
            device="cuda",
        )

        output = flash_backend.forward(query, key, value, metadata)
        assert output.shape == query.shape


class TestPagedAttentionBackend:
    """Test suite for Paged Attention backend."""

    @pytest.fixture
    def paged_backend(self):
        """Get Paged attention backend."""
        return get_attention_backend("paged")

    def test_backend_basic(self, paged_backend):
        """Test that Paged backend has correct properties."""
        assert paged_backend.name == "paged"
        assert paged_backend.supports_paged_attention is True

    def test_prefill_forward(self, paged_backend):
        """Test paged backend prefill (doesn't require KV cache)."""
        total_tokens = 64
        num_heads = 8
        head_dim = 64

        query = torch.randn(total_tokens, num_heads, head_dim)
        key = torch.randn(total_tokens, num_heads, head_dim)
        value = torch.randn(total_tokens, num_heads, head_dim)

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[32, 32],
            device="cpu",
        )

        output = paged_backend.forward(query, key, value, metadata)
        assert output.shape == query.shape

    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
    def test_cuda_prefill(self, paged_backend):
        """Test paged attention prefill on CUDA device."""
        total_tokens = 128
        num_heads = 8
        head_dim = 64

        query = torch.randn(total_tokens, num_heads, head_dim, device="cuda")
        key = torch.randn(total_tokens, num_heads, head_dim, device="cuda")
        value = torch.randn(total_tokens, num_heads, head_dim, device="cuda")

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[64, 64],
            device="cuda",
        )

        output = paged_backend.forward(query, key, value, metadata)
        assert output.shape == query.shape
        assert output.device.type == "cuda"


class TestBackendComparison:
    """Test suite comparing different backends."""

    @pytest.fixture
    def test_inputs(self):
        """Common test inputs for all backends."""
        total_tokens = 64
        num_heads = 4
        head_dim = 32

        return {
            "query": torch.randn(total_tokens, num_heads, head_dim),
            "key": torch.randn(total_tokens, num_heads, head_dim),
            "value": torch.randn(total_tokens, num_heads, head_dim),
            "attn_metadata": AttentionMetadata.for_prefill(
                seq_lens=[32, 32],
                device="cpu",
            ),
        }

    def test_all_backends_same_shape(self, test_inputs):
        """Test that all available backends produce same output shape."""
        outputs = {}

        for backend_name in ["cpu", "paged"]:
            if is_attention_backend_available(backend_name):
                backend = get_attention_backend(backend_name)
                output = backend.forward(**test_inputs)
                outputs[backend_name] = output

        # All outputs should have the same shape
        shapes = [out.shape for out in outputs.values()]
        assert len(set(shapes)) == 1  # All shapes are identical

    def test_cpu_paged_consistency(self, test_inputs):
        """Test that CPU and Paged backends produce similar results."""
        cpu_backend = get_attention_backend("cpu")
        paged_backend = get_attention_backend("paged")

        cpu_output = cpu_backend.forward(**test_inputs)
        paged_output = paged_backend.forward(**test_inputs)

        # Shapes should match
        assert cpu_output.shape == paged_output.shape

        # Both should produce valid outputs (not all NaN)
        # Note: Values may differ due to different implementations
        # Just ensure both produce reasonable outputs
        assert cpu_output.dtype == paged_output.dtype


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
