# SPDX-License-Identifier: Apache-2.0
"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2).

Covers the request adapter, the chunk-window spec/manager (registration + the
eviction policy), and the pool builder — exercised against the installed vLLM
KV stack on CPU (block bookkeeping only, no GPU tensors).
"""

import pytest
import torch
from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind
from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry
from vllm.v1.request import RequestStatus

from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
from vllm_omni.experimental.ar_diffusion.kv_cache import (
    ARDiffusionKVCache,
    ARDiffusionKVConfig,
    ARDiffusionRequestAdapter,
    ChunkWindowManager,
    ChunkWindowSpec,
    allocate_kv_pool_with_views,
    build_kv_manager,
    compute_num_blocks,
)
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens

BLOCK = 16


def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
    return ChunkWindowSpec(
        block_size=BLOCK,
        num_kv_heads=4,
        head_size=64,
        dtype=torch.float16,
        sliding_window=window_chunks * chunk_size,
        chunk_size=chunk_size,
        window_chunks=window_chunks,
        sink_chunks=sink_chunks,
        reset_at_boundary=reset_at_boundary,
    )


# --- ChunkWindowSpec registration -------------------------------------------


def test_spec_registration_resolves_to_chunk_window_manager():
    # Without explicit registration the MRO walk would fall back to the parent
    # SlidingWindowManager; assert the subclass manager wins.
    spec = make_spec()
    assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager


def test_spec_kind_is_sliding_window():
    assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW


def test_spec_rejects_inconsistent_window():
    with pytest.raises(ValueError):
        ChunkWindowSpec(
            block_size=BLOCK,
            num_kv_heads=4,
            head_size=64,
            dtype=torch.float16,
            sliding_window=99,  # != window_chunks * chunk_size
            chunk_size=BLOCK,
            window_chunks=2,
        )


# --- eviction policy (pure) -------------------------------------------------


def test_sliding_replace_keeps_window():
    # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens;
    # the snap is to chunk boundaries.
    def skip(n):
        return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)

    assert skip(32) == 0  # nothing past the window yet
    assert skip(48) == 16  # one chunk fell out of the window
    assert skip(64) == 32


def test_sliding_replace_snaps_to_chunk_boundary():
    # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted.
    skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False)
    assert skip % 16 == 0 and skip == 16


def test_sink_chunks_protected():
    # sink = 1 chunk (16 tokens) is never skipped.
    skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
    # raw overflow = 80 - 32 - 16 = 32 -> snapped 32
    assert skip == 32


def test_reset_at_boundary_drops_completed_past_sink():
    skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True)
    # completed = 48; sink = 16 -> drop 32
    assert skip == 32


# --- ARDiffusionKVConfig ------------------------------------------------------------


def test_kv_config_sliding_window_property():
    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48
    assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None


# --- ARDiffusionRequestAdapter ------------------------------------------------------


def test_adapter_advances_per_chunk_not_per_step():
    a = ARDiffusionRequestAdapter("r0", chunk_size=16)
    assert a.num_computed_tokens == 0
    assert a.num_tokens == 16  # in-flight chunk
    a.on_chunk_committed()
    assert a.num_computed_tokens == 16
    assert a.num_tokens == 32


def test_adapter_accounts_for_prefill_prefix():
    a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4)
    assert a.num_computed_tokens == 4
    assert a.num_prompt_tokens == 4
    assert a.num_tokens == 20


def test_adapter_status_is_vllm_enum():
    assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus)


# --- pool / manager ---------------------------------------------------------


def test_compute_num_blocks():
    # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks.
    assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32


def test_paged_pool_layout_exposes_flat_slot_views():
    kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
        num_blocks=4,
        block_size=BLOCK,
        num_layers=1,
        num_kv_heads=4,
        head_dim=64,
        dtype=torch.float32,
        device=torch.device("cpu"),
    )
    assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
    assert k_pools[0].shape == (4 * BLOCK, 4, 64)
    assert v_pools[0].shape == (4 * BLOCK, 4, 64)

    k_pools[0][BLOCK + 3].fill_(7)
    v_pools[0][2 * BLOCK + 5].fill_(11)
    assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
    assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])


def test_build_manager_allocate_free_roundtrip():
    """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager.

    This is the adapter conformance check — if the adapter were missing an
    attribute the manager reads, allocate_slots/free would raise here.
    """
    spec = make_spec()
    mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024)
    free_before = mgr.block_pool.get_num_free_blocks()

    adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK)
    blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True)
    assert blocks is not None
    assert mgr.block_pool.get_num_free_blocks() < free_before

    mgr.free(adapter)
    assert mgr.block_pool.get_num_free_blocks() == free_before


def test_cross_attn_pool_deducted_from_self_attn_budget():
    """The cross-attn pool is allocated directly; its bytes are subtracted from
    the self-attn paged-pool budget so the two together stay within the free
    memory budget (review: zwhzzz0821)."""
    L = 512
    avail = 1 << 30  # 1 GiB
    kv = ARDiffusionKVCache(
        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5),
        num_layers=2,
        num_kv_heads=4,
        head_size=64,
        dtype=torch.float16,
        block_size=BLOCK,
        max_model_len=4096,
        available_bytes=avail,
        kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
        session_capacity=1,
        cross_attention_lengths={"text": L},
        device=torch.device("cpu"),
    )
    cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2  # K+V, pos+neg, layers
    expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
    assert kv.num_blocks == expected
    assert expected > 1 * (2 + 1) + 2  # above the local-slot floor, so the cross deduction is what's tested


def _make_kv(
    *,
    local_branches,
    num_frame_per_block=2,
    window_chunks=9,
    max_scratch_tokens_per_branch=0,
):
    kv_branches = (
        (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0))
        if local_branches == 1
        else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
    )
    return ARDiffusionKVCache(
        ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
        num_layers=1,
        num_kv_heads=4,
        head_size=64,
        dtype=torch.float32,
        block_size=BLOCK,
        max_model_len=4096,
        available_bytes=1 << 16,  # tiny -> the floor binds
        device=torch.device("cpu"),
        kv_branches=kv_branches,
        session_capacity=1,
        frames_per_block=num_frame_per_block,
        max_scratch_tokens_per_branch=max_scratch_tokens_per_branch,
    )


def test_pool_floor_is_branch_aware():
    """CFG-parallel rank (one local kv_branch) sizes for one window + in-flight chunk;
    a single-process run (both kv_branches) sizes for two. Scratch scales the same way."""
    one = _make_kv(local_branches=1)
    two = _make_kv(local_branches=2)

    assert one.managed_num_blocks == 1 * (9 + 2) + 2  # 13
    assert two.managed_num_blocks == 2 * (9 + 2) + 2  # 24
    assert one.scratch_num_blocks == one.scratch_blocks_per_kv_branch
    assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_kv_branch
    assert one.scratch_blocks_per_kv_branch == 2
    assert one.num_blocks_total == 13 + one.scratch_blocks_per_kv_branch


def test_scratch_capacity_is_derived_from_declared_geometry(monkeypatch):
    monkeypatch.delenv("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", raising=False)
    kv = _make_kv(
        local_branches=1,
        num_frame_per_block=6,
        max_scratch_tokens_per_branch=BLOCK + 1,
    )
    assert kv.scratch_blocks_per_kv_branch == 8  # six video blocks + two auxiliary blocks


def test_scratch_env_override_cannot_reduce_declared_minimum(monkeypatch):
    monkeypatch.setenv("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "1")
    kv = _make_kv(
        local_branches=1,
        num_frame_per_block=6,
        max_scratch_tokens_per_branch=BLOCK + 1,
    )
    assert kv.scratch_blocks_per_kv_branch == 8


def test_scratch_maps_to_slot_zero_with_one_local_branch():
    """A CFG-parallel rank runs exactly one kv_branch: whichever CFG side it is,
    its scratch lands in the rank's single slot (no dead second slot)."""
    one = _make_kv(local_branches=1)
    assert one.scratch_block_ids("positive", 0, 2) == one.scratch_block_ids("negative", 0, 2)

    two = _make_kv(local_branches=2)
    assert two.scratch_block_ids("positive", 0, 2) != two.scratch_block_ids("negative", 0, 2)


def test_scratch_exhaustion_still_raises():
    one = _make_kv(local_branches=1)
    cap = one.scratch_blocks_per_kv_branch
    with pytest.raises(RuntimeError, match="scratch blocks exhausted"):
        one.scratch_block_ids("positive", 0, cap + 1)


def test_non_contiguous_branch_indices_rejected():
    with pytest.raises(ValueError, match="contiguous"):
        ARDiffusionKVCache(
            ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2),
            num_layers=1,
            num_kv_heads=4,
            head_size=64,
            dtype=torch.float32,
            block_size=BLOCK,
            max_model_len=4096,
            available_bytes=1 << 16,
            kv_branches=(ARDiffusionKVBranchSpec("main", 1),),
            session_capacity=1,
        )