"""Tests for ROCm/HIP Attention Backend.

Tests ROCm attention implementations for Hygon DCU and Muxi GPUs.
"""

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 TestROCmAttentionBackend:
    """Test ROCm attention backend."""

    def test_import(self):
        """Test that ROCm attention backend can be imported."""
        from sagellm_backend.attention.rocm import ROCmAttentionBackend

        assert ROCmAttentionBackend is not None

    def test_initialization(self):
        """Test ROCm attention backend initialization."""
        from sagellm_backend.attention.rocm import ROCmAttentionBackend

        backend = ROCmAttentionBackend(device="cuda:0", use_flash_attn=False)
        assert backend.name == "rocm_hip"
        assert not backend.supports_paged_attention

    def test_get_capabilities(self):
        """Test getting backend capabilities."""
        from sagellm_backend.attention.rocm import ROCmAttentionBackend

        backend = ROCmAttentionBackend(device="cuda:0")
        caps = backend.get_capabilities()

        assert "name" in caps
        assert "architecture" in caps
        assert "device" in caps
        assert "wavefront_size" in caps
        assert caps["wavefront_size"] == 64
        assert caps["supports_causal_mask"] is True

    def test_forward_basic(self):
        """Test basic forward pass."""
        from sagellm_backend.attention.rocm import ROCmAttentionBackend

        backend = ROCmAttentionBackend(device="cuda:0", use_flash_attn=False)

        # Create dummy tensors
        batch, num_heads, seq_len, head_dim = 2, 8, 16, 64
        query = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")
        key = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")
        value = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")

        # Mock attention metadata
        from sagellm_backend.attention.base import AttentionMetadata

        metadata = AttentionMetadata(
            num_prefills=batch,
            num_prefill_tokens=batch * seq_len,
            num_decode_tokens=0,
            seq_lens=torch.tensor([seq_len, seq_len], device="cuda:0"),
        )

        # Forward pass
        output = backend.forward(query, key, value, metadata)

        assert output.shape == (batch, seq_len, num_heads, head_dim)
        assert output.device.type == "cuda"

    def test_forward_causal(self):
        """Test forward pass with causal masking."""
        from sagellm_backend.attention.rocm import ROCmAttentionBackend

        backend = ROCmAttentionBackend(device="cuda:0", use_flash_attn=False)

        # Create dummy tensors
        batch, num_heads, seq_len, head_dim = 1, 4, 8, 32
        query = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")
        key = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")
        value = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")

        from sagellm_backend.attention.base import AttentionMetadata

        metadata = AttentionMetadata(
            num_prefills=batch,
            num_prefill_tokens=batch * seq_len,
            num_decode_tokens=0,
            seq_lens=torch.tensor([seq_len], device="cuda:0"),
        )

        # Forward pass with causal mask
        output = backend.forward(query, key, value, metadata, is_causal=True)

        assert output.shape == (batch, seq_len, num_heads, head_dim)

    @pytest.mark.skipif(True, reason="Flash attention may not be installed")
    def test_flash_attention(self):
        """Test flash attention forward pass."""
        from sagellm_backend.attention.rocm import ROCmAttentionBackend

        # Try to use flash attention
        backend = ROCmAttentionBackend(device="cuda:0", use_flash_attn=True)

        # If flash attention is not available, backend will use fallback
        batch, num_heads, seq_len, head_dim = 1, 4, 16, 64
        query = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")
        key = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")
        value = torch.randn(batch, seq_len, num_heads, head_dim, device="cuda:0")

        from sagellm_backend.attention.base import AttentionMetadata

        metadata = AttentionMetadata(
            num_prefills=batch,
            num_prefill_tokens=batch * seq_len,
            num_decode_tokens=0,
            seq_lens=torch.tensor([seq_len], device="cuda:0"),
        )

        output = backend.forward(query, key, value, metadata)
        assert output.shape == (batch, seq_len, num_heads, head_dim)


class TestROCmKernels:
    """Test ROCm kernel implementations."""

    def test_rocm_attention_kernel(self):
        """Test ROCm attention kernel."""
        from sagellm_backend.kernels.rocm.attention_kernels import rocm_attention

        batch, num_heads, seq_len, head_dim = 1, 4, 8, 32
        query = torch.randn(batch, num_heads, seq_len, head_dim, device="cuda:0")
        key = torch.randn(batch, num_heads, seq_len, head_dim, device="cuda:0")
        value = torch.randn(batch, num_heads, seq_len, head_dim, device="cuda:0")

        output = rocm_attention(query, key, value)
        assert output.shape == (batch, num_heads, seq_len, head_dim)

    def test_rocm_gqa_attention(self):
        """Test ROCm GQA attention."""
        from sagellm_backend.kernels.rocm.attention_kernels import rocm_gqa_attention

        batch, num_query_heads, num_kv_heads = 1, 8, 2
        seq_len, head_dim = 16, 64

        query = torch.randn(batch, num_query_heads, seq_len, head_dim, device="cuda:0")
        key = torch.randn(batch, num_kv_heads, seq_len, head_dim, device="cuda:0")
        value = torch.randn(batch, num_kv_heads, seq_len, head_dim, device="cuda:0")

        output = rocm_gqa_attention(query, key, value, num_kv_heads)
        assert output.shape == (batch, num_query_heads, seq_len, head_dim)

    def test_rocm_gemm(self):
        """Test ROCm GEMM kernel."""
        from sagellm_backend.kernels.rocm.gemm_kernels import rocm_gemm

        m, k, n = 16, 32, 64
        a = torch.randn(m, k, device="cuda:0")
        b = torch.randn(k, n, device="cuda:0")
        bias = torch.randn(n, device="cuda:0")

        output = rocm_gemm(a, b, bias)
        assert output.shape == (m, n)

    def test_rocm_layernorm(self):
        """Test ROCm LayerNorm kernel."""
        from sagellm_backend.kernels.rocm.normalization_kernels import rocm_layernorm

        batch, seq_len, hidden_size = 2, 16, 128
        input = torch.randn(batch, seq_len, hidden_size, device="cuda:0")
        weight = torch.randn(hidden_size, device="cuda:0")
        bias = torch.randn(hidden_size, device="cuda:0")

        output = rocm_layernorm(input, (hidden_size,), weight, bias)
        assert output.shape == (batch, seq_len, hidden_size)

    def test_rocm_rmsnorm(self):
        """Test ROCm RMSNorm kernel."""
        from sagellm_backend.kernels.rocm.normalization_kernels import rocm_rmsnorm

        batch, seq_len, hidden_size = 2, 16, 128
        input = torch.randn(batch, seq_len, hidden_size, device="cuda:0")
        weight = torch.randn(hidden_size, device="cuda:0")

        output = rocm_rmsnorm(input, weight)
        assert output.shape == (batch, seq_len, hidden_size)

    def test_hipblas_availability(self):
        """Test hipBLAS availability check."""
        from sagellm_backend.kernels.rocm.gemm_kernels import (
            check_hipblas_available,
            get_hipblas_version,
        )

        # Should be available since we're running on CUDA/ROCm
        assert check_hipblas_available() is True

        # Version may be "unknown" but should not crash
        version = get_hipblas_version()
        assert version is not None
