"""Tests for DCU Backend Provider with ROCm/HIP support.

Tests the enhanced DCU provider with attention backends.
"""

from __future__ import annotations

import pytest
import torch

# Skip all tests if ROCm/CUDA not available
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="ROCm/CUDA not available")


class TestDCUProvider:
    """Test DCU provider with ROCm support."""

    def test_import(self):
        """Test that DCU provider can be imported."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        assert DCUBackendProvider is not None

    def test_initialization(self):
        """Test DCU provider initialization."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)
        assert provider is not None

    def test_capability(self):
        """Test capability descriptor."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)
        cap = provider.capability()

        assert cap.device_type == "dcu"
        assert cap.has_stream is True
        assert cap.has_event is True
        assert "MATMUL" in [k.name for k in cap.supported_kernels]
        assert "ATTENTION" in [k.name for k in cap.supported_kernels]

    def test_get_attention_backend(self):
        """Test getting attention backend."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)

        # Default backend
        backend = provider.get_attention_backend()
        assert backend is not None
        assert backend.name == "rocm_hip"

        # Explicit ROCm backend
        backend_rocm = provider.get_attention_backend("rocm")
        assert backend_rocm is not None

        # PyTorch fallback
        backend_pytorch = provider.get_attention_backend("pytorch")
        assert backend_pytorch is not None

    def test_get_supported_attention_backends(self):
        """Test listing supported attention backends."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)
        backends = provider.get_supported_attention_backends()

        assert "rocm" in backends
        assert "flash" in backends
        assert "pytorch" in backends

    def test_invalid_attention_backend(self):
        """Test error handling for invalid backend."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)

        with pytest.raises(ValueError, match="Unsupported attention backend"):
            provider.get_attention_backend("invalid_backend")

    def test_kv_block_operations(self):
        """Test KV block allocation and operations."""
        from sagellm_protocol import DType

        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)

        # Allocate KV block
        block = provider.kv_block_alloc(num_tokens=128, dtype=DType.FP16)
        assert block.num_tokens == 128
        assert block.dtype == DType.FP16
        assert "dcu" in block.device

        # Free KV block
        provider.kv_block_free(block)

    def test_memory_stats(self):
        """Test memory statistics."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)
        stats = provider.memory_stats()

        assert "total_bytes" in stats
        assert "free_bytes" in stats
        assert "used_bytes" in stats
        assert stats["total_bytes"] > 0

    def test_stream_and_event(self):
        """Test stream and event creation."""
        from sagellm_backend.providers.dcu import DCUBackendProvider

        provider = DCUBackendProvider(device_id=0)

        # Create stream
        stream = provider.create_stream()
        assert stream is not None

        # Create event
        event = provider.create_event()
        assert event is not None

        # Synchronize
        stream.synchronize()
        event.synchronize()
