# SPDX-License-Identifier: Apache-2.0
"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1)."""

import pytest
import torch

from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
DIMS = dict(
    num_layers=2,
    num_kv_heads=4,
    head_size=64,
    dtype=torch.float16,
    block_size=16,
    kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1)),
    session_capacity=2,
)


def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24):
    cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks)
    return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS)


def test_requires_enabled_config():
    with pytest.raises(ValueError):
        ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS)


def test_requires_bounded_window():
    cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None)
    with pytest.raises(ValueError):
        ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS)


def test_full_request_lifecycle_and_eviction():
    """begin -> per-chunk allocate/slots/commit over a long rollout -> free.

    Exercises the orchestrator end-to-end and asserts the chunk window bounds
    memory (pool plateaus) and frees cleanly.
    """
    kv = make_cache(chunk_size=16, window_chunks=2)
    free_total = kv.manager.block_pool.get_num_free_blocks()

    adapter = kv.begin_request("req-0")
    free_after = []
    for k in range(10):
        block_table = kv.allocate_chunk(adapter)
        slots = kv.chunk_write_slots(adapter)
        # Slots target real blocks in this chunk's table; length == chunk_size.
        assert len(slots) == kv.spec.chunk_size
        used = {int(s) // kv.block_size for s in slots}
        assert kv.null_block_id not in used
        assert used <= set(block_table)
        # Allocation may transiently add one in-flight chunk.
        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
        kv.commit_chunk(adapter)
        # A successful commit immediately prunes physical ownership.
        assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
        free_after.append(kv.manager.block_pool.get_num_free_blocks())

    # Pool memory plateaus once the window is full (eviction recycles blocks).
    assert free_after[-1] == free_after[kv.spec.window_chunks - 1]

    kv.end_request(adapter)
    assert kv.manager.block_pool.get_num_free_blocks() == free_total


def test_num_computed_advances_per_chunk():
    kv = make_cache(chunk_size=16, window_chunks=2)
    a = kv.begin_request("r")
    assert a.num_computed_tokens == 0
    kv.allocate_chunk(a)
    kv.commit_chunk(a)
    assert a.num_computed_tokens == 16
    kv.end_request(a)


def test_state_close_frees_both_branch_blocks():
    """ARDiffusionKVState.close() returns both CFG kv_branches' pool blocks to the pool.

    This is the primitive the runner's LRU eviction relies on: when a session is
    evicted, close() must free the blocks both adapters hold, or session churn
    leaks pool ownership (review P1).
    """
    from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState

    kv = make_cache(chunk_size=16, window_chunks=2)
    free_total = kv.manager.block_pool.get_num_free_blocks()

    pos = kv.begin_request("bde__s")
    neg = kv.begin_request("bde__s__neg")
    state = ARDiffusionKVState(
        kv,
        "s",
        {"positive": pos, "negative": neg},
        num_layers=kv.num_layers,
    )
    for _ in range(3):
        for adapter in (pos, neg):
            kv.allocate_chunk(adapter)
            kv.commit_chunk(adapter)
    # Both kv_branches hold resident blocks now.
    assert kv.manager.block_pool.get_num_free_blocks() < free_total

    state.close()
    # All blocks returned — no leak across an evicted session.
    assert kv.manager.block_pool.get_num_free_blocks() == free_total


@pytest.mark.parametrize(
    "frames_per_call,block_size,max_model_len,num_calls",
    [(1, 16, 113, 1000), (3, 16, 113, 1000), (3, 1560, 1 << 20, 241)],
    ids=["single-frame", "three-frame", "lingbot-two-minutes"],
)
@pytest.mark.parametrize("sink_chunks,reset_at_boundary", [(0, False), (2, False), (0, True), (2, True)])
def test_rollout_keeps_metadata_bounded(
    frames_per_call, block_size, max_model_len, num_calls, sink_chunks, reset_at_boundary
):
    """Long rollouts retain exact global-position K/V with bounded storage indices."""
    kv = ARDiffusionKVCache(
        ARDiffusionKVConfig(
            enable=True,
            chunk_size=block_size,
            window_chunks=3,
            sink_chunks=sink_chunks,
            reset_at_boundary=reset_at_boundary,
        ),
        num_layers=1,
        num_kv_heads=1,
        head_size=4,
        dtype=torch.float32,
        block_size=block_size,
        max_model_len=max_model_len,
        available_bytes=1 << 24,
        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
        session_capacity=1,
        frames_per_block=frames_per_call,
        device=torch.device("cpu"),
    )
    adapter = kv.begin_request("long-rollout")
    state = ARDiffusionKVState(kv, "long-rollout", {"main": adapter}, num_layers=1)
    free_total = kv.manager.block_pool.get_num_free_blocks()
    pool_shape = kv.key_cache(0).shape
    retained_frames = sink_chunks + (0 if reset_at_boundary else 3)
    for call_index in range(num_calls):
        start = call_index * frames_per_call * block_size
        stop = start + frames_per_call * block_size
        if frames_per_call == 1:
            kv.allocate_chunk(adapter)
            slots = kv.chunk_write_slots(adapter)
        else:
            context = state.get_kv_caches("main", seq_len=stop - start, commit_current=True)[0].forward_ctx
            context.ensure_video_slots(torch.device("cpu"))
            slots = context.current_video_slot_mapping
        assert len(kv.block_table(adapter)) <= retained_frames + frames_per_call
        values = torch.arange(start, stop, dtype=torch.float32).view(-1, 1, 1).expand(-1, 1, 4)
        kv.key_cache(0).flatten(0, 1)[slots] = values
        kv.value_cache(0).flatten(0, 1)[slots] = -values
        if frames_per_call == 1:
            kv.commit_chunk(adapter)
        else:
            state.commit_paged_context("main")
            assert state._committed["main"] == stop
        sink_end = min(stop, sink_chunks * block_size)
        tail_start = stop if reset_at_boundary else max(sink_end, stop - 3 * block_size)
        expected = torch.cat((torch.arange(sink_end), torch.arange(tail_start, stop))).float()
        expected = expected.view(-1, 1, 1).expand(-1, 1, 4)
        resident = kv.window_block_ids(adapter)
        torch.testing.assert_close(kv.key_cache(0)[resident].flatten(0, 1), expected)
        torch.testing.assert_close(kv.value_cache(0)[resident].flatten(0, 1), -expected)
        assert adapter.completed_chunks * block_size == stop
        assert adapter.absolute_num_computed_tokens == stop
        assert len(kv.block_table(adapter)) <= retained_frames
        assert adapter.num_computed_tokens <= retained_frames * block_size
        assert kv.key_cache(0).shape == pool_shape
        assert free_total - kv.manager.block_pool.get_num_free_blocks() == min(stop // block_size, retained_frames)
    state.reset()
    fresh = state.adapter("main")
    assert fresh.compacted_tokens == fresh.absolute_num_computed_tokens == fresh.num_computed_tokens == 0
    assert kv.block_table(fresh) == []
    assert kv.manager.block_pool.get_num_free_blocks() == free_total
    kv.allocate_chunk(fresh)
    assert len(kv.chunk_write_slots(fresh)) == block_size
    state.close()
    assert kv.manager.block_pool.get_num_free_blocks() == free_total


def test_compaction_offsets_are_request_local():
    kv = make_cache()
    free_total = kv.manager.block_pool.get_num_free_blocks()
    first, second = kv.begin_request("first"), kv.begin_request("second")
    for index in range(8):
        kv.allocate_chunk(first)
        kv.commit_chunk(first)
        if index in (0, 4):
            kv.allocate_chunk(second)
            kv.commit_chunk(second)
    assert first.absolute_num_computed_tokens == 8 * 16
    assert first.compacted_tokens == 6 * 16
    assert second.absolute_num_computed_tokens == 2 * 16
    assert second.compacted_tokens == 0
    assert first.num_computed_tokens == second.num_computed_tokens == 2 * 16
    assert set(kv.block_table(first)).isdisjoint(kv.block_table(second))
    kv.end_request(first)
    kv.end_request(second)
    group = kv.manager.coordinator.single_type_managers[0]
    assert not group.req_to_blocks and not group.num_cached_block
    assert kv.manager.block_pool.get_num_free_blocks() == free_total