import sys
import unittest
from unittest.mock import MagicMock

import torch

# Stub out sgl_kernel before any sglang import so this CPU unit test does not
# require CUDA extension libraries to be installed.
for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
    if _mod not in sys.modules:
        sys.modules[_mod] = MagicMock()

from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

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


class TestCpHiCacheNodeMetadata(CustomTestCase):
    def test_split_zero_len_moves_all_positions_to_child(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([1, 3, 7], dtype=torch.int64),
            host_indices=torch.tensor([10, 11, 12], dtype=torch.int64),
        )

        parent, child = metadata.split(0)

        self.assertEqual(parent.logical_len, 0)
        self.assertEqual(parent.owned_positions.tolist(), [])
        self.assertEqual(parent.host_indices.tolist(), [])
        self.assertEqual(child.logical_len, 8)
        self.assertEqual(child.owned_positions.tolist(), [1, 3, 7])
        self.assertEqual(child.host_indices.tolist(), [10, 11, 12])

    def test_split_rebases_child_positions(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=10,
            owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
            host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
        )

        parent, child = metadata.split(5)

        self.assertEqual(parent.logical_len, 5)
        self.assertEqual(parent.owned_positions.tolist(), [0, 2])
        self.assertEqual(parent.host_indices.tolist(), [20, 21])
        self.assertEqual(child.logical_len, 5)
        self.assertEqual(child.owned_positions.tolist(), [0, 4])
        self.assertEqual(child.host_indices.tolist(), [22, 23])

    def test_zero_owned_metadata_is_valid(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=64,
            owned_positions=torch.empty((0,), dtype=torch.int32),
            host_indices=torch.empty((0,), dtype=torch.int32),
        )

        self.assertEqual(metadata.logical_len, 64)
        self.assertEqual(metadata.owned_positions.device.type, "cpu")
        self.assertEqual(metadata.host_indices.device.type, "cpu")
        self.assertEqual(metadata.owned_positions.dtype, torch.int64)
        self.assertEqual(metadata.host_indices.dtype, torch.int64)

    def test_non_int64_inputs_are_converted(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([1, 3], dtype=torch.int32),
            host_indices=torch.tensor([10, 11], dtype=torch.int32),
        )

        self.assertEqual(metadata.owned_positions.dtype, torch.int64)
        self.assertEqual(metadata.host_indices.dtype, torch.int64)

    def test_negative_logical_len_raises(self):
        with self.assertRaisesRegex(ValueError, "logical_len"):
            CpHiCacheNodeMetadata(
                logical_len=-1,
                owned_positions=torch.empty((0,), dtype=torch.int64),
                host_indices=torch.empty((0,), dtype=torch.int64),
            )

    def test_metadata_does_not_alias_input_tensors(self):
        owned_positions = torch.tensor([1, 3], dtype=torch.int64)
        host_indices = torch.tensor([10, 11], dtype=torch.int64)

        metadata = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=owned_positions,
            host_indices=host_indices,
        )
        owned_positions[0] = 2
        host_indices[0] = 12

        self.assertEqual(metadata.owned_positions.tolist(), [1, 3])
        self.assertEqual(metadata.host_indices.tolist(), [10, 11])

    def test_invalid_split_raises(self):
        metadata = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([1], dtype=torch.int64),
            host_indices=torch.tensor([9], dtype=torch.int64),
        )

        with self.assertRaisesRegex(ValueError, "split_len"):
            metadata.split(5)

    def test_unsorted_positions_raise(self):
        with self.assertRaisesRegex(ValueError, "sorted"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([2, 1], dtype=torch.int64),
                host_indices=torch.tensor([9, 10], dtype=torch.int64),
            )

    def test_duplicate_positions_raise(self):
        with self.assertRaisesRegex(ValueError, "strictly increasing"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([1, 1], dtype=torch.int64),
                host_indices=torch.tensor([9, 10], dtype=torch.int64),
            )

    def test_length_mismatch_raises(self):
        with self.assertRaisesRegex(ValueError, "same length"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([1, 2], dtype=torch.int64),
                host_indices=torch.tensor([9], dtype=torch.int64),
            )

    def test_out_of_range_positions_raise(self):
        with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"):
            CpHiCacheNodeMetadata(
                logical_len=4,
                owned_positions=torch.tensor([4], dtype=torch.int64),
                host_indices=torch.tensor([9], dtype=torch.int64),
            )


class FakeWriteFailure:
    metadata = None

    def __init__(self, required_host_slots):
        self.required_host_slots = required_host_slots


class FakeWriteSuccess:
    required_host_slots = 0

    def __init__(self, metadata):
        self.metadata = metadata


class FakeWriteController:
    def __init__(self, required_host_slots):
        self.required_host_slots = required_host_slots
        self.calls = 0
        self.write_policy = "write_through"
        self.evicted_host_indices = []

    def write(self, device_indices, node_id=-1, priority=None):
        self.calls += 1
        if self.calls == 1:
            return FakeWriteFailure(self.required_host_slots)
        return FakeWriteSuccess(
            CpHiCacheNodeMetadata(
                logical_len=len(device_indices),
                owned_positions=torch.tensor([0], dtype=torch.int64),
                host_indices=torch.tensor([99], dtype=torch.int64),
            )
        )

    def evict_host(self, host_indices):
        self.evicted_host_indices.append(host_indices.clone())
        return len(host_indices)


class FakeEvictionStrategy:
    def get_priority(self, node):
        return 0


class TestHiRadixCacheCPBackup(CustomTestCase):
    def test_node_backuped_uses_cp_metadata(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        node = TreeNode()
        node.host_len = 8
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=8,
            owned_positions=torch.tensor([1, 2], dtype=torch.int64),
            host_indices=torch.tensor([10, 11], dtype=torch.int64),
        )

        self.assertTrue(cache._node_backuped(node))

    def test_inc_hit_count_does_not_rewrite_cp_backed_node(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.write_through_threshold = 1
        cache.cache_controller = type(
            "Controller", (), {"write_policy": "write_through"}
        )()
        cache.write_backup = lambda node: (_ for _ in ()).throw(
            AssertionError("must not rewrite")
        )
        node = TreeNode()
        node.host_len = 4
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([], dtype=torch.int64),
            host_indices=torch.tensor([], dtype=torch.int64),
        )

        cache._inc_hit_count(node)

        self.assertEqual(node.hit_count, 1)

    def test_write_backup_retries_by_required_physical_slots(self):
        cache = HiRadixCache.__new__(HiRadixCache)
        cache._uses_cp_hicache = True
        cache.cache_controller = FakeWriteController(required_host_slots=1)
        cache.evictable_host_leaves = set()
        cache.eviction_strategy = FakeEvictionStrategy()
        cache.get_child_key_fn = lambda key: key.token_ids[0]
        cache._record_remove_event = lambda node: None
        cache.ongoing_write_through = {}
        cache.inc_node_lock_ref = lambda node: None

        root = TreeNode()
        root.key = RadixKey(token_ids=[], extra_key=None)
        root.value = []
        cache.root_node = root
        evictable_node = TreeNode()
        evictable_node.parent = root
        evictable_node.key = RadixKey(token_ids=[1], extra_key=None)
        evictable_node.value = None
        evictable_node.host_len = 4
        evictable_node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.tensor([0], dtype=torch.int64),
            host_indices=torch.tensor([55], dtype=torch.int64),
        )
        root.children[1] = evictable_node
        cache.evictable_host_leaves.add(evictable_node)

        node = TreeNode()
        node.value = torch.arange(16, dtype=torch.int64)

        cache.write_backup(node)

        self.assertEqual(
            cache.cache_controller.evicted_host_indices[0].tolist(), [55]
        )
        self.assertEqual(evictable_node.host_len, 0)
        self.assertIsNone(evictable_node.cp_hicache)
        self.assertNotIn(1, root.children)
        self.assertEqual(node.host_len, 16)


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