"""
Unit tests for the fused Triton kernel in normal_decode_set_metadata.

This test suite verifies:
1. Correctness against reference PyTorch implementation
2. Different page sizes (1, 16, 64)
3. With and without Sliding Window Attention (SWA)
4. Various batch sizes and sequence lengths
5. Edge cases
"""

import unittest

import torch

from sglang.srt.layers.attention.flashattention_backend import (
    normal_decode_set_metadata,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase

# Register this test for CUDA CI in stage-b (fast attention/kernel tests)
register_cuda_ci(est_time=25, suite="stage-b-test-1-gpu-large")


def reference_normal_decode_set_metadata(
    cache_seqlens_int32: torch.Tensor,
    cu_seqlens_k: torch.Tensor,
    page_table: torch.Tensor,
    req_to_token: torch.Tensor,
    req_pool_indices: torch.Tensor,
    strided_indices: torch.Tensor,
    max_seq_pages: int,
    seq_lens: torch.Tensor,
    seq_len_delta: int,
    page_size: int,
    swa_page_table: torch.Tensor = None,
    token_to_kv_pool=None,
):
    """
    Reference implementation using original PyTorch operations.
    This is the pre-Triton version for correctness comparison.
    """
    cache_seqlens_int32.copy_(seq_lens + seq_len_delta)
    cu_seqlens_k[1:].copy_(torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32))
    page_indices = req_to_token[
        req_pool_indices[:, None],
        strided_indices[:max_seq_pages][None, :],
    ]
    page_table[:, :max_seq_pages].copy_(page_indices // page_size)

    if swa_page_table is not None and token_to_kv_pool is not None:
        swa_page_indices = token_to_kv_pool.translate_loc_from_full_to_swa(page_indices)
        swa_page_table[:, :max_seq_pages].copy_(swa_page_indices // page_size)


@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestNormalDecodeSetMetadata(CustomTestCase):
    """Test fused Triton kernel in normal_decode_set_metadata."""

    def setUp(self):
        self.device = "cuda"
        self.dtype = torch.int32

    def _create_test_data(
        self,
        batch_size: int,
        max_seq_len: int,
        page_size: int,
        has_swa: bool = False,
        seq_len_delta: int = 0,
    ):
        """Create test data for normal_decode_set_metadata."""
        # Random sequence lengths for each batch
        seq_lens = torch.randint(
            max_seq_len // 2,
            max_seq_len + 1,
            (batch_size,),
            dtype=torch.int64,
            device=self.device,
        )

        # Calculate max_seq_pages
        max_len = seq_lens.max().item()
        max_seq_pages = (max_len + seq_len_delta + page_size - 1) // page_size

        # Create req_pool_indices (maps batch index to pool index)
        req_pool_indices = torch.arange(
            batch_size, dtype=torch.int32, device=self.device
        )

        # Create strided_indices for page table indexing
        if page_size == 1:
            strided_indices = torch.arange(
                max_seq_len * 2, dtype=torch.int32, device=self.device
            )
        else:
            strided_indices = torch.arange(
                0, max_seq_len * 2, page_size, dtype=torch.int32, device=self.device
            )

        # Create req_to_token pool (simulates token locations in KV cache)
        pool_size = batch_size
        max_tokens = max_seq_len * 2
        req_to_token = torch.randint(
            0, 10000, (pool_size, max_tokens), dtype=torch.int32, device=self.device
        )

        # Output tensors (to be filled by the function)
        cache_seqlens_int32 = torch.zeros(
            batch_size, dtype=torch.int32, device=self.device
        )
        cu_seqlens_k = torch.zeros(
            batch_size + 1, dtype=torch.int32, device=self.device
        )
        page_table = torch.zeros(
            (batch_size, max_seq_pages + 10), dtype=torch.int32, device=self.device
        )

        # SWA setup if needed
        swa_page_table = None
        token_to_kv_pool = None
        if has_swa:
            swa_page_table = torch.zeros(
                (batch_size, max_seq_pages + 10), dtype=torch.int32, device=self.device
            )
            # Create a simple SWA KV pool for testing
            token_to_kv_pool = self._create_swa_kv_pool(10000, page_size)

        return {
            "cache_seqlens_int32": cache_seqlens_int32,
            "cu_seqlens_k": cu_seqlens_k,
            "page_table": page_table,
            "req_to_token": req_to_token,
            "req_pool_indices": req_pool_indices,
            "strided_indices": strided_indices,
            "max_seq_pages": max_seq_pages,
            "seq_lens": seq_lens,
            "seq_len_delta": seq_len_delta,
            "page_size": page_size,
            "swa_page_table": swa_page_table,
            "token_to_kv_pool": token_to_kv_pool,
        }

    def _create_swa_kv_pool(self, size: int, page_size: int):
        """Create a mock SWA KV pool for testing that inherits from SWAKVPool."""

        # Create a minimal mock that inherits from SWAKVPool to pass isinstance check
        class MinimalSWAKVPool(SWAKVPool):
            def __init__(self, size, device):
                # Don't call super().__init__() to avoid complex initialization
                # Just set the minimal attributes needed for the test
                self.full_to_swa_index_mapping = torch.arange(
                    size, dtype=torch.int32, device=device
                )
                # Add some randomness to simulate real SWA mapping
                self.full_to_swa_index_mapping = (
                    self.full_to_swa_index_mapping
                    + torch.randint(0, 100, (size,), device=device)
                ) % size
                self.device = device

            def translate_loc_from_full_to_swa(self, page_indices):
                """Mock translation method."""
                return self.full_to_swa_index_mapping[page_indices]

        return MinimalSWAKVPool(size, self.device)

    def _run_test(
        self,
        batch_size: int,
        max_seq_len: int,
        page_size: int,
        has_swa: bool = False,
        seq_len_delta: int = 0,
    ):
        """Run a single test configuration."""
        # Create test data
        test_data = self._create_test_data(
            batch_size, max_seq_len, page_size, has_swa, seq_len_delta
        )

        # Clone data for reference implementation
        ref_data = {
            "cache_seqlens_int32": test_data["cache_seqlens_int32"].clone(),
            "cu_seqlens_k": test_data["cu_seqlens_k"].clone(),
            "page_table": test_data["page_table"].clone(),
            "swa_page_table": test_data["swa_page_table"].clone() if has_swa else None,
        }

        # Run reference implementation
        reference_normal_decode_set_metadata(
            ref_data["cache_seqlens_int32"],
            ref_data["cu_seqlens_k"],
            ref_data["page_table"],
            test_data["req_to_token"],
            test_data["req_pool_indices"],
            test_data["strided_indices"],
            test_data["max_seq_pages"],
            test_data["seq_lens"],
            test_data["seq_len_delta"],
            test_data["page_size"],
            ref_data["swa_page_table"],
            test_data["token_to_kv_pool"],
        )

        # Run fused Triton implementation
        normal_decode_set_metadata(
            test_data["cache_seqlens_int32"],
            test_data["cu_seqlens_k"],
            test_data["page_table"],
            test_data["req_to_token"],
            test_data["req_pool_indices"],
            test_data["strided_indices"],
            test_data["max_seq_pages"],
            test_data["seq_lens"],
            test_data["seq_len_delta"],
            test_data["page_size"],
            test_data["swa_page_table"],
            test_data["token_to_kv_pool"],
        )

        # Compare results
        self.assertTrue(
            torch.equal(
                test_data["cache_seqlens_int32"], ref_data["cache_seqlens_int32"]
            ),
            f"cache_seqlens_int32 mismatch. Expected:\n{ref_data['cache_seqlens_int32']}\nGot:\n{test_data['cache_seqlens_int32']}",
        )

        self.assertTrue(
            torch.equal(test_data["cu_seqlens_k"], ref_data["cu_seqlens_k"]),
            f"cu_seqlens_k mismatch. Expected:\n{ref_data['cu_seqlens_k']}\nGot:\n{test_data['cu_seqlens_k']}",
        )

        self.assertTrue(
            torch.equal(test_data["page_table"], ref_data["page_table"]),
            f"page_table mismatch at bs={batch_size}, page_size={page_size}",
        )

        if has_swa:
            self.assertTrue(
                torch.equal(test_data["swa_page_table"], ref_data["swa_page_table"]),
                f"swa_page_table mismatch at bs={batch_size}, page_size={page_size}",
            )

    # Test cases for page_size=1 (uses specialized kernel _fused_metadata_kernel_ps1_no_swa)
    def test_page_size_1_small_batch(self):
        """Test with page_size=1, small batch."""
        self._run_test(batch_size=2, max_seq_len=128, page_size=1, has_swa=False)

    def test_page_size_1_medium_batch(self):
        """Test with page_size=1, medium batch."""
        self._run_test(batch_size=16, max_seq_len=256, page_size=1, has_swa=False)

    def test_page_size_1_large_batch(self):
        """Test with page_size=1, large batch."""
        self._run_test(batch_size=64, max_seq_len=512, page_size=1, has_swa=False)

    def test_page_size_1_with_seq_len_delta(self):
        """Test with page_size=1 and seq_len_delta > 0."""
        self._run_test(
            batch_size=8, max_seq_len=200, page_size=1, has_swa=False, seq_len_delta=5
        )

    # Test cases for page_size > 1 (uses general kernel _fused_metadata_kernel_general)
    def test_page_size_16_small_batch(self):
        """Test with page_size=16, small batch."""
        self._run_test(batch_size=4, max_seq_len=256, page_size=16, has_swa=False)

    def test_page_size_16_medium_batch(self):
        """Test with page_size=16, medium batch."""
        self._run_test(batch_size=16, max_seq_len=512, page_size=16, has_swa=False)

    def test_page_size_64_small_batch(self):
        """Test with page_size=64, small batch."""
        self._run_test(batch_size=4, max_seq_len=512, page_size=64, has_swa=False)

    def test_page_size_64_medium_batch(self):
        """Test with page_size=64, medium batch."""
        self._run_test(batch_size=32, max_seq_len=1024, page_size=64, has_swa=False)

    def test_page_size_64_with_seq_len_delta(self):
        """Test with page_size=64 and seq_len_delta > 0."""
        self._run_test(
            batch_size=8, max_seq_len=512, page_size=64, has_swa=False, seq_len_delta=3
        )

    # Test cases with Sliding Window Attention (SWA)
    def test_page_size_16_with_swa(self):
        """Test with page_size=16 and SWA enabled."""
        self._run_test(batch_size=8, max_seq_len=256, page_size=16, has_swa=True)

    def test_page_size_64_with_swa(self):
        """Test with page_size=64 and SWA enabled."""
        self._run_test(batch_size=16, max_seq_len=512, page_size=64, has_swa=True)

    def test_page_size_64_with_swa_and_delta(self):
        """Test with page_size=64, SWA, and seq_len_delta."""
        self._run_test(
            batch_size=8, max_seq_len=400, page_size=64, has_swa=True, seq_len_delta=2
        )

    # Edge cases
    def test_batch_size_1(self):
        """Test with single batch."""
        self._run_test(batch_size=1, max_seq_len=128, page_size=1, has_swa=False)
        self._run_test(batch_size=1, max_seq_len=256, page_size=64, has_swa=False)

    def test_max_seq_pages_small(self):
        """Test edge case where max_seq_pages could be very small."""
        # This tests when sequences are very short
        test_data = self._create_test_data(
            batch_size=2, max_seq_len=10, page_size=64, has_swa=False
        )

        # Run fused implementation (should handle gracefully)
        normal_decode_set_metadata(
            test_data["cache_seqlens_int32"],
            test_data["cu_seqlens_k"],
            test_data["page_table"],
            test_data["req_to_token"],
            test_data["req_pool_indices"],
            test_data["strided_indices"],
            test_data["max_seq_pages"],
            test_data["seq_lens"],
            test_data["seq_len_delta"],
            test_data["page_size"],
            test_data["swa_page_table"],
            test_data["token_to_kv_pool"],
        )

        # Verify no crashes and basic properties
        self.assertEqual(
            test_data["cache_seqlens_int32"].sum().item(),
            test_data["seq_lens"].sum().item(),
        )

    def test_power_of_two_page_sizes(self):
        """Test various power-of-2 page sizes."""
        page_sizes = [1, 2, 4, 8, 16, 32, 64, 128]
        for page_size in page_sizes:
            with self.subTest(page_size=page_size):
                self._run_test(
                    batch_size=4, max_seq_len=256, page_size=page_size, has_swa=False
                )

    def test_varied_sequence_lengths(self):
        """Test with highly varied sequence lengths in the same batch."""
        batch_size = 8
        max_seq_len = 512
        page_size = 64

        test_data = self._create_test_data(
            batch_size, max_seq_len, page_size, has_swa=False
        )

        # Manually set varied sequence lengths
        test_data["seq_lens"] = torch.tensor(
            [10, 50, 100, 200, 300, 450, 500, 512],
            dtype=torch.int64,
            device=self.device,
        )
        test_data["max_seq_pages"] = (
            test_data["seq_lens"].max().item() + page_size - 1
        ) // page_size

        # Run both implementations
        ref_data = {
            "cache_seqlens_int32": test_data["cache_seqlens_int32"].clone(),
            "cu_seqlens_k": test_data["cu_seqlens_k"].clone(),
            "page_table": test_data["page_table"].clone(),
        }

        reference_normal_decode_set_metadata(
            ref_data["cache_seqlens_int32"],
            ref_data["cu_seqlens_k"],
            ref_data["page_table"],
            test_data["req_to_token"],
            test_data["req_pool_indices"],
            test_data["strided_indices"],
            test_data["max_seq_pages"],
            test_data["seq_lens"],
            0,
            page_size,
            None,
            None,
        )

        normal_decode_set_metadata(
            test_data["cache_seqlens_int32"],
            test_data["cu_seqlens_k"],
            test_data["page_table"],
            test_data["req_to_token"],
            test_data["req_pool_indices"],
            test_data["strided_indices"],
            test_data["max_seq_pages"],
            test_data["seq_lens"],
            0,
            page_size,
            None,
            None,
        )

        self.assertTrue(
            torch.equal(
                test_data["cache_seqlens_int32"], ref_data["cache_seqlens_int32"]
            )
        )
        self.assertTrue(
            torch.equal(test_data["cu_seqlens_k"], ref_data["cu_seqlens_k"])
        )
        self.assertTrue(torch.equal(test_data["page_table"], ref_data["page_table"]))


if __name__ == "__main__":
    unittest.main()
