import unittest
from importlib import util
from pathlib import Path

import numpy as np

from sglang.test.ci.ci_register import register_cpu_ci


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

_UTILS_PATH = (
    Path(__file__).resolve().parents[4]
    / "python"
    / "sglang"
    / "srt"
    / "disaggregation"
    / "common"
    / "utils.py"
)
_spec = util.spec_from_file_location("disaggregation_common_utils_under_test", _UTILS_PATH)
_utils = util.module_from_spec(_spec)
assert _spec.loader is not None
_spec.loader.exec_module(_utils)
contiguous_group_stats = _utils.contiguous_group_stats


class TestDisaggregationCommonUtils(unittest.TestCase):
    def test_contiguous_group_stats_reports_fragmentation_shape(self):
        src_indices = np.array([1, 2, 3, 10, 18, 19], dtype=np.int32)
        dst_indices = np.array([101, 102, 103, 110, 118, 119], dtype=np.int32)
        src_groups = [[1, 2, 3], [10], [18, 19]]
        dst_groups = [[101, 102, 103], [110], [118, 119]]

        stats = contiguous_group_stats(
            src_indices, dst_indices, src_groups, dst_groups
        )

        self.assertEqual(stats["pages"], 6)
        self.assertEqual(stats["groups"], 3)
        self.assertEqual(stats["min_group_pages"], 1)
        self.assertEqual(stats["max_group_pages"], 3)
        self.assertAlmostEqual(stats["avg_group_pages"], 2.0)
        self.assertEqual(stats["src_diff_head"], [1, 1, 7, 8, 1])
        self.assertEqual(stats["dst_diff_head"], [1, 1, 7, 8, 1])

    def test_contiguous_group_stats_handles_empty_transfer(self):
        stats = contiguous_group_stats(
            np.array([], dtype=np.int32),
            np.array([], dtype=np.int32),
            [],
            [],
        )

        self.assertEqual(stats["pages"], 0)
        self.assertEqual(stats["groups"], 0)
        self.assertEqual(stats["avg_group_pages"], 0.0)
        self.assertEqual(stats["src_diff_head"], [])
        self.assertEqual(stats["dst_diff_head"], [])


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