"""Tests for Kunlun XPU attention backend.

Tests Kunlun attention implementation including:
- Attention forward pass
- PyTorch fallback
- XDNN integration (mocked)
"""

from __future__ import annotations

import pytest
import torch

# Skip all tests if torch_xpu is not available
try:
    import torch_xpu  # noqa: F401
    _has_xpu = hasattr(torch, "xpu") and 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_backend.attention.base import AttentionMetadata
from sagellm_backend.attention.kunlun import KunlunAttentionBackend


class TestKunlunAttentionBackend:
    """Test cases for KunlunAttentionBackend."""

    def test_init(self) -> None:
        """Test backend initialization."""
        backend = KunlunAttentionBackend(device="xpu:0")
        assert backend.name == "kunlun_xpu"
        assert backend.supports_paged_attention is False

    def test_forward_simple(self) -> None:
        """Test simple attention forward pass."""
        backend = KunlunAttentionBackend(device="xpu:0")

        # Create simple inputs
        batch_size = 2
        seq_len = 4
        num_heads = 2
        head_dim = 8

        query = torch.randn(batch_size, num_heads, seq_len, head_dim, device="xpu:0")
        key = torch.randn(batch_size, num_heads, seq_len, head_dim, device="xpu:0")
        value = torch.randn(batch_size, num_heads, seq_len, head_dim, device="xpu:0")

        # Create metadata
        metadata = AttentionMetadata.for_prefill(
            seq_lens=[seq_len] * batch_size,
            device="xpu:0"
        )

        # Execute attention
        output = backend.forward(query, key, value, metadata)

        assert output.shape == query.shape
        assert output.device.type == "xpu"

    def test_forward_3d_input(self) -> None:
        """Test attention with 3D input (num_tokens, num_heads, head_dim)."""
        backend = KunlunAttentionBackend(device="xpu:0")

        num_tokens = 8
        num_heads = 4
        head_dim = 16

        query = torch.randn(num_tokens, num_heads, head_dim, device="xpu:0")
        key = torch.randn(num_tokens, num_heads, head_dim, device="xpu:0")
        value = torch.randn(num_tokens, num_heads, head_dim, device="xpu:0")

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[num_tokens],
            device="xpu:0"
        )

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

        # Output should have same shape as query
        assert output.shape == (num_tokens, num_heads, head_dim)
        assert output.device.type == "xpu"

    def test_forward_with_scale(self) -> None:
        """Test attention with custom scale factor."""
        backend = KunlunAttentionBackend(device="xpu:0")

        query = torch.randn(1, 2, 4, 8, device="xpu:0")
        key = torch.randn(1, 2, 4, 8, device="xpu:0")
        value = torch.randn(1, 2, 4, 8, device="xpu:0")

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[4],
            device="xpu:0"
        )

        output = backend.forward(query, key, value, metadata, scale=0.5)

        assert output.shape == query.shape

    def test_forward_cpu_to_xpu(self) -> None:
        """Test automatic device transfer from CPU to XPU."""
        backend = KunlunAttentionBackend(device="xpu:0")

        # Create inputs on CPU
        query = torch.randn(1, 2, 4, 8)
        key = torch.randn(1, 2, 4, 8)
        value = torch.randn(1, 2, 4, 8)

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[4],
            device="xpu:0"
        )

        # Should automatically transfer to XPU
        output = backend.forward(query, key, value, metadata)

        assert output.device.type == "xpu"

    def test_xdnn_fallback(self) -> None:
        """Test XDNN fallback to PyTorch implementation."""
        backend = KunlunAttentionBackend(device="xpu:0")

        # Even if XDNN is not available, should fall back gracefully
        query = torch.randn(1, 2, 4, 8, device="xpu:0")
        key = torch.randn(1, 2, 4, 8, device="xpu:0")
        value = torch.randn(1, 2, 4, 8, device="xpu:0")

        metadata = AttentionMetadata.for_prefill(
            seq_lens=[4],
            device="xpu:0"
        )

        output = backend.forward(query, key, value, metadata)
        assert output is not None


@pytest.mark.skipif(not _has_xpu, reason="XPU not available")
def test_attention_numerical_correctness() -> None:
    """Test that Kunlun attention produces reasonable outputs."""
    backend = KunlunAttentionBackend(device="xpu:0")

    # Simple test case
    batch_size = 1
    seq_len = 3
    num_heads = 1
    head_dim = 4

    # Create simple inputs with known patterns
    query = torch.ones(batch_size, num_heads, seq_len, head_dim, device="xpu:0")
    key = torch.ones(batch_size, num_heads, seq_len, head_dim, device="xpu:0")
    value = torch.arange(seq_len * head_dim, dtype=torch.float32, device="xpu:0").reshape(
        batch_size, num_heads, seq_len, head_dim
    )

    metadata = AttentionMetadata.for_prefill(
        seq_lens=[seq_len] * batch_size,
        device="xpu:0"
    )

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

    # Output should be finite
    assert torch.isfinite(output).all()

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