"""Tests for Kunlun XPU backend provider.

Tests Kunlun provider functionality including:
- Device initialization
- Stream and event creation
- KV block management
- Attention backend support
- Kernel registration
"""

from __future__ import annotations

import pytest

# Skip all tests if torch_xpu is not available
pytest.importorskip("torch")

try:
    import torch_xpu  # noqa: F401

    _has_xpu = hasattr(__import__("torch"), "xpu") and __import__("torch").xpu.is_available()
except ImportError:
    _has_xpu = False

pytestmark = pytest.mark.skipif(
    not _has_xpu, reason="Kunlun XPU (torch_xpu) not available. Tests require XPU hardware."
)

from sagellm_protocol import DType

from sagellm_backend.providers.kunlun import KunlunBackendProvider


class TestKunlunBackendProvider:
    """Test cases for KunlunBackendProvider."""

    def test_init_default_device(self) -> None:
        """Test initialization with default device."""
        provider = KunlunBackendProvider()
        assert provider is not None

        cap = provider.capability()
        assert cap.device_type == "xpu"
        assert "kunlun" in cap.compute_capability.lower()

    def test_init_specific_device(self) -> None:
        """Test initialization with specific device ID."""
        import torch

        if torch.xpu.device_count() > 0:
            provider = KunlunBackendProvider(device_id=0)
            assert provider is not None

    def test_capability_query(self) -> None:
        """Test capability descriptor query."""
        provider = KunlunBackendProvider()
        cap = provider.capability()

        # Check basic fields
        assert cap.device_name is not None
        assert cap.device_type == "xpu"
        assert cap.total_memory_bytes > 0

        # Check supported data types
        assert DType.FP32 in cap.supported_dtypes
        assert DType.FP16 in cap.supported_dtypes

        # Check features
        assert cap.has_stream is True
        assert cap.has_event is True

    def test_stream_creation(self) -> None:
        """Test stream creation."""
        provider = KunlunBackendProvider()
        stream = provider.create_stream()
        assert stream is not None

        # Test synchronization
        stream.synchronize()

    def test_event_creation(self) -> None:
        """Test event creation."""
        provider = KunlunBackendProvider()
        event = provider.create_event()
        assert event is not None

        # Test synchronization
        event.synchronize()

    def test_kv_block_alloc_free(self) -> None:
        """Test KV block allocation and deallocation."""
        provider = KunlunBackendProvider()

        # Allocate block
        block = provider.kv_block_alloc(num_tokens=128, dtype=DType.FP16)
        assert block.handle > 0
        assert block.num_tokens == 128
        assert block.dtype == DType.FP16

        # Free block
        provider.kv_block_free(block)

    def test_kv_block_copy(self) -> None:
        """Test KV block copy."""
        provider = KunlunBackendProvider()

        # Allocate two blocks
        block1 = provider.kv_block_alloc(num_tokens=64, dtype=DType.FP32)
        block2 = provider.kv_block_alloc(num_tokens=64, dtype=DType.FP32)

        # Copy block1 to block2
        provider.kv_block_copy(src=block1, dst=block2)

        # Cleanup
        provider.kv_block_free(block1)
        provider.kv_block_free(block2)

    def test_kv_block_migrate(self) -> None:
        """Test KV block migration."""
        provider = KunlunBackendProvider()

        # Allocate block on XPU
        block = provider.kv_block_alloc(num_tokens=32, dtype=DType.FP16)
        assert "xpu" in block.device

        # Migrate to CPU
        cpu_block = provider.kv_block_migrate(block, target_device="cpu")
        assert cpu_block.device == "cpu"
        assert cpu_block.num_tokens == 32

        # Cleanup
        provider.kv_block_free(block)
        provider.kv_block_free(cpu_block)

    def test_memory_stats(self) -> None:
        """Test memory statistics query."""
        provider = KunlunBackendProvider()
        stats = provider.memory_stats()

        # May return empty dict if API not available
        if stats:
            assert "total_bytes" in stats
            assert "free_bytes" in stats
            assert stats["total_bytes"] > 0

    def test_clear_cache(self) -> None:
        """Test cache clearing."""
        provider = KunlunBackendProvider()
        provider.clear_cache()  # Should not raise

    def test_kernel_registration(self) -> None:
        """Test kernel registration and retrieval."""
        provider = KunlunBackendProvider()

        # Check default kernels are registered
        supported_backends = provider.get_supported_attention_backends()
        assert "kunlun_xpu" in supported_backends

        # Test getting a kernel
        try:
            matmul_kernel = provider.get_kernel("kunlun_matmul")
            assert matmul_kernel is not None
        except ValueError:
            # Kernel might not be registered if kunlun module import failed
            pytest.skip("kunlun_matmul kernel not registered")

    def test_attention_backend_support(self) -> None:
        """Test attention backend support."""
        provider = KunlunBackendProvider()

        # Check supported backends
        backends = provider.get_supported_attention_backends()
        assert isinstance(backends, list)
        assert "kunlun_xpu" in backends

        # Get attention backend
        attn_backend = provider.get_attention_backend()
        assert attn_backend is not None
        assert attn_backend.name == "kunlun_xpu"

    def test_torch_device_access(self) -> None:
        """Test torch device access."""
        import torch

        provider = KunlunBackendProvider()
        device = provider.get_torch_device()
        assert isinstance(device, torch.device)
        assert device.type == "xpu"

    def test_to_device(self) -> None:
        """Test tensor device transfer."""
        import torch

        provider = KunlunBackendProvider()

        # Create CPU tensor
        cpu_tensor = torch.randn(4, 4)

        # Move to XPU
        xpu_tensor = provider.to_device(cpu_tensor)
        assert xpu_tensor.device.type == "xpu"

    def test_dtype_validation(self) -> None:
        """Test dtype support validation."""
        provider = KunlunBackendProvider()

        # Supported dtype should work
        block = provider.kv_block_alloc(num_tokens=16, dtype=DType.FP16)
        assert block is not None
        provider.kv_block_free(block)

        # Unsupported dtype should raise
        # (Note: Currently FP32, FP16, INT8 are supported)
        # This test validates the fail-fast behavior
