import unittest

import torch

from sglang.srt.layers.attention.nsa.transform_index import (
    transform_index_page_table_decode_ref,
    transform_index_page_table_prefill_ref,
)
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=1, suite="stage-a-test-cpu")


class TestNSATransformIndex(unittest.TestCase):
    def test_decode_ref_masks_high_out_of_bounds_topk_indices(self):
        page_table = torch.tensor([[10, 11, 12], [20, 21, 22]], dtype=torch.int32)
        topk_indices = torch.tensor([[0, 2, 3, -1], [1, 4, 0, -1]], dtype=torch.int64)

        result = transform_index_page_table_decode_ref(page_table, topk_indices)

        self.assertEqual(result.tolist(), [[10, 12, -1, -1], [21, -1, 20, -1]])

    def test_prefill_ref_masks_high_out_of_bounds_topk_indices(self):
        page_table = torch.tensor([[10, 11, 12], [20, 21, 22]], dtype=torch.int32)
        topk_indices = torch.tensor(
            [[0, 3, -1], [2, 1, 9], [1, 0, -1]], dtype=torch.int64
        )

        result = transform_index_page_table_prefill_ref(
            page_table=page_table,
            topk_indices=topk_indices,
            extend_lens_cpu=[2, 1],
        )

        self.assertEqual(result.tolist(), [[10, -1, -1], [12, 11, -1], [21, 20, -1]])


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