import sys
import types
from unittest import main
from unittest.mock import patch

import torch

if "sgl_kernel" not in sys.modules:
    sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
if "sgl_kernel.kvcacheio" not in sys.modules:
    kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio")
    for name in (
        "transfer_kv_all_layer",
        "transfer_kv_all_layer_direct_lf_pf",
        "transfer_kv_all_layer_lf_pf",
        "transfer_kv_all_layer_lf_ph",
        "transfer_kv_all_layer_mla",
        "transfer_kv_all_layer_mla_lf_pf",
        "transfer_kv_direct",
        "transfer_kv_per_layer",
        "transfer_kv_per_layer_direct_pf_lf",
        "transfer_kv_per_layer_mla",
        "transfer_kv_per_layer_mla_pf_lf",
        "transfer_kv_per_layer_pf_lf",
        "transfer_kv_per_layer_ph_lf",
    ):
        setattr(kvcacheio_stub, name, lambda *args, **kwargs: None)
    sys.modules["sgl_kernel.kvcacheio"] = kvcacheio_stub

from sglang.srt.managers.cache_controller import HiCacheController
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
from sglang.srt.mem_cache.radix_cache import 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 FakeHostPool:
    def __init__(self, alloc_result):
        self.alloc_result = alloc_result
        self.alloc_calls = []
        self.backups = []
        self.loads = []
        self.frees = []
        self.page_size = 4
        self.layout = "page_first_direct"

    def alloc(self, need_size):
        self.alloc_calls.append(need_size)
        if self.alloc_result is None:
            return None
        return self.alloc_result[:need_size].clone()

    def backup_from_device_all_layer(
        self, device_pool, host_indices, device_indices, io_backend
    ):
        self.backups.append((host_indices.clone(), device_indices.clone()))

    def load_to_device_per_layer(
        self, device_pool, host_indices, device_indices, layer_id, io_backend
    ):
        self.loads.append((host_indices.clone(), device_indices.clone(), layer_id))

    def free(self, indices):
        self.frees.append(indices.clone())
        return len(indices)


class FakeDevicePool:
    device = "cpu"
    layer_num = 1

    def register_layer_transfer_counter(self, counter):
        self.counter = counter


class FakeAllocator:
    def __init__(self, alloc_result=None):
        self.alloc_result = alloc_result
        self.alloc_calls = []
        self.cp_size = 4
        self.cp_rank = 1

    def get_kvcache(self):
        return FakeDevicePool()

    def alloc(self, need_size):
        self.alloc_calls.append(need_size)
        if self.alloc_result is None:
            return None
        return self.alloc_result[:need_size].clone()


class HostIndicesTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, data):
        return torch.Tensor._make_subclass(cls, data, require_grad=False)

    def to(self, *args, **kwargs):
        raise AssertionError("load_cp should not move host indices before queuing")


class DummyEvent:
    def record(self):
        pass

    def wait(self, stream):
        pass

    def query(self):
        return True

    def synchronize(self):
        pass


class DummyStream:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class DummyDeviceModule:
    Event = DummyEvent
    Stream = DummyStream

    @staticmethod
    def stream(stream):
        return stream


class DummyLayerDoneCounter:
    def __init__(self):
        self.events = [
            type(
                "ProducerEvent",
                (),
                {
                    "start_event": DummyEvent(),
                    "finish_event": DummyEvent(),
                    "complete": lambda self, layer_id: None,
                },
            )()
        ]

    def update_producer(self):
        return 0


class TestHiCacheControllerCPWrite(CustomTestCase):
    def setUp(self):
        self.device_module_patcher = patch(
            "sglang.srt.managers.cache_controller.device_module",
            DummyDeviceModule,
        )
        self.nsa_pool_patcher = patch(
            "sglang.srt.managers.cache_controller.NSATokenToKVPool",
            FakeDevicePool,
        )
        self.device_module_patcher.start()
        self.nsa_pool_patcher.start()
        self.addCleanup(self.device_module_patcher.stop)
        self.addCleanup(self.nsa_pool_patcher.stop)

    def make_controller(self, host_pool, allocator=None, cp_rank=1):
        allocator = allocator or FakeAllocator()
        controller = HiCacheController(
            token_to_kv_pool_allocator=allocator,
            mem_pool_host=host_pool,
            page_size=4,
            tp_group=None,
            load_cache_event=__import__("threading").Event(),
            io_backend="direct",
            cp_shared_kv_layout=CpSharedKVLayout(
                page_size=4, cp_size=4, cp_rank=cp_rank
            ),
        )
        controller.layer_done_counter = DummyLayerDoneCounter()
        return controller

    def test_cp_write_filters_to_owned_physical_locs(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=7)

        self.assertEqual(result.metadata.logical_len, 16)
        self.assertEqual(result.metadata.owned_positions.tolist(), [4, 5, 6, 7])
        self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(host_pool.alloc_calls, [4])
        self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7])

    def test_cp_write_rejects_incomplete_owned_physical_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)

        with self.assertRaisesRegex(ValueError, "physical_device_indices.*whole pages"):
            controller.write(logical_locs, node_id=21)

    def test_cp_write_rejects_non_contiguous_owned_physical_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.tensor([8, 9, 11, 10], dtype=torch.int64)

        with self.assertRaisesRegex(
            ValueError, "physical_device_indices.*contiguous page spans"
        ):
            controller.write(logical_locs, node_id=22)

    def test_cp_write_rejects_non_contiguous_host_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 103, 102], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.arange(8, 12, dtype=torch.int64)

        with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
            controller.write(logical_locs, node_id=23)

    def test_cp_write_zero_owned_returns_metadata_and_noop_ack(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(host_pool, cp_rank=3)
        logical_locs = torch.arange(4, 8, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=8)

        self.assertEqual(result.metadata.logical_len, 4)
        self.assertEqual(result.metadata.host_indices.tolist(), [])
        self.assertEqual(host_pool.alloc_calls, [])
        self.assertEqual(len(controller.ack_write_queue), 1)

    def test_cp_write_allocation_failure_reports_required_host_slots(self):
        host_pool = FakeHostPool(None)
        controller = self.make_controller(host_pool, cp_rank=1)
        logical_locs = torch.arange(4, 20, dtype=torch.int64)

        result = controller.write(logical_locs, node_id=9)

        self.assertEqual(result.required_host_slots, 4)

    def test_generate_storage_config_constructs_config_at_runtime(self):
        controller = HiCacheController.__new__(HiCacheController)
        controller.mem_pool_device = FakeDevicePool()
        controller.mem_pool_host = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller.pp_rank = 1
        controller.pp_size = 2
        controller.enable_storage_metrics = True

        with patch(
            "sglang.srt.managers.cache_controller.is_dp_attention_enabled",
            return_value=False,
        ), patch(
            "sglang.srt.managers.cache_controller.get_tensor_model_parallel_rank",
            return_value=3,
        ), patch(
            "sglang.srt.managers.cache_controller.get_tensor_model_parallel_world_size",
            return_value=4,
        ):
            config = controller._generate_storage_config(
                model_name="test-model",
                storage_backend_extra_config={"tp_lcm_size": 8},
            )

        self.assertEqual(config.tp_rank, 3)
        self.assertEqual(config.tp_size, 4)
        self.assertEqual(config.pp_rank, 1)
        self.assertEqual(config.pp_size, 2)
        self.assertEqual(config.model_name, "test-model")
        self.assertEqual(config.tp_lcm_size, 8)

    def test_attach_storage_backend_rejects_cp_hicache(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        controller = self.make_controller(host_pool)

        with self.assertRaisesRegex(RuntimeError, "CP shared KV.*storage backend"):
            controller.attach_storage_backend("mooncake")


class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
    def test_cp_load_allocates_full_logical_locs_and_transfers_owned_physical_locs(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
        )

        device_indices = controller.load_cp([node], node_id=11)
        controller.start_loading()

        self.assertEqual(device_indices.tolist(), list(range(64, 80)))
        self.assertEqual(allocator.alloc_calls, [16])
        self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])

    def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self):
        host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=3)
        node = TreeNode()
        node.host_len = 4
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=4,
            owned_positions=torch.empty((0,), dtype=torch.int64),
            host_indices=torch.empty((0,), dtype=torch.int64),
        )

        device_indices = controller.load_cp([node], node_id=12)

        self.assertEqual(device_indices.tolist(), [64, 65, 66, 67])
        self.assertEqual(host_pool.loads, [])
        self.assertEqual(len(controller.ack_load_queue), 1)

    def test_cp_load_queues_cpu_host_indices_before_backend_moves(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        host_indices = HostIndicesTensor(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = type(
            "CpHiCacheMetadataStub",
            (),
            {
                "owned_positions": torch.tensor([4, 5, 6, 7], dtype=torch.int64),
                "host_indices": host_indices,
            },
        )()

        controller.load_cp([node], node_id=13)

        queued_op = controller.load_queue[0]
        self.assertEqual(queued_op.host_indices.device.type, "cpu")
        self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103])
        self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23])

    def test_cp_load_rejects_non_contiguous_physical_device_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(
            alloc_result=torch.tensor(
                [64, 65, 66, 67, 68, 69, 71, 70, 72, 73, 74, 75, 76, 77, 78, 79],
                dtype=torch.int64,
            )
        )
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
        )

        with self.assertRaisesRegex(
            ValueError, "physical_device_indices.*contiguous page spans"
        ):
            controller.load_cp([node], node_id=31)

    def test_cp_load_rejects_non_contiguous_host_page(self):
        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
        allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
        node = TreeNode()
        node.host_len = 16
        node.cp_hicache = CpHiCacheNodeMetadata(
            logical_len=16,
            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
            host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
        )

        with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
            controller.load_cp([node], node_id=32)


if __name__ == "__main__":
    main()