# SPDX-License-Identifier: Apache-2.0
"""End-to-end benchmark: CacheGen encode/decode on GPU (CUDA or XPU).

Measures wall-clock to serialize and deserialize a synthetic KV-cache
blob, and records the compression ratio against the uncompressed baseline.

Run with: ``pytest tests/benchmarks/test_cachegen.py --benchmark-only``
"""

# Third Party
import pytest
import torch

# First Party
from lmcache import torch_device_type
from lmcache.storage_backend.serde.cachegen_decoder import CacheGenDeserializer
from lmcache.storage_backend.serde.cachegen_encoder import CacheGenSerializer
from lmcache.v1.config import LMCacheEngineConfig
from lmcache.v1.metadata import LMCacheMetadata

# CacheGen has hand-written CUDA and SYCL kernels only; gate the test on
# those backends explicitly (whitelist), so HPU / CPU-only CI / future
# backends without a CacheGen kernel are skipped automatically.
# torch_device_type is set to "cuda"/"xpu" only after is_available() passes
# in lmcache.__init__, so no extra availability check is needed here.
pytestmark = pytest.mark.skipif(
    torch_device_type not in ("cuda", "xpu"),
    reason="CacheGen kernels only exist for CUDA and XPU",
)


def _generate_kv(num_tokens, device):
    num_layers, num_heads, head_size = 32, 8, 128
    shape = [num_tokens, num_heads, head_size]
    pairs = []
    for _ in range(num_layers):
        k = torch.rand(shape, dtype=torch.bfloat16, device=device)
        v = torch.rand(shape, dtype=torch.bfloat16, device=device)
        pairs.append((k, v))
    return torch.stack([torch.stack(p, dim=0) for p in pairs], dim=0)


def _make_serde(chunk_size):
    config = LMCacheEngineConfig.from_defaults(chunk_size=chunk_size)
    metadata = LMCacheMetadata(
        model_name="mistralai/Mistral-7B-Instruct-v0.2",
        world_size=1,
        local_world_size=1,
        worker_id=0,
        local_worker_id=0,
        kv_dtype=torch.bfloat16,
        kv_shape=None,
    )
    return (
        CacheGenSerializer(config, metadata),
        CacheGenDeserializer(config, metadata, torch.bfloat16),
    )


@pytest.mark.benchmark(group="cachegen_encode")
@pytest.mark.parametrize("chunk_size", [64, 128, 256, 768])
def test_cachegen_encoder_bench(benchmark, chunk_size):
    serializer, _ = _make_serde(chunk_size)
    kv = _generate_kv(chunk_size, torch_device_type)

    def run():
        return serializer.to_bytes(kv)

    out = benchmark(run)
    raw_bytes = kv.element_size() * kv.numel()
    print(
        f"\n[chunk={chunk_size}] raw={raw_bytes} compressed={len(out)} "
        f"ratio={raw_bytes / max(len(out), 1):.2f}x"
    )


@pytest.mark.benchmark(group="cachegen_decode")
@pytest.mark.parametrize("chunk_size", [64, 128, 256, 768])
def test_cachegen_decoder_bench(benchmark, chunk_size):
    serializer, deserializer = _make_serde(chunk_size)
    kv = _generate_kv(chunk_size, torch_device_type)
    payload = serializer.to_bytes(kv)

    benchmark(deserializer.from_bytes, payload)