# SPDX-License-Identifier: Apache-2.0
# Standard
from pathlib import Path
import asyncio
import contextlib
import functools
import os
import shutil
import tempfile
import threading
import uuid

# Third Party
import pytest
import torch

pytest.importorskip("nixl", reason="nixl package is required for nixl tests")

# Third Party
from nixl._api import nixl_agent as NixlAgent
from nixl._api import nixl_agent_config as NixlAgentConfig

# First Party
from lmcache.utils import CacheEngineKey
from lmcache.v1.config import LMCacheEngineConfig
from lmcache.v1.memory_allocators.paged_tensor_memory_allocator import (
    PagedTensorMemoryAllocator,
)
from lmcache.v1.metadata import LMCacheMetadata
from lmcache.v1.storage_backend import CreateStorageBackends
from lmcache.v1.storage_backend.nixl_storage_backend import (
    NixlStorageBackend,
    NixlStorageConfig,
)
from lmcache.v1.transfer_channel.transfer_utils import get_correct_device

# cuFile-based backends (GDS, GDS_MT) need a GDS-capable filesystem
_TEST_TMPDIR = os.environ.get("LMCACHE_TEST_TMPDIR") or None


@functools.lru_cache(maxsize=None)
def _can_register_file_with_nixl_backend(backend: str) -> bool:
    """Probe ``cuFileHandleRegister`` via NIXL on the test scratch dir."""

    probe_dir = tempfile.mkdtemp(prefix="nixl_gds_probe_", dir=_TEST_TMPDIR)
    probe_path = os.path.join(probe_dir, "probe.bin")
    fd = -1
    try:
        agent = NixlAgent(
            f"NixlGdsProbe_{uuid.uuid4().hex}",
            NixlAgentConfig(backends=[]),
        )
        agent.create_backend(backend, {})
        fd = os.open(probe_path, os.O_CREAT | os.O_RDWR, 0o600)
        os.write(fd, b"\x00" * 4096)
        agent.register_memory([(0, 4096, fd, "")], mem_type="FILE")
        return True
    except Exception:
        return False
    finally:
        if fd >= 0:
            with contextlib.suppress(OSError):
                os.close(fd)
        shutil.rmtree(probe_dir, ignore_errors=True)


_GDS_SKIP_REASON = (
    "NIXL {backend} cannot register file handles in this environment; "
    "set LMCACHE_TEST_TMPDIR to a GDS-capable mount (ext4/xfs) to enable."
)


@pytest.fixture
def nixl_tmp_path():
    """Per-test scratch dir, honoring ``LMCACHE_TEST_TMPDIR``."""
    path = tempfile.mkdtemp(prefix="nixl_test_", dir=_TEST_TMPDIR)
    try:
        yield path
    finally:
        shutil.rmtree(path, ignore_errors=True)


def create_key(chunk_hash: str):
    return CacheEngineKey(
        model_name="meta-llama/Llama-3.1-70B-Instruct",
        world_size=8,
        worker_id=0,
        chunk_hash=int(chunk_hash, base=16),
        dtype=torch.bfloat16,
    )


def run(config: LMCacheEngineConfig, shape, dtype):
    BACKEND_NAME = "NixlStorageBackend"
    keys = []
    objs = []
    keys.append(
        create_key("e3229141e680fb413d2c5d3ebb416c4ad300d381e309fc9e417757b91406c157")
    )
    keys.append(
        create_key("e3229141e680fb413d2c5d3ebb416c4ad300d381e309fc9e417757b91406d268")
    )
    keys.append(
        create_key("e3229141e680fb413d2c5d3ebb416c4ad300d381e309fc9e417757b91406e379")
    )
    bad_key = create_key("deadbeefdeadbeef")

    thread_loop = None
    thread = None
    try:
        thread_loop = asyncio.new_event_loop()
        thread = threading.Thread(target=thread_loop.run_forever)
        thread.start()

        metadata = LMCacheMetadata(
            model_name="Llama-3.1-70B-Instruct",
            world_size=1,
            local_world_size=1,
            worker_id=0,
            local_worker_id=0,
            kv_dtype=dtype,
            kv_shape=shape,
        )

        backends = CreateStorageBackends(
            config,
            metadata,
            thread_loop,
            dst_device=get_correct_device(
                config.nixl_buffer_device, metadata.worker_id
            ),
        )
        assert len(backends) == 2  # NixlStorageBackend + LocalCPUBackend
        assert BACKEND_NAME in backends

        nixl_backend = backends[BACKEND_NAME]
        assert isinstance(nixl_backend, NixlStorageBackend)
        assert isinstance(nixl_backend.memory_allocator, PagedTensorMemoryAllocator)
        assert nixl_backend is not None
        assert nixl_backend.memory_allocator is not None

        # Allocate via the chunk shape that LMCacheMetadata derives from
        # kv_shape (LocalCPUBackend / NIXL share the same paged pool sized
        # by metadata.get_shapes()); passing the raw 5D kv_shape would
        # produce a same-byte-count but differently-indexed tensor.
        alloc_shape = metadata.get_shapes()[0]
        for key in keys:
            assert not nixl_backend.contains(key, False)
            assert not nixl_backend.exists_in_put_tasks(key)

            obj = nixl_backend.memory_allocator.allocate(alloc_shape, dtype)
            assert obj is not None
            assert obj.tensor is not None
            objs.append(obj)

        # small tensor changes for data validation (chunk shape is 4D:
        # [kv_size, num_layers, num_tokens, num_heads * head_size] =
        # [2, 4, 256, 1024] for kv_shape (4, 2, 256, 8, 128))
        objs[0].tensor[0, 0, 100, 200] = 1e-3
        objs[0].tensor[1, 0, 200, 100] = 1e-4

        objs[1].tensor[0, 1, 150, 400] = 1e-2
        objs[1].tensor[1, 1, 100, 300] = 1e-5

        objs[2].tensor[0, 2, 50, 200] = 3e-2
        objs[2].tensor[1, 3, 200, 100] = 4e-5

        # Insert first 2 keys
        first_keys = keys[0:2]
        first_objs = objs[0:2]
        nixl_backend.batched_submit_put_task(first_keys, first_objs)

        for key, obj in zip(first_keys, first_objs, strict=False):
            returned_memory_obj = nixl_backend.get_blocking(key)
            assert returned_memory_obj is not None
            assert returned_memory_obj.get_size() == obj.get_size()
            assert returned_memory_obj.get_shape() == obj.get_shape()
            assert returned_memory_obj.get_dtype() == obj.get_dtype()
            assert returned_memory_obj.metadata.address != obj.metadata.address

            returned_tensor = returned_memory_obj.tensor
            obj_tensor = obj.tensor
            assert returned_tensor is not None
            assert obj_tensor is not None
            assert torch.equal(returned_tensor, obj_tensor)

        obj_list = asyncio.run(
            nixl_backend.batched_get_non_blocking(lookup_id="test", keys=first_keys)
        )

        for i, obj in enumerate(first_objs):
            returned_memory_obj = obj_list[i]
            assert returned_memory_obj is not None
            assert returned_memory_obj.get_size() == obj.get_size()
            assert returned_memory_obj.get_shape() == obj.get_shape()
            assert returned_memory_obj.get_dtype() == obj.get_dtype()
            assert returned_memory_obj.metadata.address != obj.metadata.address

            returned_tensor = returned_memory_obj.tensor
            obj_tensor = obj.tensor
            assert returned_tensor is not None
            assert obj_tensor is not None
            assert torch.equal(returned_tensor, obj_tensor)

        def test_eviction(new_idx, old_idx):
            nixl_backend.batched_submit_put_task([keys[new_idx]], [objs[new_idx]])

            obj = nixl_backend.get_blocking(keys[new_idx])
            assert obj is not None
            assert obj.tensor is not None
            assert torch.equal(obj.tensor, objs[new_idx].tensor)

            obj = nixl_backend.get_blocking(keys[old_idx])
            assert obj is None

        ######## Test bad key lookup #########
        obj = nixl_backend.get_blocking(bad_key)
        assert obj is None

        ######## Test eviction #########
        obj = nixl_backend.get_blocking(keys[0])
        assert obj is not None

        # At this point, key 0 & key 1 are cached. Key 1 is LRU key.
        # Submitting key 2 should evict key 1.

        test_eviction(new_idx=2, old_idx=1)

        ######## Test pin #########
        val = nixl_backend.pin(keys[2])
        assert val is True

        obj = nixl_backend.get_blocking(keys[0])
        assert obj is not None

        # At this point, key 0 & key 2 are cached.
        # Key 2 is LRU key, but is pinned.
        # Submitting key 1 should evict key 0.

        test_eviction(new_idx=1, old_idx=0)

        ######## Test unpin #########
        val = nixl_backend.unpin(keys[2])
        assert val is True

        obj = nixl_backend.get_blocking(keys[1])
        assert obj is not None

        # At this point, key 1 & key 2 are cached.
        # Key 2 is LRU key, and is now unpinned.
        # Submitting key 0 should evict key 2.

        test_eviction(new_idx=0, old_idx=2)

        for backend in backends.values():
            backend.close()

    except Exception:
        raise
    finally:
        if thread_loop and thread_loop.is_running():
            thread_loop.call_soon_threadsafe(thread_loop.stop)
        if thread and thread.is_alive():
            thread.join()


@pytest.mark.no_shared_allocator
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
@pytest.mark.skipif(
    not _can_register_file_with_nixl_backend("GDS_MT"),
    reason=_GDS_SKIP_REASON.format(backend="GDS_MT"),
)
def test_nixl_gds_mt_cuda_backend(nixl_tmp_path):
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")

    dtype = torch.bfloat16
    shape = torch.Size([4, 2, 256, 8, 128])

    config.nixl_buffer_device = "cuda"
    # data/nixl.yaml is CPU-mode (no nixl_buffer_size); CUDA mode sizes the
    # NIXL buffer via nixl_buffer_size, so restore it when flipping the device.
    config.nixl_buffer_size = 1024**3
    config.extra_config["nixl_backend"] = "GDS_MT"
    config.extra_config["enable_cuda"] = True
    config.extra_config["nixl_path"] = nixl_tmp_path

    run(config, shape, dtype)


@pytest.mark.no_shared_allocator
@pytest.mark.skipif(
    not _can_register_file_with_nixl_backend("GDS_MT"),
    reason=_GDS_SKIP_REASON.format(backend="GDS_MT"),
)
def test_nixl_gds_mt_cpu_backend(nixl_tmp_path):
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")

    dtype = torch.bfloat16
    shape = torch.Size([4, 2, 256, 8, 128])

    config.nixl_buffer_device = "cpu"
    config.extra_config["nixl_backend"] = "GDS_MT"
    config.extra_config["enable_cuda"] = False
    config.extra_config["nixl_path"] = nixl_tmp_path

    run(config, shape, dtype)


@pytest.mark.no_shared_allocator
@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
@pytest.mark.skipif(
    not _can_register_file_with_nixl_backend("GDS"),
    reason=_GDS_SKIP_REASON.format(backend="GDS"),
)
def test_nixl_gds_cuda_backend(nixl_tmp_path):
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")

    dtype = torch.bfloat16
    shape = torch.Size([4, 2, 256, 8, 128])

    config.nixl_buffer_device = "cuda"
    # data/nixl.yaml is CPU-mode (no nixl_buffer_size); CUDA mode sizes the
    # NIXL buffer via nixl_buffer_size, so restore it when flipping the device.
    config.nixl_buffer_size = 1024**3
    config.extra_config["nixl_backend"] = "GDS"
    config.extra_config["enable_cuda"] = True
    config.extra_config["nixl_path"] = nixl_tmp_path

    run(config, shape, dtype)


@pytest.mark.no_shared_allocator
@pytest.mark.skipif(
    not _can_register_file_with_nixl_backend("GDS"),
    reason=_GDS_SKIP_REASON.format(backend="GDS"),
)
def test_nixl_gds_cpu_backend(nixl_tmp_path):
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")

    dtype = torch.bfloat16
    shape = torch.Size([4, 2, 256, 8, 128])

    config.nixl_buffer_device = "cpu"
    config.extra_config["nixl_backend"] = "GDS"
    config.extra_config["enable_cuda"] = False
    config.extra_config["nixl_path"] = nixl_tmp_path

    run(config, shape, dtype)


@pytest.mark.no_shared_allocator
def test_nixl_endpoint_list_empty_raises():
    """nixl_endpoint_list=[] should raise ValueError before any nixl ops."""
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")
    config.extra_config["nixl_endpoint_list"] = []

    metadata = LMCacheMetadata(
        model_name="Llama-3.1-70B-Instruct",
        world_size=1,
        local_world_size=1,
        worker_id=0,
        local_worker_id=0,
        kv_dtype=torch.bfloat16,
        kv_shape=(4, 2, 256, 8, 128),
    )

    with pytest.raises(ValueError, match="nixl_endpoint_list is set but empty"):
        NixlStorageConfig.from_cache_engine_config(config, metadata)


@pytest.mark.no_shared_allocator
def test_nixl_endpoint_list_malformed_url_raises():
    """A non-http(s) entry in nixl_endpoint_list should raise ValueError."""
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")
    config.extra_config["nixl_endpoint_list"] = ["htps://typo.example.com"]
    config.extra_config["nixl_backend"] = "OBJ"

    metadata = LMCacheMetadata(
        model_name="Llama-3.1-70B-Instruct",
        world_size=1,
        local_world_size=1,
        worker_id=0,
        local_worker_id=0,
        kv_dtype=torch.bfloat16,
        kv_shape=(4, 2, 256, 8, 128),
    )

    with pytest.raises(ValueError, match="is not a valid URL"):
        NixlStorageConfig.from_cache_engine_config(config, metadata)


def _presence_cache_metadata() -> LMCacheMetadata:
    return LMCacheMetadata(
        model_name="Llama-3.1-70B-Instruct",
        world_size=1,
        local_world_size=1,
        worker_id=0,
        local_worker_id=0,
        kv_dtype=torch.bfloat16,
        kv_shape=(4, 2, 256, 8, 128),
    )


@pytest.mark.no_shared_allocator
def test_nixl_presence_cache_only_defaults_false():
    """presence_cache_only defaults to False when the config key is absent."""
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")

    nixl_config = NixlStorageConfig.from_cache_engine_config(
        config, _presence_cache_metadata()
    )

    assert nixl_config.presence_cache_only is False


@pytest.mark.no_shared_allocator
def test_nixl_presence_cache_only_parsed():
    """nixl_presence_cache_only=True is parsed onto NixlStorageConfig."""
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")
    config.extra_config["nixl_presence_cache"] = True
    config.extra_config["nixl_presence_cache_only"] = True

    nixl_config = NixlStorageConfig.from_cache_engine_config(
        config, _presence_cache_metadata()
    )

    assert nixl_config.presence_cache_only is True


@pytest.mark.no_shared_allocator
def test_nixl_presence_cache_only_requires_presence_cache():
    """nixl_presence_cache_only=True without nixl_presence_cache raises ValueError."""
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")
    config.extra_config["nixl_presence_cache_only"] = True

    with pytest.raises(ValueError, match="nixl_presence_cache must be true"):
        config.validate()


@pytest.mark.no_shared_allocator
def test_nixl_posix_backend(nixl_tmp_path):
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl.yaml")

    dtype = torch.bfloat16
    shape = torch.Size([4, 2, 256, 8, 128])

    config.nixl_buffer_device = "cpu"
    config.extra_config["nixl_backend"] = "POSIX"
    config.extra_config["enable_cuda"] = False
    config.extra_config["nixl_path"] = nixl_tmp_path

    run(config, shape, dtype)


@pytest.mark.no_shared_allocator
def test_nixl_posix_backend_multipath():
    """Test NIXL backend with multipath support and path sharding."""
    BASE_DIR = Path(__file__).parent
    config = LMCacheEngineConfig.from_file(BASE_DIR / "data/nixl_multipath.yaml")

    dtype = torch.bfloat16
    shape = torch.Size([4, 2, 256, 8, 128])

    config.nixl_buffer_device = "cpu"
    config.extra_config["nixl_backend"] = "POSIX"
    config.extra_config["enable_cuda"] = False

    # Test that multipath configuration is properly handled
    assert isinstance(config.extra_config["nixl_path"], list)
    assert len(config.extra_config["nixl_path"]) == 3
    assert config.extra_config["nixl_path_sharding"] == "by_gpu"

    run(config, shape, dtype)