import unittest

import torch

from sglang.srt.mem_cache.page_index_utils import (
    validate_page_aligned_token_indices,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

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


class TestPageIndexUtils(CustomTestCase):
    def test_valid_empty_indices(self):
        validate_page_aligned_token_indices(
            torch.empty((0,), dtype=torch.int64), 4, "indices"
        )

    def test_valid_single_page(self):
        validate_page_aligned_token_indices(
            torch.tensor([8, 9, 10, 11], dtype=torch.int64), 4, "indices"
        )

    def test_valid_multiple_pages(self):
        validate_page_aligned_token_indices(
            torch.tensor([4, 5, 6, 7, 16, 17, 18, 19], dtype=torch.int64),
            4,
            "indices",
        )

    def test_rejects_non_1d_tensor(self):
        with self.assertRaisesRegex(ValueError, "must be a 1-D tensor"):
            validate_page_aligned_token_indices(
                torch.tensor([[0, 1], [2, 3]], dtype=torch.int64), 4, "indices"
            )

    def test_rejects_non_positive_page_size(self):
        with self.assertRaisesRegex(ValueError, "page_size must be positive"):
            validate_page_aligned_token_indices(
                torch.tensor([0, 1, 2, 3], dtype=torch.int64), 0, "indices"
            )

    def test_rejects_partial_page(self):
        with self.assertRaisesRegex(ValueError, "whole pages"):
            validate_page_aligned_token_indices(
                torch.tensor([8, 9, 10], dtype=torch.int64), 4, "indices"
            )

    def test_rejects_non_page_start(self):
        with self.assertRaisesRegex(ValueError, "start at page boundaries"):
            validate_page_aligned_token_indices(
                torch.tensor([9, 10, 11, 12], dtype=torch.int64), 4, "indices"
            )

    def test_rejects_non_contiguous_group(self):
        with self.assertRaisesRegex(ValueError, "contiguous page spans"):
            validate_page_aligned_token_indices(
                torch.tensor([8, 9, 11, 10], dtype=torch.int64), 4, "indices"
            )


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