"""Unit tests for CP HiCache L3 config (cp_l3_config): parse, rank->disk mapping, validation, PLP probe."""

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

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


def _load(name, rel):
    module_name = f"_test_{name}_module"
    spec = importlib.util.spec_from_file_location(module_name, _REPO_ROOT / rel)
    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_config", "python/sglang/srt/mem_cache/cp_l3_config.py")
_GIB = 1024 ** 3


def _cfg(**over):
    d = {
        "backend": "posix",
        "disks": [{"path": "/cache/l3", "budget_gb": 2000}, {"path": "/mnt/nvme3/l3", "budget_gb": 1000}],
        "require_plp": True,
        "index_map_gb": 16,
    }
    d.update(over)
    return m.CpL3Config.from_dict(d)


class TestConfigParse(unittest.TestCase):
    def test_from_dict_basic(self):
        c = _cfg()
        self.assertEqual(c.backend, "posix")
        self.assertEqual(len(c.disks), 2)
        self.assertEqual(c.disks[0].budget_bytes, 2000 * _GIB)
        self.assertTrue(c.require_plp)
        self.assertEqual(c.index_map_bytes, 16 * _GIB)
        self.assertEqual(c.index_dir, "/cache/l3/l3_index")  # default = first disk
        self.assertIsNone(c.rank_disk_map)

    def test_accepts_cp_l3_wrapper(self):
        c = m.CpL3Config.from_dict({"cp_l3": {"disks": [{"path": "/d", "budget_gb": 1}]}})
        self.assertEqual(len(c.disks), 1)

    def test_balanced_mapping_8ranks_4disks(self):
        c = _cfg(disks=[{"path": f"/d{i}", "budget_gb": 1} for i in range(4)])
        got = [c.disk_for_rank(r, 8) for r in range(8)]
        self.assertEqual(got, [0, 1, 2, 3, 0, 1, 2, 3])  # 2 ranks/disk, balanced

    def test_balanced_8ranks_3disks_graceful(self):
        c = _cfg(disks=[{"path": f"/d{i}", "budget_gb": 1} for i in range(3)])
        got = [c.disk_for_rank(r, 8) for r in range(8)]
        self.assertEqual(got, [0, 1, 2, 0, 1, 2, 0, 1])  # ranks share disks gracefully

    def test_explicit_rank_disk_map(self):
        c = _cfg(disks=[{"path": "/a", "budget_gb": 1}, {"path": "/b", "budget_gb": 1}],
                 rank_disk_map=[0, 0, 1, 1])
        self.assertEqual([c.disk_for_rank(r, 4) for r in range(4)], [0, 0, 1, 1])

    def test_explicit_map_wrong_length_fails(self):
        c = _cfg(rank_disk_map=[0, 1])
        with self.assertRaises(ValueError):
            c.disk_for_rank(0, 8)  # map len 2 != cp_size 8

    def test_explicit_map_oob_fails(self):
        with self.assertRaises(ValueError):
            _cfg(disks=[{"path": "/a", "budget_gb": 1}], rank_disk_map=[0, 5]).disk_for_rank(0, 2)

    def test_empty_disks_fails(self):
        with self.assertRaises(ValueError):
            m.CpL3Config.from_dict({"disks": []})

    def test_bad_backend_fails(self):
        with self.assertRaises(ValueError):
            _cfg(backend="leveldb")

    def test_from_file_json(self):
        with tempfile.TemporaryDirectory() as td:
            p = os.path.join(td, "l3.json")
            Path(p).write_text('{"cp_l3": {"disks": [{"path": "/x", "budget_gb": 5}], "index_map_gb": 8}}')
            c = m.CpL3Config.from_file(p)
            self.assertEqual(c.disks[0].budget_bytes, 5 * _GIB)
            self.assertEqual(c.index_map_bytes, 8 * _GIB)


class TestPlpProbe(unittest.TestCase):
    def test_probe_returns_bool_and_latency(self):
        with tempfile.TemporaryDirectory() as td:
            is_plp, median_us = m.probe_disk_plp(td, samples=4)
            self.assertIsInstance(is_plp, bool)
            self.assertGreater(median_us, 0.0)


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