"""Unit tests for CP HiCache L3 on-disk primitives (cp_l3_disk): blob header + CRC + slot pool.

Pure-Python (no torch); loaded directly so the sglang import chain isn't needed.
"""

import importlib.util
import sys
import unittest
from pathlib import Path
from unittest.mock import patch

_REPO_ROOT = Path(__file__).resolve().parents[4]


def _load_cp_l3_disk_module():
    module_name = "_test_cp_l3_disk_module"
    module_path = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache" / "cp_l3_disk.py"
    spec = importlib.util.spec_from_file_location(module_name, module_path)
    module = importlib.util.module_from_spec(spec)
    with patch.dict(sys.modules, {module_name: module}):
        spec.loader.exec_module(module)
    return module


m = _load_cp_l3_disk_module()


class TestBlobHeader(unittest.TestCase):
    def _hash(self, seed: int) -> bytes:
        return bytes((seed + i) & 0xFF for i in range(m.CONTENT_HASH_BYTES))

    def test_header_is_64_bytes(self):
        self.assertEqual(m.CP_L3_HEADER_BYTES, 64)

    def test_pack_unpack_roundtrip(self):
        h = self._hash(7)
        payload = bytes(range(256)) * 1024  # 256 KiB
        crc = m.compute_crc(payload)
        raw = m.pack_blob_header(
            payload_kind="target_kv", content_hash=h, n_layers=78, page_bytes=len(payload), crc=crc
        )
        self.assertEqual(len(raw), 64)
        hdr = m.unpack_blob_header(raw)
        self.assertEqual(hdr.payload_kind, "target_kv")
        self.assertEqual(bytes(hdr.content_hash), h)
        self.assertEqual(hdr.n_layers, 78)
        self.assertEqual(hdr.page_bytes, len(payload))
        self.assertEqual(hdr.crc, crc)
        self.assertTrue(m.verify_blob(hdr, payload, expect_hash=h))

    def test_all_payload_kinds_roundtrip(self):
        for kind in ("target_kv", "draft_kv", "index_k"):
            raw = m.pack_blob_header(
                payload_kind=kind, content_hash=self._hash(1), n_layers=21, page_bytes=4096, crc=0
            )
            self.assertEqual(m.unpack_blob_header(raw).payload_kind, kind)

    def test_verify_detects_corruption(self):
        h = self._hash(3)
        payload = b"\xAB" * 4096
        raw = m.pack_blob_header(
            payload_kind="index_k", content_hash=h, n_layers=21, page_bytes=4096, crc=m.compute_crc(payload)
        )
        hdr = m.unpack_blob_header(raw)
        self.assertTrue(m.verify_blob(hdr, payload, expect_hash=h))
        # torn same-length write (bit flip) -> CRC mismatch
        torn = bytearray(payload)
        torn[100] ^= 0x01
        self.assertFalse(m.verify_blob(hdr, bytes(torn)))
        # wrong length
        self.assertFalse(m.verify_blob(hdr, payload[:-1]))
        # wrong expected hash
        self.assertFalse(m.verify_blob(hdr, payload, expect_hash=self._hash(99)))

    def test_bad_magic_fails_loud(self):
        raw = bytearray(m.pack_blob_header(
            payload_kind="target_kv", content_hash=self._hash(0), n_layers=1, page_bytes=4096, crc=0))
        raw[0] ^= 0xFF
        with self.assertRaises(ValueError):
            m.unpack_blob_header(bytes(raw))

    def test_content_hash_hex_conversion(self):
        hexstr = "ab" * 32
        self.assertEqual(m.content_hash_to_bytes(hexstr), bytes([0xAB]) * 32)
        with self.assertRaises(ValueError):
            m.content_hash_to_bytes("zz" * 32)        # not hex
        with self.assertRaises(ValueError):
            m.content_hash_to_bytes("ab" * 16)        # wrong length

    def test_slot_bytes_alignment(self):
        # target_kv page = 2,875,392 B (702*4096); slot = header+payload rounded to 4K
        slot = m.slot_bytes_for_page(2_875_392)
        self.assertEqual(slot % m.CP_L3_DIO_ALIGN, 0)
        self.assertGreaterEqual(slot, 64 + 2_875_392)
        self.assertEqual(slot, 703 * 4096)


class TestSlotPool(unittest.TestCase):
    def test_alloc_order_and_full(self):
        p = m.CpL3SlotPool(4)
        self.assertEqual([p.alloc() for _ in range(4)], [0, 1, 2, 3])
        self.assertEqual(p.num_free, 0)
        self.assertEqual(p.alloc(), -1)  # full -> -1, no raise
        self.assertEqual(p.num_allocated, 4)
        self.assertEqual(p.occupancy, 1.0)

    def test_free_and_realloc(self):
        p = m.CpL3SlotPool(4)
        a, b, c = p.alloc(), p.alloc(), p.alloc()  # 0,1,2
        p.free(b)  # free 1
        self.assertEqual(p.num_free, 2)  # {1,3}
        self.assertEqual(p.alloc(), 1)   # reuse the freed slot
        self.assertEqual(p.alloc(), 3)

    def test_double_free_and_oob_fail_loud(self):
        p = m.CpL3SlotPool(2)
        i = p.alloc()
        p.free(i)
        with self.assertRaises(ValueError):
            p.free(i)            # double free
        with self.assertRaises(ValueError):
            p.free(5)            # out of range
        with self.assertRaises(ValueError):
            p.free(1)            # never allocated

    def test_reset(self):
        p = m.CpL3SlotPool(3)
        p.alloc(); p.alloc()
        p.reset()
        self.assertEqual(p.num_free, 3)
        self.assertEqual(p.num_allocated, 0)
        self.assertEqual([p.alloc() for _ in range(3)], [0, 1, 2])

    def test_construct_invalid(self):
        with self.assertRaises(ValueError):
            m.CpL3SlotPool(0)


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