"""Tests for attention backend registry system.

Tests:
- register_attention_backend / get_attention_backend
- list_attention_backends returns registered backends
- KeyError for unregistered backends
- is_attention_backend_available checks
"""

import pytest
import torch

from sagellm_backend.attention import (
    AttentionBackend,
    AttentionMetadata,
    get_attention_backend,
    is_attention_backend_available,
    list_attention_backends,
    register_attention_backend,
)


class DummyAttentionBackend(AttentionBackend):
    """Dummy backend for testing registration."""

    @property
    def name(self) -> str:
        return "dummy"

    @property
    def supports_paged_attention(self) -> bool:
        return False

    @property
    def supports_prefix_caching(self) -> bool:
        return False

    @property
    def supports_sliding_window(self) -> bool:
        return False

    @property
    def supports_cross_attention(self) -> bool:
        return False

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        metadata: AttentionMetadata,
        scale: float | None = None,
    ) -> torch.Tensor:
        """Dummy forward implementation."""
        batch_size = query.shape[0]
        num_heads = query.shape[1]
        head_dim = query.shape[2]
        return torch.zeros(batch_size, num_heads, head_dim)

    @classmethod
    def is_available(cls) -> bool:
        return True


class UnavailableBackend(AttentionBackend):
    """Backend that is not available."""

    @property
    def name(self) -> str:
        return "unavailable"

    @property
    def supports_paged_attention(self) -> bool:
        return False

    @property
    def supports_prefix_caching(self) -> bool:
        return False

    @property
    def supports_sliding_window(self) -> bool:
        return False

    @property
    def supports_cross_attention(self) -> bool:
        return False

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        metadata: AttentionMetadata,
        scale: float | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

    @classmethod
    def is_available(cls) -> bool:
        return False


class TestAttentionRegistry:
    """Test suite for attention backend registry."""

    def test_list_attention_backends(self):
        """Test listing registered backends."""
        backends = list_attention_backends()
        assert isinstance(backends, list)
        assert len(backends) > 0
        # Default backends should be registered
        assert "cpu" in backends

    def test_get_default_backends(self):
        """Test getting default registered backends."""
        # CPU backend should always be available
        cpu_backend = get_attention_backend("cpu")
        assert cpu_backend.name == "cpu"

        # Flash backend should be registered (may not be available)
        backends = list_attention_backends()
        assert "flash" in backends

        # Paged backend should be registered
        assert "paged" in backends

    def test_register_and_get_backend(self):
        """Test registering and retrieving a custom backend."""
        # Register dummy backend class (not instance)
        register_attention_backend("dummy_test", DummyAttentionBackend)

        # Verify it's in the list
        backends = list_attention_backends()
        assert "dummy_test" in backends

        # Retrieve it (will be instantiated)
        retrieved = get_attention_backend("dummy_test")
        assert retrieved.name == "dummy"

    def test_get_unregistered_backend_raises_keyerror(self):
        """Test that getting an unregistered backend raises KeyError."""
        with pytest.raises(KeyError, match="not found"):
            get_attention_backend("nonexistent_backend_xyz")

    def test_is_attention_backend_available(self):
        """Test checking if a backend is available."""
        # CPU should be available
        assert is_attention_backend_available("cpu")

        # Unregistered backend should return False
        assert not is_attention_backend_available("nonexistent_backend_xyz_123")

    def test_register_overwrites_existing(self):
        """Test that re-registering a backend with overwrite=True works."""
        # Register first backend
        register_attention_backend("overwrite_test", DummyAttentionBackend)

        # Register second backend with same name (need overwrite=True)
        register_attention_backend("overwrite_test", DummyAttentionBackend, overwrite=True)

        # Should get a backend instance
        retrieved = get_attention_backend("overwrite_test")
        assert retrieved.name == "dummy"

    def test_backend_capabilities(self):
        """Test checking backend capabilities."""
        cpu_backend = get_attention_backend("cpu")

        # Check basic capabilities
        assert hasattr(cpu_backend, "supports_paged_attention")
        assert isinstance(cpu_backend.supports_paged_attention, bool)
        assert cpu_backend.supports_paged_attention is False  # CPU doesn't support paged

    def test_flash_backend_registration(self):
        """Test that flash backend is registered."""
        backends = list_attention_backends()
        assert "flash" in backends

        # Flash backend initialization requires flash-attn
        # Just check it's registered, don't instantiate
        assert "flash" in backends

    def test_paged_backend_registration(self):
        """Test that paged backend is registered."""
        backends = list_attention_backends()
        assert "paged" in backends

        # Get the backend
        paged_backend = get_attention_backend("paged")
        assert paged_backend.name == "paged"
        assert paged_backend.supports_paged_attention is True

    def test_backend_forward_signature(self):
        """Test that available backends have forward method."""
        for backend_name in ["cpu", "paged"]:
            if not is_attention_backend_available(backend_name):
                continue

            backend = get_attention_backend(backend_name)

            # Check forward method exists
            assert hasattr(backend, "forward")
            assert callable(backend.forward)

            # Check method signature
            import inspect

            sig = inspect.signature(backend.forward)
            params = list(sig.parameters.keys())
            assert "query" in params
            assert "key" in params
            assert "value" in params
            # attn_metadata is the actual parameter name
            assert "attn_metadata" in params


class TestBackendSelection:
    """Test suite for backend selection logic."""

    def test_get_best_available_backend(self):
        """Test automatic selection of best available backend."""
        # This tests the provider's logic for selecting backends
        # The order should be: flash > paged > cpu

        available_backends = [
            name for name in list_attention_backends() if is_attention_backend_available(name)
        ]

        assert len(available_backends) > 0
        # CPU should always be available
        assert "cpu" in available_backends

    def test_backend_priority(self):
        """Test that backends are selected in priority order."""
        # Get all available backends
        backends = list_attention_backends()

        # Check that priority backends are registered
        # (availability depends on environment)
        assert "cpu" in backends  # Always registered
        assert "flash" in backends  # Registered but may not be available
        assert "paged" in backends  # Registered and should be available

    def test_fallback_to_cpu(self):
        """Test that CPU backend is always available as fallback."""
        cpu_backend = get_attention_backend("cpu")

        # CPU backend should work on any device
        query = torch.randn(2, 4, 64)
        key = torch.randn(2, 4, 64)
        value = torch.randn(2, 4, 64)

        # Create minimal metadata for prefill
        metadata = AttentionMetadata.for_prefill(
            seq_lens=[2, 2],
            device="cpu",
        )

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


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