# SPDX-License-Identifier: Apache-2.0
"""Tests for Device-DAX-backed L1 allocation.

The tests use a regular mmap-able file rather than requiring real
``/dev/dax`` hardware. That exercises the allocator contract and storage
manager wiring while keeping CI portable.
"""

# Standard
from typing import Any, cast
import argparse
import gc
import json
import os

# Third Party
import pytest
import torch

# First Party
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
from lmcache.v1.distributed.config import (
    EvictionConfig,
    L1ManagerConfig,
    L1MemoryManagerConfig,
    StorageManagerConfig,
    add_storage_manager_args,
    parse_args_to_config,
)
from lmcache.v1.distributed.error import L1Error
from lmcache.v1.distributed.l1_manager import L1Manager
from lmcache.v1.distributed.l2_adapters.config import (
    L2AdapterConfigBase,
    L2AdaptersConfig,
    get_type_name_for_config,
)
from lmcache.v1.distributed.memory_manager.devdax_l1_memory_manager import (
    DevDaxL1MemoryManager,
)
from lmcache.v1.memory_allocators.devdax_memory_allocator import DevDaxMemoryAllocator
from lmcache.v1.multiprocess.config import add_mp_server_args
from lmcache.v1.multiprocess.engine_context import MPCacheServerContext
import lmcache.v1.memory_management as memory_management


def _make_mmap_file(
    tmp_path, size: int = 4 * 1024 * 1024, name: str = "l1-devdax-test.bin"
) -> str:
    path = tmp_path / name
    with open(path, "wb") as f:
        f.truncate(size)
    return str(path)


def _key(seed: int = 0) -> ObjectKey:
    return ObjectKey(
        chunk_hash=seed.to_bytes(4, "big") + b"\0" * 28,
        model_name="devdax-l1-test",
        kv_rank=0,
    )


def _layout(num_bytes: int = 4096) -> MemoryLayoutDesc:
    return MemoryLayoutDesc(shapes=[torch.Size([num_bytes])], dtypes=[torch.uint8])


def _parse_mp_storage_args(args: list[str]) -> StorageManagerConfig:
    parser = argparse.ArgumentParser()
    add_mp_server_args(parser)
    add_storage_manager_args(parser)
    return parse_args_to_config(parser.parse_args(args))


class _FakeMooncakeL2Config:
    def __init__(self, setup_config: dict[str, str]) -> None:
        self.setup_config = setup_config


class _FakeExt:
    is_pin_supported = True

    def __init__(self, fake_runtime: "_FakeCudaRuntime") -> None:
        self._runtime = fake_runtime

    def pin_memory(self, ptr: int, size: int, flags: int = 0) -> bool:
        self._runtime.register_calls.append((ptr, size, flags))
        return self._runtime.register_error == 0

    def unpin_memory(self, ptr: int) -> bool:
        self._runtime.unregister_calls.append(ptr)
        return True


class _FakeCudaRuntime:
    def __init__(self, register_error: int = 0) -> None:
        self.register_error = register_error
        self.register_calls: list[tuple[int, int, int]] = []
        self.unregister_calls: list[int] = []
        self.synchronize_calls = 0
        self.ext = _FakeExt(self)

    def is_available(self) -> bool:
        return True

    def synchronize(self) -> None:
        self.synchronize_calls += 1

    def cudart(self) -> "_FakeCudaRuntime":
        return self

    def cudaHostRegister(self, ptr: int, size: int, flags: int) -> int:
        self.register_calls.append((ptr, size, flags))
        return self.register_error

    def cudaHostUnregister(self, ptr: int) -> int:
        self.unregister_calls.append(ptr)
        return 0


def _hybrid_storage_config(path: str, adapter_config: object) -> StorageManagerConfig:
    config = StorageManagerConfig(
        l1_manager_config=L1ManagerConfig(
            memory_config=L1MemoryManagerConfig(
                size_in_bytes=1024 * 1024,
                use_lazy=False,
                shm_name="",
                devdax_path=path,
                devdax_size_in_bytes=1024 * 1024,
            )
        ),
        eviction_config=EvictionConfig(eviction_policy="LRU"),
        l2_adapter_config=L2AdaptersConfig(
            adapters=[cast(L2AdapterConfigBase, adapter_config)]
        ),
    )
    return config


def test_devdax_config_rejects_lazy_allocation(tmp_path):
    path = _make_mmap_file(tmp_path)

    with pytest.raises(ValueError, match="--no-l1-use-lazy"):
        L1MemoryManagerConfig(
            size_in_bytes=1024 * 1024,
            use_lazy=True,
            shm_name="",
            devdax_path=path,
        )


def test_devdax_config_rejects_shm(tmp_path):
    path = _make_mmap_file(tmp_path)

    with pytest.raises(ValueError, match="--shm-name"):
        L1MemoryManagerConfig(
            size_in_bytes=1024 * 1024,
            use_lazy=False,
            shm_name="lmcache_l1_pool_test",
            devdax_path=path,
            devdax_size_in_bytes=2 * 1024 * 1024,
        )


def test_devdax_config_accepts_explicit_lazy_and_shm_disable(tmp_path):
    path = _make_mmap_file(tmp_path)

    cfg = L1MemoryManagerConfig(
        size_in_bytes=1024 * 1024,
        use_lazy=False,
        shm_name="",
        devdax_path=path,
    )

    assert cfg.devdax_path == path
    assert cfg.use_lazy is False
    assert cfg.shm_name == ""


@pytest.mark.parametrize(
    ("adapter_name", "adapter_config"),
    [
        ("nixl_store", object()),
        ("nixl_store_dynamic", object()),
        ("mooncake_store", _FakeMooncakeL2Config({"protocol": "rdma"})),
    ],
)
def test_devdax_overflow_rejects_single_region_l2_adapters(
    tmp_path, monkeypatch, adapter_name, adapter_config
):
    path = _make_mmap_file(tmp_path)
    monkeypatch.setattr(
        "lmcache.v1.distributed.config.get_type_name_for_config",
        lambda _: adapter_name,
    )

    with pytest.raises(ValueError, match=adapter_name):
        _hybrid_storage_config(path, adapter_config)


def test_devdax_overflow_allows_mooncake_without_rdma(tmp_path, monkeypatch):
    path = _make_mmap_file(tmp_path)
    monkeypatch.setattr(
        "lmcache.v1.distributed.config.get_type_name_for_config",
        lambda _: "mooncake_store",
    )

    config = _hybrid_storage_config(path, _FakeMooncakeL2Config({"protocol": "tcp"}))

    assert config.l1_manager_config.memory_config.devdax_size_in_bytes == 1024 * 1024


def test_devdax_allocator_uses_mmap_backing_file(tmp_path):
    path = _make_mmap_file(tmp_path)
    allocator = DevDaxMemoryAllocator(
        size=1024 * 1024,
        device_path=path,
        align_bytes=4096,
    )

    objs = allocator.batched_allocate(torch.Size([4096]), torch.uint8, 2)
    assert objs is not None
    first = objs[0]
    assert first.data_ptr == allocator.buffer.data_ptr()
    assert first.shm_offset == 0

    first.raw_tensor.fill_(0x5A)
    allocator.batched_free(objs)
    del first
    del objs
    gc.collect()
    allocator.close()

    with open(path, "rb") as f:
        assert f.read(4096) == bytes([0x5A]) * 4096


def test_devdax_allocator_registers_cuda_host_mapping(tmp_path, monkeypatch):
    path = _make_mmap_file(tmp_path)
    cuda_runtime = _FakeCudaRuntime()
    monkeypatch.setattr(memory_management, "torch_device_type", "cuda")
    monkeypatch.setattr(memory_management, "torch_dev", cuda_runtime)
    monkeypatch.setattr(memory_management, "current_device_spec", cuda_runtime.ext)

    allocator = DevDaxMemoryAllocator(
        size=1024 * 1024,
        device_path=path,
        align_bytes=4096,
    )
    ptr = allocator.buffer.data_ptr()

    assert cuda_runtime.register_calls == [(ptr, 1024 * 1024, 0)]
    allocator.close()
    assert cuda_runtime.unregister_calls == [ptr]


def test_devdax_allocator_falls_back_when_cuda_host_register_fails(
    tmp_path, monkeypatch
):
    path = _make_mmap_file(tmp_path)
    cuda_runtime = _FakeCudaRuntime(register_error=1)
    monkeypatch.setattr(memory_management, "torch_device_type", "cuda")
    monkeypatch.setattr(memory_management, "torch_dev", cuda_runtime)
    monkeypatch.setattr(memory_management, "current_device_spec", cuda_runtime.ext)

    allocator = DevDaxMemoryAllocator(
        size=1024 * 1024,
        device_path=path,
        align_bytes=4096,
    )
    obj = allocator.allocate(torch.Size([4096]), torch.uint8)

    assert cuda_runtime.register_calls == [
        (allocator.buffer.data_ptr(), 1024 * 1024, 0)
    ]
    assert cuda_runtime.unregister_calls == []
    assert obj is not None
    allocator.free(obj)
    del obj
    gc.collect()
    allocator.close()
    assert cuda_runtime.unregister_calls == []


def test_devdax_close_failure_preserves_allocator_state(tmp_path):
    path = _make_mmap_file(tmp_path)
    allocator = DevDaxMemoryAllocator(
        size=1024 * 1024,
        device_path=path,
        align_bytes=4096,
    )
    obj = allocator.allocate(torch.Size([4096]), torch.uint8)
    assert obj is not None

    with pytest.raises(BufferError):
        allocator.close()

    assert allocator.devdax_allocator is not None
    assert allocator.devdax_buffer.numel() == 1024 * 1024

    allocator.free(obj)
    del obj
    gc.collect()
    allocator.close()


def test_l1_manager_round_trip_on_devdax_mapping(tmp_path):
    path = _make_mmap_file(tmp_path)
    cfg = L1ManagerConfig(
        memory_config=L1MemoryManagerConfig(
            size_in_bytes=1024 * 1024,
            use_lazy=False,
            shm_name="",
            devdax_path=path,
        )
    )
    manager = L1Manager(cfg)
    key = _key(1)

    write = manager.reserve_write([key], [False], _layout())
    assert write[key][0] == L1Error.SUCCESS
    obj = write[key][1]
    assert obj is not None
    obj.tensor.fill_(0x23)
    assert manager.finish_write([key])[key] == L1Error.SUCCESS

    read = manager.reserve_read([key])
    assert read[key][0] == L1Error.SUCCESS
    read_obj = read[key][1]
    assert read_obj is not None
    assert int(read_obj.tensor[0]) == 0x23
    assert manager.finish_read([key])[key] == L1Error.SUCCESS

    del write
    del read
    del obj
    del read_obj
    gc.collect()
    manager.close()

    with open(path, "rb") as f:
        assert f.read(1) == bytes([0x23])


def test_devdax_l1_memory_manager_spills_from_dram_to_devdax(tmp_path):
    path = _make_mmap_file(tmp_path, size=8192)
    manager = DevDaxL1MemoryManager(
        L1MemoryManagerConfig(
            size_in_bytes=8192,
            use_lazy=False,
            shm_name="",
            align_bytes=4096,
            devdax_path=path,
            devdax_size_in_bytes=8192,
        )
    )

    error, objs = manager.allocate(_layout(4096), count=3)

    assert error == L1Error.SUCCESS
    assert len(objs) == 3
    assert isinstance(manager._allocator, DevDaxMemoryAllocator)
    assert manager._allocator.local_allocator is not None
    assert objs[0].parent() is manager._allocator.local_allocator
    assert objs[1].parent() is manager._allocator.local_allocator
    assert objs[2].parent() is manager._allocator
    assert objs[0].data_ptr == manager._allocator.local_allocator.buffer.data_ptr()
    assert (
        objs[1].data_ptr == manager._allocator.local_allocator.buffer.data_ptr() + 4096
    )
    assert objs[2].data_ptr == manager._allocator.devdax_buffer.data_ptr()
    used, total = manager.get_memory_usage()
    assert used == 3 * 4096
    assert total == 4 * 4096

    objs[2].raw_tensor.fill_(0x6D)
    manager.free(objs)
    used, total = manager.get_memory_usage()
    assert used == 0
    assert total == 4 * 4096
    manager.close()

    with open(path, "rb") as f:
        assert f.read(4096) == bytes([0x6D]) * 4096


def test_devdax_l1_memory_manager_reports_devdax_desc(tmp_path):
    path = _make_mmap_file(tmp_path)
    manager = DevDaxL1MemoryManager(
        L1MemoryManagerConfig(
            size_in_bytes=1024 * 1024,
            use_lazy=False,
            shm_name="",
            devdax_path=path,
        )
    )

    desc = manager.get_l1_memory_desc()
    used, total = manager.get_memory_usage()

    assert desc.ptr != 0
    assert desc.size == 1024 * 1024
    assert desc.align_bytes == 4096
    assert used == 0
    assert total == 1024 * 1024
    manager.close()


def test_cli_parses_l1_devdax_path(tmp_path):
    path = _make_mmap_file(tmp_path)
    config = _parse_mp_storage_args(
        [
            "--l1-size-gb",
            "1",
            "--eviction-policy",
            "LRU",
            "--no-l1-use-lazy",
            "--shm-name",
            "",
            "--l1-devdax-path",
            path,
        ]
    )

    mem_cfg = config.l1_manager_config.memory_config
    assert mem_cfg.devdax_path == path
    assert mem_cfg.use_lazy is False
    assert mem_cfg.shm_name == ""


def test_cli_rejects_devdax_l1_with_gds_l1(tmp_path):
    path = _make_mmap_file(tmp_path)

    with pytest.raises(ValueError, match="gds-l1-path"):
        _parse_mp_storage_args(
            [
                "--l1-size-gb",
                "1",
                "--eviction-policy",
                "LRU",
                "--no-l1-use-lazy",
                "--shm-name",
                "",
                "--l1-devdax-path",
                path,
                "--gds-l1-path",
                str(tmp_path),
            ]
        )


def test_cli_infers_l1_devdax_overflow_from_registered_dax_adapter(tmp_path):
    path = _make_mmap_file(tmp_path)
    config = _parse_mp_storage_args(
        [
            "--l1-size-gb",
            "1",
            "--eviction-policy",
            "LRU",
            "--no-l1-use-lazy",
            "--shm-name",
            "",
            "--l1-devdax-path",
            path,
            "--l2-adapter",
            ('{"type":"dax","device_path":"%s","max_dax_size_gb":2,"slot_bytes":4096}')
            % path,
        ]
    )

    mem_cfg = config.l1_manager_config.memory_config
    assert mem_cfg.size_in_bytes == 1 << 30
    assert mem_cfg.devdax_path == path
    assert mem_cfg.devdax_size_in_bytes == 2 << 30
    assert mem_cfg.use_lazy is False
    assert mem_cfg.shm_name == ""
    assert config.l2_adapter_config.adapters == []


@pytest.mark.parametrize(
    ("adapter_spec", "expected_adapter_type"),
    [
        (
            {
                "type": "raw_block",
                "device_path": "rawblock-l2.bin",
                "slot_bytes": 8192,
                "capacity_bytes": 16384,
                "meta_total_bytes": 4096,
                "use_odirect": False,
                "meta_enable_periodic": False,
                "load_checkpoint_on_init": False,
                "meta_verify_on_load": False,
            },
            "raw_block",
        ),
    ],
)
def test_cli_hybrid_l1_keeps_ordinary_l2_adapters(
    tmp_path, adapter_spec, expected_adapter_type
):
    path = _make_mmap_file(tmp_path)
    adapter_spec = {
        key: str(tmp_path / value) if key in ("base_path", "device_path") else value
        for key, value in adapter_spec.items()
    }

    config = _parse_mp_storage_args(
        [
            "--l1-size-gb",
            "1",
            "--eviction-policy",
            "LRU",
            "--no-l1-use-lazy",
            "--shm-name",
            "",
            "--l1-devdax-path",
            path,
            "--l2-adapter",
            json.dumps(
                {
                    "type": "dax",
                    "device_path": path,
                    "max_dax_size_gb": 2,
                    "slot_bytes": 4096,
                }
            ),
            "--l2-adapter",
            json.dumps(adapter_spec),
        ]
    )

    mem_cfg = config.l1_manager_config.memory_config
    assert mem_cfg.devdax_size_in_bytes == 2 << 30
    assert len(config.l2_adapter_config.adapters) == 1
    assert (
        get_type_name_for_config(config.l2_adapter_config.adapters[0])
        == expected_adapter_type
    )


def test_cli_hybrid_l1_splits_matching_dax_device_and_keeps_other_l2(tmp_path):
    l1_dax_path = _make_mmap_file(tmp_path, name="l1-devdax.bin")
    l2_dax_path = _make_mmap_file(tmp_path, name="l2-devdax.bin")

    config = _parse_mp_storage_args(
        [
            "--l1-size-gb",
            "1",
            "--eviction-policy",
            "LRU",
            "--no-l1-use-lazy",
            "--shm-name",
            "",
            "--l1-devdax-path",
            l1_dax_path,
            "--l2-adapter",
            json.dumps(
                {
                    "type": "dax",
                    "devices": [
                        {"device_path": l1_dax_path, "max_dax_size_gb": 2},
                        {"device_path": l2_dax_path, "max_dax_size_gb": 3},
                    ],
                    "slot_bytes": 4096,
                    "hotplug_enabled": True,
                    "num_store_workers": 2,
                    "num_lookup_workers": 3,
                    "num_load_workers": 4,
                }
            ),
        ]
    )

    mem_cfg = config.l1_manager_config.memory_config
    assert mem_cfg.devdax_size_in_bytes == 2 << 30
    assert len(config.l2_adapter_config.adapters) == 1

    dax_adapter = cast(Any, config.l2_adapter_config.adapters[0])
    assert get_type_name_for_config(dax_adapter) == "dax"
    assert [device.device_path for device in dax_adapter.devices] == [l2_dax_path]
    assert dax_adapter.max_dax_size_gb == 3
    assert dax_adapter.hotplug_enabled is True
    assert dax_adapter.num_store_workers == 2
    assert dax_adapter.num_lookup_workers == 3
    assert dax_adapter.num_load_workers == 4


def test_devdax_l1_does_not_advertise_shm_pool(tmp_path):
    path = _make_mmap_file(tmp_path)
    config = StorageManagerConfig(
        l1_manager_config=L1ManagerConfig(
            memory_config=L1MemoryManagerConfig(
                size_in_bytes=1024 * 1024,
                use_lazy=False,
                shm_name="",
                devdax_path=path,
            )
        ),
        eviction_config=EvictionConfig(eviction_policy="LRU"),
    )
    context = MPCacheServerContext(config)

    try:
        assert context.shm_pool_info == {"shm_name": "", "pool_size": 0}
        assert os.path.exists(path)
    finally:
        context.storage_manager.close()