# SPDX-License-Identifier: Apache-2.0
# Standard
from collections import OrderedDict
from copy import deepcopy
from unittest.mock import MagicMock, patch
import os
import random
import shlex
import subprocess
import tempfile
import time

# Third Party
import pytest
import torch

# First Party
from lmcache.utils import (
    CacheEngineKey,
    mock_up_broadcast_fn,
    mock_up_broadcast_object_fn,
)
from lmcache.v1.cache_engine import LMCacheEngine, LMCacheEngineBuilder
from lmcache.v1.config import LMCacheEngineConfig
from lmcache.v1.event_manager import EventStatus, EventType

# Local
from .utils import (
    DummyLMCacheAsyncLookupServer,
    check_paged_kv_cache_equal,
    create_gpu_connector,
    create_test_memory_obj,
    dumb_metadata,
    generate_kv_cache_paged_list_tensors,
    generate_tokens,
    has_cufile,
    recover_engine_states,
)

# Optional override for tempfile root. In CI we point this at a GDS-capable
# host-backed mount (see .buildkite/k3_tests/unit/run.sh); locally it's
# unset and tempfile falls back to its default. Direct-I/O-backed paths are
# required for GDS tests (cuFile err=5027 on overlayfs/tmpfs).
_TEST_TMPDIR = os.environ.get("LMCACHE_TEST_TMPDIR") or None


def get_expected_count(token_len, save_unfull_chunk, chunk_size):
    """Calculate expected token count based on save_unfull_chunk setting.

    Args:
        token_len: Total token length
        save_unfull_chunk: Whether to save partial chunks
        chunk_size: Chunk size for alignment

    Returns:
        If save_unfull_chunk is True, returns token_len as-is.
        Otherwise, returns chunk-aligned count (rounded down).
    """
    if save_unfull_chunk:
        return token_len
    return (token_len // chunk_size) * chunk_size


@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_same_retrieve_store(save_unfull_chunk, autorelease_v1):
    device = "cuda"
    num_tokens = 2000
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16

    chunk_size = 256
    kv_shape = (32, 2, chunk_size, 8, 128)

    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_tokens, device)

    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )

    original_retrieved_cache = deepcopy(retrieved_cache)

    slot_mapping = random.sample(range(0, num_blocks * block_size), num_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)

    # Check the kv cache and the retrieval buffer are not the same
    check_paged_kv_cache_equal(retrieved_cache, original_retrieved_cache, slot_mapping)
    with pytest.raises(AssertionError):
        check_paged_kv_cache_equal(retrieved_cache, kv_cache, slot_mapping)
    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size, remote_url=None, save_unfull_chunk=save_unfull_chunk
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )
    """ test retrieve empty """
    ret_mask = engine.retrieve(
        tokens, kvcaches=retrieved_cache, slot_mapping=slot_mapping
    )
    recover_engine_states(engine)

    length = torch.sum(ret_mask)
    assert length == 0
    check_paged_kv_cache_equal(retrieved_cache, original_retrieved_cache, slot_mapping)
    """ test store """
    engine.store(tokens=tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
    recover_engine_states(engine)

    """ Store is async. Need to wait for the store to finish """
    expected_count = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)
    timeout = 1.5
    start_time = time.time()
    while engine.lookup(tokens) < expected_count:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ test retrieve """
    ret_mask = engine.retrieve(
        tokens, kvcaches=retrieved_cache, slot_mapping=slot_mapping
    )
    recover_engine_states(engine)

    length = torch.sum(ret_mask)
    assert length == expected_count
    check_paged_kv_cache_equal(retrieved_cache, kv_cache, slot_mapping[:expected_count])


@pytest.mark.parametrize("chunk_size", [128, 256])
@pytest.mark.parametrize("backend", ["cpu", "local_disk", "remote", "remote_cachegen"])
@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.parametrize("lmserver_v1_process", ["cpu"], indirect=True)
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_retrieve_prefix(
    chunk_size, backend, save_unfull_chunk, lmserver_v1_process, autorelease_v1
):
    url = None
    remote_serde = None
    check_equality = True
    if "remote" in backend:
        url = lmserver_v1_process.server_url
        if backend == "remote_cachegen":
            backend = "remote"
            remote_serde = "cachegen"
            check_equality = False
        else:
            remote_serde = "naive"
    device = "cuda"
    num_tokens = 2000
    new_num_tokens = 1000
    kv_shape = (32, 2, chunk_size, 8, 128)
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16
    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    new_tokens = generate_tokens(new_num_tokens, device)
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    slot_mapping_full = random.sample(
        range(0, num_blocks * block_size), num_tokens + new_num_tokens
    )
    slot_mapping = torch.tensor(slot_mapping_full[:num_tokens], device=device)

    new_slot_mapping = torch.tensor(slot_mapping_full[-new_num_tokens:], device=device)
    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend=backend,
        remote_url=url,
        remote_serde=remote_serde,
        save_unfull_chunk=save_unfull_chunk,
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )
    """ test store """
    t1 = time.perf_counter()
    engine.store(tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
    recover_engine_states(engine)
    t2 = time.perf_counter()
    print(f"store {len(tokens)} takes {t2 - t1}")
    """ Compute expected length """
    expected_length = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)
    """ Store is async. Need to wait for the store to finish """
    if backend == "cpu":
        timeout = 1
        search_range = "LocalCPUBackend"
    elif backend == "local_disk":
        timeout = 30
        search_range = "LocalDiskBackend"
    elif backend == "remote":
        timeout = 30
        search_range = "RemoteBackend"
    start_time = time.time()
    while engine.lookup(tokens, search_range=search_range) < expected_length:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ test retrieve """
    # Get actual stored length - may be less than expected if is_last_prefill=False
    # even when save_unfull_chunk=True
    actual_stored_tokens = engine.lookup(torch.cat([tokens, new_tokens]))
    t4 = time.perf_counter()
    ret_mask = engine.retrieve(
        torch.cat([tokens, new_tokens]),
        kvcaches=retrieved_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)

    length = torch.sum(ret_mask)
    t5 = time.perf_counter()
    print(f"retrieve {length} takes {t5 - t4}")

    # Use actual stored length (may be chunk-aligned even if save_unfull_chunk=True
    # if is_last_prefill=False)
    assert length == actual_stored_tokens

    if check_equality:
        check_paged_kv_cache_equal(
            kv_cache,
            retrieved_cache,
            torch.cat([slot_mapping, new_slot_mapping])[:actual_stored_tokens],
        )

    if backend in ["local_disk"]:
        subprocess.run(shlex.split("rm -rf local/disk_test/local_disk/"))


@pytest.mark.parametrize("chunk_size", [256])
@pytest.mark.parametrize(
    "backend",
    ["cpu", "local_disk", "remote"],
)
@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.parametrize("lmserver_v1_process", ["cpu"], indirect=True)
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_store_offset(
    chunk_size, backend, save_unfull_chunk, lmserver_v1_process, autorelease_v1
):
    url = None
    if backend == "remote":
        url = lmserver_v1_process.server_url
    device = "cuda"
    num_tokens = 2000
    num_suffix_tokens = 500
    num_total_tokens = 3000
    kv_shape = (32, 2, chunk_size, 8, 128)
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16
    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_total_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    slot_mapping = random.sample(range(0, num_blocks * block_size), num_total_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)

    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend=backend,
        remote_url=url,
        save_unfull_chunk=save_unfull_chunk,
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )
    """ test store """
    engine.store(
        tokens[:num_tokens],
        kvcaches=kv_cache,
        slot_mapping=slot_mapping[:num_tokens],
    )

    offset_chunk_cnt = num_tokens // chunk_size
    offset_length = offset_chunk_cnt * chunk_size
    mask = torch.ones(num_tokens + num_suffix_tokens, device=device)
    mask[:offset_length] = 0
    engine.store(
        tokens[: num_tokens + num_suffix_tokens],
        kvcaches=kv_cache,
        mask=mask,
        slot_mapping=slot_mapping[: num_tokens + num_suffix_tokens],
    )
    recover_engine_states(engine)

    """ Compute expected length """
    total_tokens = num_tokens + num_suffix_tokens
    expected_length = (total_tokens // chunk_size) * chunk_size
    """ Store is async. Need to wait for the store to finish """
    if backend == "cpu":
        timeout = 1
    elif backend == "local_disk":
        timeout = 30
    start_time = time.time()
    while engine.lookup(tokens[: num_tokens + num_suffix_tokens]) < expected_length:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ test retrieve """
    t4 = time.perf_counter()
    ret_mask = engine.retrieve(
        tokens, kvcaches=retrieved_cache, slot_mapping=slot_mapping
    )
    recover_engine_states(engine)

    length = torch.sum(ret_mask)
    t5 = time.perf_counter()
    print(f"retrieve {length} takes {t5 - t4}")

    assert length == expected_length
    check_paged_kv_cache_equal(
        kv_cache,
        retrieved_cache,
        slot_mapping[:expected_length],
    )

    if backend in ["local_disk"]:
        subprocess.run(shlex.split("rm -rf local/disk_test/local_disk/"))


@pytest.mark.parametrize("chunk_size", [128])  # , 256])
@pytest.mark.parametrize(
    "backend",
    [
        # "cpu",
        "local_disk"
    ],
)
@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_mixed_retrieve(chunk_size, backend, save_unfull_chunk, autorelease_v1):
    device = "cuda"
    num_tokens = 2000
    new_num_tokens = 1000
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16

    kv_shape = (32, 2, chunk_size, 8, 128)
    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    new_tokens = generate_tokens(new_num_tokens, device)
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )

    slot_mapping_full = random.sample(
        range(0, num_blocks * block_size), num_tokens + new_num_tokens
    )
    slot_mapping = torch.tensor(slot_mapping_full[:num_tokens], device=device)

    new_slot_mapping = torch.tensor(slot_mapping_full[-new_num_tokens:], device=device)

    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size, backend=backend, save_unfull_chunk=save_unfull_chunk
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )
    """ test store """
    engine.store(tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
    engine.store(new_tokens, kvcaches=kv_cache, slot_mapping=new_slot_mapping)
    recover_engine_states(engine)
    """ Store is async. Need to wait for the store to finish """
    expected_length = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)
    if backend == "cpu":
        timeout = 1
        search_range = "LocalCPUBackend"
    elif backend == "local_disk":
        timeout = 30
        search_range = "LocalDiskBackend"
    start_time = time.time()
    while engine.lookup(tokens, search_range=search_range) < expected_length:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ test retrieve """
    # Check actual stored tokens for the combined tokens
    # When tokens are stored separately, the total may be chunk-aligned
    actual_stored_total = engine.lookup(
        torch.cat([tokens, new_tokens]), search_range=search_range
    )
    ret_mask = engine.retrieve(
        torch.cat([tokens, new_tokens]),
        kvcaches=retrieved_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)
    length = torch.sum(ret_mask)
    # Use actual stored total (may be chunk-aligned even if save_unfull_chunk=True
    # if is_last_prefill=False)
    assert length == actual_stored_total
    check_paged_kv_cache_equal(
        retrieved_cache,
        kv_cache,
        torch.cat([slot_mapping, new_slot_mapping])[:length],
    )

    """Wait for store to finish"""
    expected_length = get_expected_count(new_num_tokens, save_unfull_chunk, chunk_size)
    start_time = time.time()
    while engine.lookup(new_tokens, search_range=search_range) < expected_length:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ test another retrieve """
    ret_mask = engine.retrieve(
        new_tokens, kvcaches=retrieved_cache, slot_mapping=new_slot_mapping
    )
    recover_engine_states(engine)
    length = torch.sum(ret_mask)
    assert length == expected_length
    check_paged_kv_cache_equal(
        retrieved_cache, kv_cache, new_slot_mapping[:expected_length]
    )

    """ insert the mixed kv cache """
    final_tokens = torch.cat([tokens, new_tokens])
    engine.store(
        final_tokens,
        kvcaches=kv_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)

    """Wait until store finishes"""
    expected_length = get_expected_count(
        num_tokens + new_num_tokens, save_unfull_chunk, chunk_size
    )
    start_time = time.time()
    while (
        engine.lookup(torch.cat([tokens, new_tokens]), search_range=search_range)
        < expected_length
    ):
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ should retrieve the mixed version """
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    ret_mask = engine.retrieve(
        final_tokens,
        kvcaches=retrieved_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)
    length = torch.sum(ret_mask)
    assert length == expected_length

    # Only check chunk-aligned tokens when save_unfull_chunk=False
    check_paged_kv_cache_equal(
        retrieved_cache,
        kv_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping])[:expected_length],
    )
    """destroy local disk path"""
    if backend in ["local_disk"]:
        subprocess.run(shlex.split("rm -rf local/disk_test/local_disk/"))


@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_store_kv_tensors_mask(save_unfull_chunk, autorelease_v1):
    device = "cuda"
    num_tokens = 1000
    new_num_tokens = 2000
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16

    chunk_size = 256
    kv_shape = (32, 2, chunk_size, 8, 128)
    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype=dtype
    )

    new_tokens = generate_tokens(new_num_tokens, device)
    final_tokens = torch.cat([tokens, new_tokens])

    slot_mapping_full = random.sample(
        range(0, num_blocks * block_size), num_tokens + new_num_tokens
    )
    slot_mapping = torch.tensor(slot_mapping_full[:num_tokens], device=device)

    new_slot_mapping = torch.tensor(slot_mapping_full[-new_num_tokens:], device=device)

    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size, save_unfull_chunk=save_unfull_chunk
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )
    """ Store some tokens with mask """
    engine.store(tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
    recover_engine_states(engine)
    """Wait until store finishes"""
    expected_count = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)
    timeout = 1
    start_time = time.time()
    while engine.lookup(tokens) < expected_count:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)

    prefix_length = engine.lookup(tokens)
    assert prefix_length == expected_count, (
        f"Expected {expected_count} prefix tokens, but got {prefix_length}"
    )
    """ Store more tokens """
    # Re-query prefix_length for final_tokens (original flow)
    prefix_length = engine.lookup(final_tokens)
    # Store requires mask False count to be chunk-aligned
    # When save_unfull_chunk=True, prefix_length may not be chunk-aligned,
    # so we need to round it down to chunk boundary for the mask
    num_falses_for_store = (prefix_length // chunk_size) * chunk_size
    kv_tensor_mask = torch.ones_like(final_tokens, dtype=torch.bool)
    kv_tensor_mask[:num_falses_for_store] = False

    engine.store(
        final_tokens,
        mask=kv_tensor_mask,
        kvcaches=kv_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)
    """Wait until store finishes"""
    expected_final_count = get_expected_count(
        num_tokens + new_num_tokens, save_unfull_chunk, chunk_size
    )
    timeout = 1
    start_time = time.time()
    while engine.lookup(final_tokens) < expected_final_count:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)

    prefix_length = engine.lookup(final_tokens)
    assert prefix_length == expected_final_count, (
        f"Expected {expected_final_count} prefix tokens, but got {prefix_length}"
    )
    """ retrieve the whole cache """
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype=dtype
    )
    ret_mask = engine.retrieve(
        final_tokens,
        kvcaches=retrieved_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)
    length = torch.sum(ret_mask)
    check_paged_kv_cache_equal(
        retrieved_cache,
        kv_cache,
        torch.cat([slot_mapping, new_slot_mapping])[:length],
    )

    """ retrieve cache with some mask:
    """
    # Retrieve requires mask False count to be chunk-aligned
    # Original used chunk_size * 3 (768), which is tokens' chunk-aligned length
    # When save_unfull_chunk=True, we need to ensure chunk alignment
    num_falses = (num_tokens // chunk_size) * chunk_size
    mask = torch.ones_like(final_tokens, dtype=torch.bool)
    mask[:num_falses] = False
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype=dtype
    )
    ret_mask = engine.retrieve(
        final_tokens,
        mask=mask,
        kvcaches=retrieved_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)
    length = torch.sum(ret_mask)
    full_length = num_tokens + new_num_tokens
    expected_length = full_length - num_falses
    # When save_unfull_chunk=False, retrieved length may be chunk-aligned
    expected_retrieved_length = get_expected_count(
        expected_length, save_unfull_chunk, chunk_size
    )
    assert length == expected_retrieved_length

    with pytest.raises(AssertionError):
        check_paged_kv_cache_equal(
            retrieved_cache,
            kv_cache,
            torch.cat([slot_mapping, new_slot_mapping])[:full_length],
        )
    check_paged_kv_cache_equal(
        retrieved_cache,
        kv_cache,
        torch.cat([slot_mapping, new_slot_mapping])[num_falses : num_falses + length],
    )

    mask[: num_falses + 5] = False
    with pytest.raises(ValueError):
        engine.retrieve(
            final_tokens,
            mask=mask,
            kvcaches=retrieved_cache,
            slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
        )
        recover_engine_states(engine)


@pytest.mark.parametrize("chunk_size", [128])
@pytest.mark.parametrize(
    "backend",
    [
        "local_cpu_disk_remote",
    ],
)
@pytest.mark.parametrize(
    "retrieve_from",
    [
        "local_cpu",
        "local_disk",
        "remote",
    ],
)
@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.parametrize("lmserver_v1_process", ["cpu"], indirect=True)
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_hierarchy_retrieve(
    chunk_size,
    backend,
    retrieve_from,
    save_unfull_chunk,
    lmserver_v1_process,
    autorelease_v1,
):
    url = None
    if backend == "local_cpu_disk_remote":
        url = lmserver_v1_process.server_url
    device = "cuda"
    num_tokens = 2000
    new_num_tokens = 1000
    kv_shape = (32, 2, chunk_size, 8, 128)
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16

    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype=dtype
    )

    new_tokens = generate_tokens(new_num_tokens, device)
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype=dtype
    )

    slot_mapping = random.sample(
        range(0, num_blocks * block_size), num_tokens + new_num_tokens
    )
    slot_mapping = torch.tensor(slot_mapping[:num_tokens], device=device)

    new_slot_mapping = torch.tensor(slot_mapping[-new_num_tokens:], device=device)

    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend=backend,
        remote_url=url,
        save_unfull_chunk=save_unfull_chunk,
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )
    """ test store """
    t1 = time.perf_counter()
    engine.store(tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
    recover_engine_states(engine)
    t2 = time.perf_counter()
    print(f"store {len(tokens)} takes {t2 - t1}")
    """ Compute expected length """
    expected_length = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)
    """ Store is async. Need to wait for the store to finish """
    timeout = 1
    start_time = time.time()
    while engine.lookup(tokens) < expected_length:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ Wait until disk save is finished """
    if retrieve_from in ["local_disk", "remote"]:
        engine.storage_manager.clear(locations=["LocalCPUBackend"])
        timeout = 30
        start_time = time.time()
        while (
            engine.lookup(tokens, search_range=["LocalDiskBackend"]) < expected_length
        ):
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)
    """ Wait until remote save is finished """
    if retrieve_from == "remote":
        engine.storage_manager.clear(locations=["LocalCPUBackend"])
        # FIXME: change this `clear`
        engine.storage_manager.storage_backends["LocalDiskBackend"].dict.clear()
        timeout = 30
        start_time = time.time()
        while engine.lookup(tokens, search_range=["RemoteBackend"]) < expected_length:
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)
    """ test retrieve """
    t4 = time.perf_counter()
    # Get actual stored length
    actual_stored = engine.lookup(torch.cat([tokens, new_tokens]))
    ret_mask = engine.retrieve(
        torch.cat([tokens, new_tokens]),
        kvcaches=retrieved_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
    )
    recover_engine_states(engine)

    length = torch.sum(ret_mask)
    t5 = time.perf_counter()
    print(f"retrieve {length} takes {t5 - t4}")

    # Use actual stored length for assertion
    assert length == actual_stored
    check_paged_kv_cache_equal(
        retrieved_cache,
        kv_cache,
        torch.cat([slot_mapping, new_slot_mapping])[:actual_stored],
    )

    """ Wait until disk save is finished before deleting the directory"""
    if backend in ["local_cpu_disk"]:
        engine.storage_manager.clear(locations=["LocalCPUBackend"])
        timeout = 30
        start_time = time.time()
        while engine.lookup(tokens) < expected_length:
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)

    if backend in ["local_cpu_disk"]:
        subprocess.run(shlex.split("rm -rf local/disk_test/local_disk/"))


@pytest.mark.parametrize(
    "backend",
    [
        "local_cpu_disk",
    ],
)
@pytest.mark.parametrize(
    "prefetch_from",
    [
        "local_disk",
    ],
)
@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_prefetch_retrieve(
    backend, prefetch_from, save_unfull_chunk, autorelease_v1
):
    device = "cuda"
    num_tokens = 2000
    new_num_tokens = 1000
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16
    test_lookup_id = "test_lookup_id"

    chunk_size = 256
    kv_shape = (32, 2, chunk_size, 8, 128)
    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype=dtype
    )
    new_tokens = generate_tokens(new_num_tokens, device)
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype=dtype
    )

    slot_mapping = random.sample(
        range(0, num_blocks * block_size), num_tokens + new_num_tokens
    )
    slot_mapping = torch.tensor(slot_mapping[:num_tokens], device=device)

    new_slot_mapping = torch.tensor(slot_mapping[-new_num_tokens:], device=device)

    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend=backend,
        enable_async_loading=True,
        save_unfull_chunk=save_unfull_chunk,
    )

    async_lookup_server = DummyLMCacheAsyncLookupServer()
    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        ),
        async_lookup_server=async_lookup_server,
    )

    """ test store """
    t1 = time.perf_counter()
    engine.store(tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
    recover_engine_states(engine)
    t2 = time.perf_counter()
    print(f"store {len(tokens)} takes {t2 - t1}")
    """ Compute expected length """
    # For prefetch retrieve, we need to check what was actually stored
    # Since this test uses async operations, we check the actual lookup result
    expected_length = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)
    """ Wait for cpu store to finish """
    timeout = 1
    start_time = time.time()
    actual_lookup = engine.lookup(tokens)
    while actual_lookup < expected_length:
        if time.time() - start_time > timeout:
            raise TimeoutError(f"Operation timed out after {timeout} seconds.")
        time.sleep(0.01)
    """ Delete cpu cache and wait until disk save finishes."""
    if prefetch_from == "local_disk":
        engine.storage_manager.clear(locations=["LocalCPUBackend"])
        timeout = 30
        start_time = time.time()
        while engine.lookup(tokens) < expected_length:
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.1)
    """ Wait until disk load (prefetch) finishes and delete disk cache"""
    engine.async_lookup_and_prefetch(
        lookup_id=test_lookup_id, tokens=torch.cat([tokens, new_tokens])
    )

    if prefetch_from == "local_disk":
        timeout = 60
        start_time = time.time()
        while (
            engine.event_manager.get_event_status(EventType.LOADING, test_lookup_id)
            != EventStatus.DONE
        ):
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)
        engine.storage_manager.storage_backends["LocalDiskBackend"].dict.clear()
    """ test retrieve """
    t4 = time.perf_counter()

    # Get actual stored length for retrieve
    actual_stored = engine.lookup(torch.cat([tokens, new_tokens]))
    ret_mask = engine.retrieve(
        torch.cat([tokens, new_tokens])[:actual_stored],
        kvcaches=retrieved_cache,
        slot_mapping=torch.cat([slot_mapping, new_slot_mapping]),
        req_id=test_lookup_id,
    )
    recover_engine_states(engine)

    length = torch.sum(ret_mask)
    t5 = time.perf_counter()
    print(f"retrieve {length} takes {t5 - t4}")

    assert length == actual_stored
    check_paged_kv_cache_equal(
        retrieved_cache,
        kv_cache,
        torch.cat([slot_mapping, new_slot_mapping])[:actual_stored],
    )

    if backend in ["local_cpu_disk"]:
        subprocess.run(shlex.split("rm -rf local/disk_test/local_disk/"))


@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_async_lookup_and_prefetch_layerwise(autorelease_v1):
    # Regression: before the fix, async_lookup_and_prefetch sent chunk-level
    # CacheEngineKey objects into batched_async_contains, but store_layer
    # populates hot_cache with per-layer LayerCacheEngineKey objects, so every
    # lookup reported 0 hits. We stage hot_cache the way store_layer would and
    # assert async_lookup_and_prefetch reports the full token count.
    chunk_size = 256
    num_layers = 4
    num_chunks = 3
    num_tokens = chunk_size * num_chunks
    kv_shape = (num_layers, 2, chunk_size, 8, 128)
    lookup_id = "layerwise-async-1"

    captured: dict[str, int] = {}

    class _RecordingAsyncLookupServer:
        def send_response_to_scheduler(
            self, lookup_id: str, retrieved_length: int
        ) -> None:
            captured[lookup_id] = retrieved_length

    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend="cpu",
        enable_async_loading=True,
    )
    cfg.use_layerwise = True

    connector = create_gpu_connector(1024, num_layers)
    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        ),
        async_lookup_server=_RecordingAsyncLookupServer(),
    )

    tokens = generate_tokens(num_tokens, "cuda")
    chunk_keys = [
        key for _, _, key in engine.token_database.process_tokens(tokens=tokens)
    ]
    assert len(chunk_keys) == num_chunks

    # Populate LocalCPUBackend.hot_cache the way store_layer would: one entry
    # per (chunk, layer) keyed by LayerCacheEngineKey.
    cpu_backend = engine.storage_manager.storage_backends["LocalCPUBackend"]
    for chunk_key in chunk_keys:
        for layer_key in chunk_key.split_layers(num_layers):
            cpu_backend.submit_put_task(layer_key, create_test_memory_obj())

    engine.async_lookup_and_prefetch(lookup_id=lookup_id, tokens=tokens)

    deadline = time.time() + 10
    while (
        engine.event_manager.get_event_status(EventType.LOADING, lookup_id)
        != EventStatus.DONE
    ):
        if time.time() > deadline:
            raise TimeoutError("layerwise async lookup did not finish in time")
        time.sleep(0.01)

    assert captured.get(lookup_id) == num_tokens, (
        f"Expected retrieved_length={num_tokens}, got {captured.get(lookup_id)}"
    )


@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_async_lookup_and_prefetch_layerwise_partial_layer_missing(autorelease_v1):
    # When a single per-layer key is missing for one chunk, that chunk must be
    # rounded down to a miss (the `// keys_per_chunk` round-down path).
    chunk_size = 256
    num_layers = 4
    num_chunks = 3
    num_tokens = chunk_size * num_chunks
    kv_shape = (num_layers, 2, chunk_size, 8, 128)
    lookup_id = "layerwise-async-2"

    captured: dict[str, int] = {}

    class _RecordingAsyncLookupServer:
        def send_response_to_scheduler(
            self, lookup_id: str, retrieved_length: int
        ) -> None:
            captured[lookup_id] = retrieved_length

    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend="cpu",
        enable_async_loading=True,
    )
    cfg.use_layerwise = True

    connector = create_gpu_connector(1024, num_layers)
    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        ),
        async_lookup_server=_RecordingAsyncLookupServer(),
    )

    tokens = generate_tokens(num_tokens, "cuda")
    chunk_keys = [
        key for _, _, key in engine.token_database.process_tokens(tokens=tokens)
    ]

    # Stage chunks 0 and 1 fully; drop the last per-layer key of chunk 2.
    cpu_backend = engine.storage_manager.storage_backends["LocalCPUBackend"]
    for i, chunk_key in enumerate(chunk_keys):
        per_layer_keys = chunk_key.split_layers(num_layers)
        if i == len(chunk_keys) - 1:
            per_layer_keys = per_layer_keys[:-1]
        for layer_key in per_layer_keys:
            cpu_backend.submit_put_task(layer_key, create_test_memory_obj())

    engine.async_lookup_and_prefetch(lookup_id=lookup_id, tokens=tokens)

    deadline = time.time() + 10
    while (
        engine.event_manager.get_event_status(EventType.LOADING, lookup_id)
        != EventStatus.DONE
    ):
        if time.time() > deadline:
            raise TimeoutError("layerwise async lookup did not finish in time")
        time.sleep(0.01)

    # First two chunks are complete; the partially-evicted third chunk is a
    # miss, so the prefix-match retrieval pattern reports 2 chunks worth.
    assert captured.get(lookup_id) == 2 * chunk_size, (
        f"Expected retrieved_length={2 * chunk_size}, got {captured.get(lookup_id)}"
    )


@pytest.mark.parametrize("chunk_size", [256])
@pytest.mark.parametrize(
    "backend",
    [
        "cpu",
        "local_disk",
        "remote",
        "local_disk_remote",
        "local_cpu_disk_remote",
    ],
)
@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.no_shared_allocator
@pytest.mark.parametrize("lmserver_v1_process", ["cpu"], indirect=True)
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_mem_leak(
    chunk_size, backend, save_unfull_chunk, lmserver_v1_process, autorelease_v1
):
    url = None
    if "remote" in backend:
        url = lmserver_v1_process.server_url

    device = "cuda"
    num_tokens = 2000
    kv_shape = (32, 2, chunk_size, 8, 128)
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16
    connector = create_gpu_connector(1024, 32)

    tokens = generate_tokens(num_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    slot_mapping = random.sample(range(0, num_blocks * block_size), num_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)
    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend=backend,
        remote_url=url,
        save_unfull_chunk=save_unfull_chunk,
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )

    engine.store(tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
    recover_engine_states(engine)

    expected_length = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)
    timeout = 30
    """Wait until cpu store finishes"""
    if "cpu" in backend:
        start_time = time.time()
        while engine.lookup(tokens, search_range=["LocalCPUBackend"]) < expected_length:
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)
    """Wait until disk store finishes"""
    if "disk" in backend:
        start_time = time.time()
        while (
            engine.lookup(tokens, search_range=["LocalDiskBackend"]) < expected_length
        ):
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)

    if "remote" in backend:
        start_time = time.time()
        while engine.lookup(tokens, search_range=["RemoteBackend"]) < expected_length:
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)

    tensor_memory_allocator = (
        engine.storage_manager.allocator_backend.memory_allocator.pin_allocator
    )
    if "cpu" not in backend:
        assert tensor_memory_allocator.total_allocated_size == 0
    else:
        assert tensor_memory_allocator.total_allocated_size > 0

    if "disk" in backend:
        subprocess.run(shlex.split("rm -rf local/disk_test/local_disk/"))


@pytest.mark.parametrize("chunk_size", [256])
@pytest.mark.parametrize(
    "backend",
    [
        "cpu",
        "local_disk",
    ],
)
@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.no_shared_allocator
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_paged_retrieve_after_eviction(
    chunk_size, backend, save_unfull_chunk, autorelease_v1
):
    device = "cuda"
    # NOTE: The default backend cache size is 2 GB.
    # 10000 tokens ia around 1.3 GB so a second retrieve will cause an eviction.
    num_tokens = 10000
    kv_shape = (32, 2, chunk_size, 8, 128)
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16
    connector = create_gpu_connector(1024, 32)

    tokens_1 = generate_tokens(num_tokens, device)
    tokens_2 = generate_tokens(num_tokens, device)
    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    slot_mapping_1 = random.sample(range(0, num_blocks * block_size), num_tokens)
    slot_mapping_1 = torch.tensor(slot_mapping_1, device=device)
    slot_mapping_2 = random.sample(range(0, num_blocks * block_size), num_tokens)
    slot_mapping_2 = torch.tensor(slot_mapping_2, device=device)
    """ initialize the engine """
    cfg = LMCacheEngineConfig.from_legacy(
        chunk_size=chunk_size,
        backend=backend,
        save_unfull_chunk=save_unfull_chunk,
    )

    engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            "test",
            cfg,
            dumb_metadata(kv_shape),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )

    expected_length = get_expected_count(num_tokens, save_unfull_chunk, chunk_size)

    engine.store(tokens_1, kvcaches=kv_cache, slot_mapping=slot_mapping_1)
    recover_engine_states(engine)

    timeout = 30
    if "disk" in backend:
        start_time = time.time()
        while (
            engine.lookup(tokens_1, search_range=["LocalDiskBackend"]) < expected_length
        ):
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)

    engine.store(tokens_2, kvcaches=kv_cache, slot_mapping=slot_mapping_2)
    recover_engine_states(engine)

    """Wait until cpu store finishes"""
    if "cpu" in backend:
        start_time = time.time()
        while (
            engine.lookup(tokens_2, search_range=["LocalCPUBackend"]) < expected_length
        ):
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)
        assert (
            engine.lookup(tokens_1, search_range=["LocalCPUBackend"]) < expected_length
        )

    """Wait until disk store finishes"""
    if "disk" in backend:
        start_time = time.time()
        while (
            engine.lookup(tokens_2, search_range=["LocalDiskBackend"]) < expected_length
        ):
            if time.time() - start_time > timeout:
                raise TimeoutError(f"Operation timed out after {timeout} seconds.")
            time.sleep(0.01)
        assert (
            engine.lookup(tokens_1, search_range=["LocalDiskBackend"]) < expected_length
        )

    ret_mask = engine.retrieve(
        tokens_1,
        kvcaches=retrieved_cache,
        slot_mapping=slot_mapping_1,
    )
    recover_engine_states(engine)
    length = torch.sum(ret_mask)
    assert length < num_tokens

    ret_mask = engine.retrieve(
        tokens_2,
        kvcaches=retrieved_cache,
        slot_mapping=slot_mapping_2,
    )
    recover_engine_states(engine)
    length = torch.sum(ret_mask)
    assert length == expected_length

    if backend in ["local_disk"]:
        subprocess.run(shlex.split("rm -rf local/disk_test/local_disk/"))


def test_builder(autorelease_v1):
    instance_id = "test"
    cfg = LMCacheEngineConfig.from_legacy(chunk_size=256)
    cfg2 = LMCacheEngineConfig.from_legacy(chunk_size=512)
    connector = None
    should_be_none = LMCacheEngineBuilder.get(instance_id)
    assert should_be_none is None

    _engine = autorelease_v1(
        LMCacheEngineBuilder.get_or_create(
            instance_id,
            cfg,
            dumb_metadata(),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )
    )
    _engine2 = autorelease_v1(LMCacheEngineBuilder.get(instance_id))  # noqa

    with pytest.raises(ValueError):
        LMCacheEngineBuilder.get_or_create(
            instance_id,
            cfg2,
            dumb_metadata(),
            connector,
            mock_up_broadcast_fn,
            mock_up_broadcast_object_fn,
        )


@pytest.mark.no_shared_allocator
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_force_store_wait(autorelease_v1):
    device = "cuda"
    num_tokens = 10000
    num_blocks = 5000
    block_size = 16
    dtype = torch.bfloat16

    chunk_size = 256
    kv_shape = (32, 2, chunk_size, 8, 128)

    connector = create_gpu_connector(1024, 32)

    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )

    num_requests = 8

    def generate_random_slot_mapping(num_blocks, block_size, num_tokens, device):
        slot_mapping = random.sample(range(0, num_blocks * block_size), num_tokens)
        return torch.tensor(slot_mapping, device=device)

    list_tokens = [generate_tokens(num_tokens, device) for _ in range(num_requests)]
    list_slot_mappings = [
        generate_random_slot_mapping(num_blocks, block_size, num_tokens, device)
        for _ in range(num_requests)
    ]

    with tempfile.TemporaryDirectory(
        dir=_TEST_TMPDIR, ignore_cleanup_errors=True
    ) as temp_dir:
        cfg = LMCacheEngineConfig.from_defaults(
            local_cpu=False,
            max_local_cpu_size=2,  # small cpu buffer
            local_disk=temp_dir,
            max_local_disk_size=20,
            extra_config={"force_store_wait": True},
        )

        engine = autorelease_v1(
            LMCacheEngineBuilder.get_or_create(
                "test",
                cfg,
                dumb_metadata(kv_shape),
                connector,
                mock_up_broadcast_fn,
                mock_up_broadcast_object_fn,
            )
        )

        # Store kv cache into slow devices
        for t, s in zip(list_tokens, list_slot_mappings, strict=False):
            engine.store(t, kvcaches=kv_cache, slot_mapping=s)

        # Sleep 10 seconds for the last request
        time.sleep(20)

        # No KV cache should be skipped
        # With default save_unfull_chunk=False, we expect chunk-aligned count
        chunk_size = 256
        for t in list_tokens:
            expected_count = (len(t) // chunk_size) * chunk_size
            assert engine.lookup(t) == expected_count


@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_builder_destroy(autorelease_v1):
    """Test the destroy method of LMCacheEngineBuilder"""
    instance_id = "test_destroy"
    cfg = LMCacheEngineConfig.from_legacy(chunk_size=256)
    connector = create_gpu_connector(1024, 32)

    # Verify instance doesn't exist initially
    should_be_none = LMCacheEngineBuilder.get(instance_id)
    assert should_be_none is None

    # Create an engine instance
    engine = LMCacheEngineBuilder.get_or_create(
        instance_id,
        cfg,
        dumb_metadata(),
        connector,
        mock_up_broadcast_fn,
        mock_up_broadcast_object_fn,
    )

    # Verify instance exists
    retrieved_engine = LMCacheEngineBuilder.get(instance_id)
    assert retrieved_engine is not None
    assert retrieved_engine is engine

    # Verify internal state is populated
    assert instance_id in LMCacheEngineBuilder._instances
    assert instance_id in LMCacheEngineBuilder._cfgs
    assert instance_id in LMCacheEngineBuilder._metadatas
    assert instance_id in LMCacheEngineBuilder._stat_loggers

    # Destroy the instance
    LMCacheEngineBuilder.destroy(instance_id)

    # Verify instance is completely removed
    should_be_none_after_destroy = LMCacheEngineBuilder.get(instance_id)
    assert should_be_none_after_destroy is None

    # Verify all internal state is cleaned up
    assert instance_id not in LMCacheEngineBuilder._instances
    assert instance_id not in LMCacheEngineBuilder._cfgs
    assert instance_id not in LMCacheEngineBuilder._metadatas
    assert instance_id not in LMCacheEngineBuilder._stat_loggers

    # Verify destroying non-existent instance doesn't raise error
    LMCacheEngineBuilder.destroy("non_existent_id")  # Should not raise


@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="TODO: Add non-CUDA implementation to VLLMPagedMemGPUConnectorV2",
)
def test_builder_destroy_multiple_instances(autorelease_v1):
    """Test destroying one instance doesn't affect others"""
    instance_id1 = "test_destroy_1"
    instance_id2 = "test_destroy_2"
    cfg = LMCacheEngineConfig.from_legacy(chunk_size=256)
    connector = create_gpu_connector(1024, 32)

    # Create two engine instances
    engine1 = LMCacheEngineBuilder.get_or_create(
        instance_id1,
        cfg,
        dumb_metadata(),
        connector,
        mock_up_broadcast_fn,
        mock_up_broadcast_object_fn,
    )

    engine2 = LMCacheEngineBuilder.get_or_create(
        instance_id2,
        cfg,
        dumb_metadata(),
        connector,
        mock_up_broadcast_fn,
        mock_up_broadcast_object_fn,
    )

    # Verify both instances exist
    assert LMCacheEngineBuilder.get(instance_id1) is engine1
    assert LMCacheEngineBuilder.get(instance_id2) is engine2

    # Destroy only the first instance
    LMCacheEngineBuilder.destroy(instance_id1)

    # Verify first instance is destroyed but second remains
    assert LMCacheEngineBuilder.get(instance_id1) is None
    assert LMCacheEngineBuilder.get(instance_id2) is engine2

    # Verify internal state for first instance is cleaned up
    assert instance_id1 not in LMCacheEngineBuilder._instances
    assert instance_id1 not in LMCacheEngineBuilder._cfgs
    assert instance_id1 not in LMCacheEngineBuilder._metadatas
    assert instance_id1 not in LMCacheEngineBuilder._stat_loggers

    # Verify internal state for second instance remains
    assert instance_id2 in LMCacheEngineBuilder._instances
    assert instance_id2 in LMCacheEngineBuilder._cfgs
    assert instance_id2 in LMCacheEngineBuilder._metadatas
    assert instance_id2 in LMCacheEngineBuilder._stat_loggers

    # Clean up second instance
    LMCacheEngineBuilder.destroy(instance_id2)


@pytest.mark.parametrize("save_unfull_chunk", [False, True])
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="Requires CUDA for test_multi_device_backends",
)
@pytest.mark.skipif(
    not has_cufile(),
    reason="Requires NVIDIA cuFile (libcufile.so). "
    "Skipping on systems without GDS/cuFile (e.g., AMD ROCm).",
)
def test_multi_device_backends(save_unfull_chunk, autorelease_v1):
    """Test running GPU-related backend with local CPU backends
    together
    """
    device = "cuda"
    num_tokens = 2000
    chunk_size = 256  # Default chunk size for this test
    num_blocks = 1000
    block_size = 16
    dtype = torch.bfloat16

    connector = create_gpu_connector(1024, 32)
    metadata = dumb_metadata()
    metadata.model_name = "test-model"  # NOTE: Gds does not accept name with '_'

    tokens = generate_tokens(num_tokens, device)

    kv_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )
    retrieved_cache = generate_kv_cache_paged_list_tensors(
        num_blocks, device, block_size, dtype
    )

    original_retrieved_cache = deepcopy(retrieved_cache)

    slot_mapping = random.sample(range(0, num_blocks * block_size), num_tokens)
    slot_mapping = torch.tensor(slot_mapping, device=device)

    # Check the kv cache and the retrieval buffer are not the same
    check_paged_kv_cache_equal(retrieved_cache, original_retrieved_cache, slot_mapping)
    with pytest.raises(AssertionError):
        check_paged_kv_cache_equal(retrieved_cache, kv_cache, slot_mapping)

    with tempfile.TemporaryDirectory(
        dir=_TEST_TMPDIR, ignore_cleanup_errors=True
    ) as temp_dir:
        cfg = LMCacheEngineConfig.from_dict(
            {
                "local_cpu": True,
                "max_local_cpu_size": 5,
                "gds_path": temp_dir,
                "gds_buffer_size": 1024,
                "save_unfull_chunk": save_unfull_chunk,
                "extra_config": {
                    "use_direct_io": True,
                },
            }
        )

        connector = create_gpu_connector(1024, 32)

        engine = autorelease_v1(
            LMCacheEngineBuilder.get_or_create(
                "engine",
                cfg,
                metadata,
                connector,
                mock_up_broadcast_fn,
                mock_up_broadcast_object_fn,
            )
        )

        """ test store """
        engine.store(tokens=tokens, kvcaches=kv_cache, slot_mapping=slot_mapping)
        recover_engine_states(engine)
        time.sleep(3)  # wait a bit to finish the store

        """ Test lookup """
        expected_count = get_expected_count(len(tokens), save_unfull_chunk, chunk_size)
        ret = engine.lookup(tokens)
        assert ret == expected_count

        ret_cpu = engine.lookup(tokens, search_range=["LocalCPUBackend"])
        assert ret_cpu == expected_count

        ret_gds = engine.lookup(tokens, search_range=["GdsBackend"])
        assert ret_gds == expected_count

        """ Test retrieve """
        ret_mask = engine.retrieve(
            tokens, kvcaches=retrieved_cache, slot_mapping=slot_mapping
        )
        recover_engine_states(engine)
        length = torch.sum(ret_mask)
        assert length == expected_count
        # Only check chunk-aligned tokens when save_unfull_chunk=False
        check_paged_kv_cache_equal(
            retrieved_cache, kv_cache, slot_mapping[:expected_count]
        )

        LMCacheEngineBuilder.destroy("engine")


def _make_key(chunk_hash: int) -> CacheEngineKey:
    """Create a CacheEngineKey for testing."""
    return CacheEngineKey("test", 1, 0, chunk_hash, torch.bfloat16)


def _make_mock_memory_obj(size: int = 1024) -> MagicMock:
    """Create a mock MemoryObj that tracks ref_count_down calls."""
    mock = MagicMock()
    mock.get_size.return_value = size
    return mock


def _make_mock_engine(
    process_tokens_results: list,
    block_mapping: dict,
    batched_get_side_effect: list,
) -> MagicMock:
    """Create a mock engine with the attributes needed by
    _process_tokens_internal.

    Args:
        process_tokens_results: list of (start, end, key) tuples that
            token_database.process_tokens will yield.
        block_mapping: dict returned by storage_manager.get_block_mapping.
        batched_get_side_effect: list of return values for successive
            storage_manager.batched_get calls (one per location).

    Returns:
        A MagicMock configured as a minimal LMCacheEngine.
    """
    engine = MagicMock()
    engine.token_database.process_tokens.return_value = process_tokens_results
    engine.storage_manager.get_block_mapping.return_value = block_mapping
    engine.storage_manager.batched_get.side_effect = batched_get_side_effect
    engine.lookup_pins = {}
    return engine


def test_process_tokens_single_location_boundary_failure():
    """The block whose end equals last_failed_block_start covers
    [start, last_failed_block_start) — entirely before the gap — and
    must be kept."""
    k0, k1 = _make_key(0), _make_key(1)
    mem0 = _make_mock_memory_obj()

    engine = _make_mock_engine(
        process_tokens_results=[(0, 10, k0), (10, 20, k1)],
        block_mapping=OrderedDict(
            [
                ("LocationA", [(k0, 0, 10), (k1, 10, 20)]),
            ]
        ),
        batched_get_side_effect=[[mem0, None]],
    )

    ret_mask = torch.zeros(20, dtype=torch.bool)
    chunks, tot_kv_size = LMCacheEngine._process_tokens_internal(
        engine, torch.zeros(20, dtype=torch.long), None, ret_mask
    )

    assert len(chunks) == 1
    assert chunks[0][0] == k0
    assert chunks[0][1] is mem0
    assert tot_kv_size == 1024
    assert ret_mask[:10].all()
    assert not ret_mask[10:].any()
    mem0.ref_count_down.assert_not_called()


def test_process_tokens_early_failure_truncates_later_location():
    """When an early location fails, blocks successfully retrieved from
    a later location (covering higher positions) must be discarded and
    freed because they are past the gap."""
    k0, k1 = _make_key(0), _make_key(1)
    k2, k3 = _make_key(2), _make_key(3)
    mem0 = _make_mock_memory_obj()
    mem2 = _make_mock_memory_obj()
    mem3 = _make_mock_memory_obj()

    engine = _make_mock_engine(
        process_tokens_results=[
            (0, 10, k0),
            (10, 20, k1),
            (20, 30, k2),
            (30, 40, k3),
        ],
        block_mapping=OrderedDict(
            [
                ("LocationA", [(k0, 0, 10), (k1, 10, 20)]),
                ("LocationB", [(k2, 20, 30), (k3, 30, 40)]),
            ]
        ),
        batched_get_side_effect=[
            [mem0, None],
            [mem2, mem3],
        ],
    )

    ret_mask = torch.zeros(40, dtype=torch.bool)
    chunks, tot_kv_size = LMCacheEngine._process_tokens_internal(
        engine, torch.zeros(40, dtype=torch.long), None, ret_mask
    )

    assert len(chunks) == 1
    assert chunks[0][0] == k0
    assert tot_kv_size == 1024
    assert ret_mask[:10].all()
    assert not ret_mask[10:].any()
    mem0.ref_count_down.assert_not_called()
    mem2.ref_count_down.assert_called_once()
    mem3.ref_count_down.assert_called_once()


def test_process_tokens_multi_location_both_fail_takes_min():
    """When failures occur in multiple locations, the earliest failure
    start (MIN) should be used so that everything after the first gap
    is discarded."""
    k0, k1 = _make_key(0), _make_key(1)
    k2, k3 = _make_key(2), _make_key(3)
    mem0 = _make_mock_memory_obj()
    mem2 = _make_mock_memory_obj()

    engine = _make_mock_engine(
        process_tokens_results=[
            (0, 10, k0),
            (10, 20, k1),
            (20, 30, k2),
            (30, 40, k3),
        ],
        block_mapping=OrderedDict(
            [
                ("LocationA", [(k0, 0, 10), (k1, 10, 20)]),
                ("LocationB", [(k2, 20, 30), (k3, 30, 40)]),
            ]
        ),
        batched_get_side_effect=[
            [mem0, None],
            [mem2, None],
        ],
    )

    ret_mask = torch.zeros(40, dtype=torch.bool)
    chunks, tot_kv_size = LMCacheEngine._process_tokens_internal(
        engine, torch.zeros(40, dtype=torch.long), None, ret_mask
    )

    assert len(chunks) == 1
    assert chunks[0][0] == k0
    assert tot_kv_size == 1024
    assert ret_mask[:10].all()
    assert not ret_mask[10:].any()
    mem0.ref_count_down.assert_not_called()
    mem2.ref_count_down.assert_called_once()


def test_process_tokens_no_failure():
    """When all blocks are retrieved successfully, every chunk should
    be returned and no ref_count_down should be called."""
    k0, k1 = _make_key(0), _make_key(1)
    mem0 = _make_mock_memory_obj()
    mem1 = _make_mock_memory_obj()

    engine = _make_mock_engine(
        process_tokens_results=[(0, 10, k0), (10, 20, k1)],
        block_mapping=OrderedDict(
            [
                ("LocationA", [(k0, 0, 10), (k1, 10, 20)]),
            ]
        ),
        batched_get_side_effect=[[mem0, mem1]],
    )

    ret_mask = torch.zeros(20, dtype=torch.bool)
    chunks, tot_kv_size = LMCacheEngine._process_tokens_internal(
        engine, torch.zeros(20, dtype=torch.long), None, ret_mask
    )

    assert len(chunks) == 2
    assert tot_kv_size == 2048
    assert ret_mask[:20].all()
    mem0.ref_count_down.assert_not_called()
    mem1.ref_count_down.assert_not_called()


def test_process_tokens_unused_keys_no_double_free():
    """A key returned non-None by batched_get but coming after a None
    (unused) should be freed exactly once in the per-location cleanup
    and never again in post-processing."""
    k0, k1, k2 = _make_key(0), _make_key(1), _make_key(2)
    mem0 = _make_mock_memory_obj()
    mem2 = _make_mock_memory_obj()

    engine = _make_mock_engine(
        process_tokens_results=[(0, 10, k0), (10, 20, k1), (20, 30, k2)],
        block_mapping=OrderedDict(
            [
                ("LocationA", [(k0, 0, 10), (k1, 10, 20), (k2, 20, 30)]),
            ]
        ),
        batched_get_side_effect=[[mem0, None, mem2]],
    )

    ret_mask = torch.zeros(30, dtype=torch.bool)
    chunks, tot_kv_size = LMCacheEngine._process_tokens_internal(
        engine, torch.zeros(30, dtype=torch.long), None, ret_mask
    )

    assert len(chunks) == 1
    assert chunks[0][0] == k0
    assert tot_kv_size == 1024
    assert ret_mask[:10].all()
    assert not ret_mask[10:].any()
    mem0.ref_count_down.assert_not_called()
    mem2.ref_count_down.assert_called_once()


def test_process_tokens_first_block_fails():
    """When the very first block fails, no chunks should be returned
    and ret_mask should be all False."""
    k0, k1 = _make_key(0), _make_key(1)
    mem1 = _make_mock_memory_obj()

    engine = _make_mock_engine(
        process_tokens_results=[(0, 10, k0), (10, 20, k1)],
        block_mapping=OrderedDict(
            [
                ("LocationA", [(k0, 0, 10), (k1, 10, 20)]),
            ]
        ),
        batched_get_side_effect=[[None, mem1]],
    )

    ret_mask = torch.zeros(20, dtype=torch.bool)
    chunks, tot_kv_size = LMCacheEngine._process_tokens_internal(
        engine, torch.zeros(20, dtype=torch.long), None, ret_mask
    )

    assert len(chunks) == 0
    assert tot_kv_size == 0
    assert not ret_mask.any()
    mem1.ref_count_down.assert_called_once()


def test_compress_decompress_unpin_not_pinned() -> None:
    """Verify that compress and decompress check is_pinned before calling unpin()

    This prevents negative pin counts and double unpin warnings on backends like
    Remote/S3.
    """
    # Create mock memory objects
    mock_mem_obj = MagicMock()
    mock_mem_obj.is_pinned = False  # Not pinned (e.g. from Remote/S3 backend)

    mock_compressed_mem_obj = MagicMock()
    mock_compressed_mem_obj.is_pinned = False  # Not pinned

    # Mock serializer and deserializer
    mock_serializer = MagicMock()
    mock_serializer.serialize.return_value = mock_compressed_mem_obj
    mock_deserializer = MagicMock()
    mock_deserializer.deserialize.return_value = mock_mem_obj

    # Create a mock engine
    engine = MagicMock(spec=LMCacheEngine)
    engine.metadata = MagicMock()
    engine.config = MagicMock()
    engine.lookup_pins = {"event_123": {"remote": ["key1"]}}

    # Mock engine.lookup to return number of tokens (non-zero)
    engine.lookup.return_value = 100

    # Mock storage_manager methods
    engine.storage_manager = MagicMock()
    # For compress: batched_get returns mock_mem_obj
    engine.storage_manager.batched_get.return_value = [mock_mem_obj]

    with patch(
        "lmcache.v1.storage_backend.naive_serde.CreateSerde",
        return_value=(mock_serializer, mock_deserializer),
    ):
        # Call compress on the engine
        res = LMCacheEngine.compress(
            engine,
            tokens=[1, 2, 3],
            method="cachegen",
            location="remote",
            event_id="event_123",
        )

        assert res == 100
        # Verify that serialize was called
        mock_serializer.serialize.assert_called_once_with(mock_mem_obj)
        # Verify that unpin was NOT called since is_pinned = False
        mock_mem_obj.unpin.assert_not_called()

        # Verify batched_remove and batched_put were called on storage_manager
        engine.storage_manager.batched_remove.assert_called_once_with(
            ["key1"], locations=["remote"]
        )
        engine.storage_manager.batched_put.assert_called_once_with(
            keys=["key1"],
            memory_objs=[mock_compressed_mem_obj],
            location="remote",
        )

    # Reset storage_manager mock
    engine.storage_manager.reset_mock()
    # For decompress: batched_get returns mock_compressed_mem_obj
    engine.storage_manager.batched_get.return_value = [mock_compressed_mem_obj]

    with patch(
        "lmcache.v1.storage_backend.naive_serde.CreateSerde",
        return_value=(mock_serializer, mock_deserializer),
    ):
        res_decomp = LMCacheEngine.decompress(
            engine,
            tokens=[1, 2, 3],
            method="cachegen",
            location="remote",
            event_id="event_123",
        )

        assert res_decomp == 100
        # Verify that deserialize was called
        mock_deserializer.deserialize.assert_called_once_with(mock_compressed_mem_obj)
        # Verify that unpin was NOT called since is_pinned = False
        mock_compressed_mem_obj.unpin.assert_not_called()

        # Verify batched_remove and batched_put were called on storage_manager
        engine.storage_manager.batched_remove.assert_called_once_with(
            ["key1"], locations=["remote"]
        )
        engine.storage_manager.batched_put.assert_called_once_with(
            keys=["key1"],
            memory_objs=[mock_mem_obj],
            location="remote",
        )


def test_compress_decompress_unpin_when_pinned() -> None:
    """Verify that compress and decompress call unpin() if is_pinned is True"""
    # Create mock memory objects
    mock_mem_obj = MagicMock()
    mock_mem_obj.is_pinned = True

    mock_compressed_mem_obj = MagicMock()
    mock_compressed_mem_obj.is_pinned = True

    # Mock serializer and deserializer
    mock_serializer = MagicMock()
    mock_serializer.serialize.return_value = mock_compressed_mem_obj
    mock_deserializer = MagicMock()
    mock_deserializer.deserialize.return_value = mock_mem_obj

    # Create a mock engine
    engine = MagicMock(spec=LMCacheEngine)
    engine.metadata = MagicMock()
    engine.config = MagicMock()
    engine.lookup_pins = {"event_123": {"local_cpu": ["key1"]}}
    engine.lookup.return_value = 100
    engine.storage_manager = MagicMock()

    # Test compress
    engine.storage_manager.batched_get.return_value = [mock_mem_obj]
    with patch(
        "lmcache.v1.storage_backend.naive_serde.CreateSerde",
        return_value=(mock_serializer, mock_deserializer),
    ):
        LMCacheEngine.compress(
            engine,
            tokens=[1, 2, 3],
            method="cachegen",
            location="local_cpu",
            event_id="event_123",
        )
        mock_mem_obj.unpin.assert_called_once()

    # Test decompress
    engine.storage_manager.reset_mock()
    engine.storage_manager.batched_get.return_value = [mock_compressed_mem_obj]
    with patch(
        "lmcache.v1.storage_backend.naive_serde.CreateSerde",
        return_value=(mock_serializer, mock_deserializer),
    ):
        LMCacheEngine.decompress(
            engine,
            tokens=[1, 2, 3],
            method="cachegen",
            location="local_cpu",
            event_id="event_123",
        )
        mock_compressed_mem_obj.unpin.assert_called_once()


def test_retrieve_cleanup_ref_count_and_unpin() -> None:
    """Verify that retrieve() unpins and ref_count_downs all retrieved chunks.

    Specifically in the else branch when remove_after_retrieve is False.
    """
    # Create mock memory objects
    mem_obj_pinned = MagicMock()
    mem_obj_pinned.is_pinned = True

    mem_obj_not_pinned = MagicMock()
    mem_obj_not_pinned.is_pinned = False

    # Mock engine instance
    engine = MagicMock(spec=LMCacheEngine)
    engine.is_healthy.return_value = True
    engine.remove_after_retrieve = False
    engine._is_passive.return_value = False
    engine.save_only_first_rank = False
    engine._get_req_id.return_value = "req_123"

    # We want retrieve to process chunks
    k0 = _make_key(0)
    k1 = _make_key(1)
    reordered_chunks = [
        (k0, mem_obj_pinned, 0, 10),
        (k1, mem_obj_not_pinned, 10, 20),
    ]

    engine.async_loading = False
    engine._process_tokens_internal.return_value = (reordered_chunks, 1024)
    engine._is_sync_pd_backend.return_value = False

    # Mock stats monitor
    engine.stats_monitor = MagicMock()
    mock_stats = MagicMock()
    mock_stats.time_to_retrieve.return_value = 1.0
    engine.stats_monitor.on_retrieve_request.return_value = mock_stats

    # Mock gpu_connector
    engine.gpu_connector = MagicMock()

    # Call retrieve
    tokens = torch.zeros(20, dtype=torch.long)
    kvcaches = torch.zeros(20)
    slot_mapping = torch.zeros(20, dtype=torch.long)

    LMCacheEngine.retrieve(
        engine,
        tokens=tokens,
        kvcaches=kvcaches,
        slot_mapping=slot_mapping,
    )

    # Assertions
    mem_obj_pinned.ref_count_down.assert_called_once()
    mem_obj_pinned.unpin.assert_called_once()

    mem_obj_not_pinned.ref_count_down.assert_called_once()
    mem_obj_not_pinned.unpin.assert_not_called()