# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU contracts for capability-driven AR-Diffusion sessions."""

from __future__ import annotations

from collections import OrderedDict
from contextlib import contextmanager
from types import SimpleNamespace

import pytest
import torch

from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.experimental.ar_diffusion.capability import (
    ARDiffusionCrossAttentionKVSpec,
    ARDiffusionKVBranchSpec,
    ARDiffusionKVCacheSpec,
)
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest

BLOCK = 16
POS = "positive"
NEG = "negative"

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


@pytest.fixture(autouse=True)
def platform_synchronize_calls(monkeypatch) -> list[None]:
    """Stub the platform barrier the runner takes at every chunk boundary.

    These tests run on a CPU host, where the resolved platform is
    ``UnspecifiedOmniPlatform`` and ``synchronize()`` raises. Tests that need
    a failing barrier override this with their own fake.
    """
    calls: list[None] = []
    monkeypatch.setattr(
        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
        SimpleNamespace(synchronize=lambda: calls.append(None)),
    )
    return calls


def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
    """Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV."""
    return ARDiffusionKVCacheSpec(
        num_layers=2,
        num_kv_heads=4,  # already TP-local
        head_size=64,
        tokens_per_frame=BLOCK,
        frames_per_block=3,
        window_frames=5,
        sink_frames=1,
        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
        session_capacity=capacity,
        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 8),),
    )


def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
    return ARDiffusionKVCacheSpec(
        num_layers=2,
        num_kv_heads=4,
        head_size=64,
        tokens_per_frame=BLOCK,
        frames_per_block=4,
        window_frames=6,
        kv_branches=(ARDiffusionKVBranchSpec(POS, 0), ARDiffusionKVBranchSpec(NEG, 1)),
        session_capacity=capacity,
        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 8),),
    )


def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
    return ARDiffusionKVCacheSpec(
        num_layers=1,
        num_kv_heads=1,
        head_size=1,
        tokens_per_frame=1,
        frames_per_block=3,
        window_frames=3,
        sink_frames=3,
        kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
        session_capacity=capacity,
        cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
    )


class CapablePipeline:
    def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
        self.spec = spec
        self.bound_state = None
        self.binds: list[str] = []
        self.resets: list[str] = []
        self.closes: list[str] = []

    def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec:
        return self.spec

    @contextmanager
    def bind_ar_diffusion_state(self, session_id: str, state):
        assert self.bound_state is None
        self.bound_state = state
        self.binds.append(session_id)
        try:
            yield
        finally:
            self.bound_state = None

    def reset_ar_diffusion_session(self, session_id: str) -> None:
        self.resets.append(session_id)

    def close_ar_diffusion_session(self, session_id: str) -> None:
        self.closes.append(session_id)


class WarmupPipeline(CapablePipeline):
    def __init__(self, spec: ARDiffusionKVCacheSpec, requests: list[object]) -> None:
        super().__init__(spec)
        self.requests = requests

    def ar_diffusion_warmup_requests(self, session_id: str):
        assert session_id == ARDiffusionModelRunner._WARMUP_SID
        return iter(self.requests)


class BatchCapablePipeline(CapablePipeline):
    supports_request_batch = True


class StepCapablePipeline(CapablePipeline):
    supports_step_execution = True

    def prepare_encode(self, state, **kwargs):
        del kwargs
        return state

    def denoise_step(self, input_batch, *, states=None, **kwargs):
        del input_batch, states, kwargs
        return None

    def step_scheduler(self, state, noise_pred, **kwargs):
        del state, noise_pred, kwargs

    def post_decode(self, state, **kwargs):
        del state, kwargs
        from vllm_omni.diffusion.data import DiffusionOutput

        return DiffusionOutput()


def scheduler_output(
    request_ids: tuple[str, ...] = ("req-1",),
    *,
    finished: tuple[str, ...] = (),
) -> DiffusionSchedulerOutput:
    """Build the real scheduler payload; stepwise requests arrive as cached rows."""
    return DiffusionSchedulerOutput(
        step_id=0,
        scheduled_new_reqs=[],
        scheduled_cached_reqs=CachedRequestData(request_ids=list(request_ids)),
        finished_req_ids=set(finished),
        num_running_reqs=len(request_ids),
        num_waiting_reqs=0,
    )


def make_runner(
    pipeline: object,
    *,
    available_bytes: int = 1 << 28,
    step_execution: bool = False,
    gpu_memory_fraction: float = 0.1,
) -> ARDiffusionModelRunner:
    runner = object.__new__(ARDiffusionModelRunner)
    runner.od_config = SimpleNamespace(
        max_num_seqs=1,
        dtype=torch.float32,
        enforce_eager=True,
        step_execution=step_execution,
    )
    runner.device = torch.device("cpu")
    runner.pipeline = pipeline
    runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
        enable=True,
        gpu_memory_fraction=gpu_memory_fraction,
    )
    runner.kv_cache = None
    runner._ar_diffusion_capability = None
    runner._ar_diffusion_kv_cache_spec = None
    runner._sessions = OrderedDict()
    runner._session_capacity = 0
    runner._perf_e2e_times = []
    runner._stepwise_chunk_started = {}
    runner._preallocate_kv_cache(available_bytes=available_bytes)
    return runner


def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
    state = runner._get_or_create_session(session_id)
    ctx = state.get_kv_caches(kv_branch, seq_len=BLOCK, commit_current=True)[0].forward_ctx
    ctx.ensure_video_slots(torch.device("cpu"))
    state.commit_paged_context(kv_branch)
    return state


def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
    state = runner._get_or_create_session(session_id)
    assert runner.kv_cache is not None
    seq_len = BLOCK * runner.kv_cache.frames_per_block
    ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
    ctx.ensure_video_slots(torch.device("cpu"))
    state.commit_paged_context(kv_branch)
    return state


def test_ar_runner_rejects_pipeline_without_capability():
    runner = object.__new__(ARDiffusionModelRunner)
    runner.od_config = SimpleNamespace(max_num_seqs=1)
    runner.pipeline = object()
    with pytest.raises(TypeError, match="SupportsARDiffusionPipeline"):
        runner._preallocate_kv_cache(available_bytes=1 << 20)


def test_runner_uses_typed_tick_as_authoritative_session_contract():
    tick = ARDiffusionTickRequest(
        session_id="world-7",
        request_id="request-3",
        chunk_index=3,
        reset=True,
    )
    req = SimpleNamespace(
        request_id="request-3",
        sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
    )

    session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)

    assert session_id == "world-7"
    assert extra_args == tick.to_extra_args()
    assert parsed == tick


def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
    tick = ARDiffusionTickRequest(
        session_id="world-7",
        request_id="client-request-3",
        chunk_index=3,
    )
    req = SimpleNamespace(
        request_id="engine-request-uuid",
        sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
    )

    session_id, _, parsed = ARDiffusionModelRunner._request_session(req)

    assert session_id == "world-7"
    assert req.request_id == "engine-request-uuid"
    assert parsed.request_id == "client-request-3"


def test_lingbot_like_single_branch_session_reuse_reset_and_close():
    pipeline = CapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline)
    kv = runner.kv_cache
    assert kv is not None
    assert kv.num_local_kv_branches == 1
    assert kv.frames_per_block == 3
    assert kv.spec.window_chunks == 5
    assert kv.spec.sink_chunks == 1
    assert kv.cross_attention_lengths == {"text": 8}

    first = commit_one_frame(runner, "s1", "main")
    assert runner._get_or_create_session("s1") is first
    k = torch.randn(1, 8, 4, 64)
    v = torch.randn(1, 8, 4, 64)
    first.populate_cross_attention("main", "text", [(k, v)] * first.num_layers)
    assert first.get_cross_attention_kv("main", "text")[0]["k"].shape == k.shape

    runner.reset_session("s1")
    assert "s1" not in runner._sessions
    assert "s1" not in kv._cross_sessions
    assert pipeline.resets == ["s1"]
    second = runner._get_or_create_session("s1")
    assert second is not first

    runner.close_session("s1")
    assert "s1" not in runner._sessions
    assert pipeline.closes == ["s1"]


def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
    runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
    kv = runner.kv_cache
    assert kv is not None

    session_a = commit_one_block(runner, "world-a", "main")
    a0_blocks = kv.window_block_ids(session_a.adapter("main"))
    session_b = commit_one_block(runner, "world-b", "main")
    b0_blocks = kv.window_block_ids(session_b.adapter("main"))
    session_a_again = commit_one_block(runner, "world-a", "main")
    a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))

    assert session_a_again is session_a
    assert session_a.adapter("main").request_id == "ar::world-a::main"
    assert session_b.adapter("main").request_id == "ar::world-b::main"
    assert session_a.adapter("main").completed_chunks == 6
    assert session_b.adapter("main").completed_chunks == 3
    assert len(a0_blocks) == 3
    assert len(b0_blocks) == 3
    assert len(a1_blocks) == 6
    assert set(a0_blocks) <= set(a1_blocks)
    assert set(a1_blocks).isdisjoint(b0_blocks)
    assert tuple(runner._sessions) == ("world-b", "world-a")


def test_lingbot_like_sink_survives_sliding_window_eviction():
    runner = make_runner(CapablePipeline(lingbot_like_spec()))
    kv = runner.kv_cache
    assert kv is not None

    state = commit_one_frame(runner, "s1", "main")
    sink_block = kv.block_table(state.adapter("main"))[0]
    for _ in range(7):
        state = commit_one_frame(runner, "s1", "main")

    adapter = state.adapter("main")
    table = kv.block_table(adapter)
    assert len(table) == 6  # one sink + five recent frames
    assert table[0] == sink_block
    assert kv.null_block_id not in table
    assert adapter.absolute_num_computed_tokens == 8 * BLOCK
    assert adapter.num_computed_tokens == 6 * BLOCK

    ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=False)[0].forward_ctx
    visible, _ = ctx.video_block_table(torch.device("cpu"))
    assert visible[0] == table[0]
    assert len(visible) == 6  # sink + recent window, including current scratch


def test_dreamzero_like_two_branches_are_independent():
    runner = make_runner(CapablePipeline(dreamzero_like_spec()))
    kv = runner.kv_cache
    assert kv is not None and kv.num_local_kv_branches == 2
    state = commit_one_frame(runner, "s1", POS)
    assert len(kv.window_block_ids(state.adapter(POS))) == 1
    assert kv.window_block_ids(state.adapter(NEG)) == []


def test_lru_eviction_releases_blocks_and_notifies_pipeline():
    pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
    runner = make_runner(pipeline)
    kv = runner.kv_cache
    assert kv is not None
    assert runner._session_capacity == 2
    assert kv.session_capacity == 2
    free_total = kv.manager.block_pool.get_num_free_blocks()
    old = commit_one_frame(runner, "old", "main")
    k = torch.randn(1, 8, 4, 64)
    old.populate_cross_attention("main", "text", [(k, k)] * old.num_layers)
    assert kv.manager.block_pool.get_num_free_blocks() < free_total

    runner._get_or_create_session("new")
    assert tuple(runner._sessions) == ("old", "new")
    runner._get_or_create_session("newest")

    assert tuple(runner._sessions) == ("new", "newest")
    assert pipeline.closes == ["old"]
    assert "old" not in kv._cross_sessions
    assert kv.manager.block_pool.get_num_free_blocks() == free_total


def test_budget_reduced_capacity_drives_runner_lru():
    pipeline = CapablePipeline(tiny_spec(capacity=2))
    # One tiny LingBot-like session requires 1,808 bytes; two require 2,592.
    #
    # These read 128 and 192 while the paging unit was the frame. tiny_spec
    # declares one token per frame, so it used to page a single token at a
    # time -- a block size no attention kernel accepts. Paging at 16 makes
    # each page sixteen times larger. The block *counts* are unchanged, and
    # so is what this test is about: a budget that fits fewer sessions than
    # the pipeline asked for must drive the runner's LRU.
    runner = make_runner(
        pipeline,
        available_bytes=2048,
        gpu_memory_fraction=1.0,
    )
    kv = runner.kv_cache
    assert kv is not None
    assert kv.requested_session_capacity == 2
    assert kv.session_capacity == 1
    assert runner._session_capacity == 1

    runner._get_or_create_session("old")
    runner._get_or_create_session("new")

    assert tuple(runner._sessions) == ("new",)
    assert pipeline.closes == ["old"]


def test_dreamzero_like_requested_capacity_is_capped_by_budget():
    spec = dreamzero_like_spec(capacity=64)
    # Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
    # managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
    # cross-attention=1 page/session, for 44 pages total.
    page_bytes = 65_536
    pipeline = CapablePipeline(spec)
    runner = make_runner(
        pipeline,
        available_bytes=44 * page_bytes,
        gpu_memory_fraction=1.0,
    )
    kv = runner.kv_cache
    assert kv is not None
    assert kv.requested_session_capacity == 64
    assert kv.session_capacity == 2
    assert runner._session_capacity == 2
    assert kv.cross_attention_reserved_bytes == 2 * page_bytes


def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
    pipeline = CapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline)
    kv = runner.kv_cache
    assert kv is not None
    free_total = kv.manager.block_pool.get_num_free_blocks()

    def boom(self, req, kv_prefetch_job=None):
        state = pipeline.bound_state
        ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=True)[0].forward_ctx
        ctx.ensure_video_slots(torch.device("cpu"))
        raise RuntimeError("layer exploded")

    monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
    request = SimpleNamespace(
        request_id="broken-request",
        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
    )

    with pytest.raises(RuntimeError, match="layer exploded"):
        runner.execute_model(request)

    assert pipeline.bound_state is None
    assert pipeline.closes == ["broken"]
    assert not runner._sessions
    assert not kv._adapters
    assert kv.manager.block_pool.get_num_free_blocks() == free_total


def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
    pipeline = CapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline)
    kv = runner.kv_cache
    assert kv is not None
    free_total = kv.manager.block_pool.get_num_free_blocks()

    def return_after_allocation(self, req, kv_prefetch_job=None):
        state = pipeline.bound_state
        ctx = state.get_kv_caches("main", seq_len=BLOCK, commit_current=True)[0].forward_ctx
        ctx.ensure_video_slots(torch.device("cpu"))
        return object()

    def synchronize_boom():
        raise RuntimeError("asynchronous kernel failed")

    monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
    # The runner defers the device barrier to the platform, so the failure is
    # injected there rather than on a device-specific torch entry point.
    monkeypatch.setattr(
        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
        SimpleNamespace(synchronize=synchronize_boom),
    )
    request = SimpleNamespace(
        request_id="broken-request",
        sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
    )

    with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
        runner.execute_model(request)

    assert pipeline.bound_state is None
    assert pipeline.closes == ["broken"]
    assert not runner._sessions
    assert not kv._adapters
    assert not runner._perf_e2e_times
    assert kv.manager.block_pool.get_num_free_blocks() == free_total


def test_ar_runner_rejects_step_and_request_batch_modes():
    with pytest.raises(ValueError, match="step_execution=True"):
        make_runner(CapablePipeline(lingbot_like_spec()), step_execution=True)
    with pytest.raises(ValueError, match="request-batch execution"):
        make_runner(BatchCapablePipeline(lingbot_like_spec()))


def test_ar_runner_allows_step_execution_when_pipeline_implements_the_contract():
    runner = make_runner(StepCapablePipeline(lingbot_like_spec()), step_execution=True)
    assert runner.kv_cache is not None


def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
    runner = object.__new__(ARDiffusionModelRunner)
    with pytest.raises(RuntimeError, match="request-batch execution"):
        runner.execute_model_batch(None, None)
    with pytest.raises(RuntimeError, match="step execution"):
        runner.execute_stepwise(None)


def test_execute_stepwise_binds_request_id_session_and_unbinds_after_call(monkeypatch):
    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

    pipeline = StepCapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline, step_execution=True)
    bound_during = []

    def fake_execute(self, scheduler_output):
        del self, scheduler_output
        bound_during.append(pipeline.bound_state is not None)
        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False)])

    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
    output = runner.execute_stepwise(scheduler_output())

    assert bound_during == [True]
    assert pipeline.bound_state is None
    assert pipeline.binds == ["req-1"]
    assert "req-1" in runner._sessions
    assert output.get_request_output("req-1") is not None
    assert not pipeline.closes


def test_execute_stepwise_closes_session_when_request_finishes(monkeypatch):
    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

    pipeline = StepCapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline, step_execution=True)

    def fake_execute(self, scheduler_output):
        del self, scheduler_output
        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=True)])

    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
    runner.execute_stepwise(scheduler_output())

    assert pipeline.bound_state is None
    assert pipeline.closes == ["req-1"]
    assert "req-1" not in runner._sessions


def test_execute_stepwise_exception_releases_kv_and_does_not_resume_pages(monkeypatch):
    pipeline = StepCapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline, step_execution=True)
    first = runner._get_or_create_session("req-1")

    def boom(self, scheduler_output):
        del self, scheduler_output
        raise RuntimeError("step failed")

    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", boom)
    with pytest.raises(RuntimeError, match="step failed"):
        runner.execute_stepwise(scheduler_output())

    assert pipeline.bound_state is None
    assert pipeline.closes == ["req-1"]
    assert "req-1" not in runner._sessions
    second = runner._get_or_create_session("req-1")
    assert second is not first


def test_execute_stepwise_closes_session_for_scheduler_aborted_request(monkeypatch):
    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput

    pipeline = StepCapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline, step_execution=True)
    runner._get_or_create_session("req-1")

    def fake_execute(self, scheduler_output):
        del self, scheduler_output
        return BatchRunnerOutput.from_list([])

    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
    # An aborted request is never scheduled again and never reports finished,
    # so the retire path is the only thing that can free its KV.
    runner.execute_stepwise(scheduler_output(request_ids=(), finished=("req-1",)))

    assert "req-1" not in runner._sessions
    assert pipeline.closes == ["req-1"]


def test_execute_stepwise_times_once_per_chunk_not_once_per_step(monkeypatch, platform_synchronize_calls):
    from vllm_omni.diffusion.data import DiffusionOutput
    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

    pipeline = StepCapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline, step_execution=True)
    results = [None, None, None, DiffusionOutput()]

    def fake_execute(self, scheduler_output):
        del self, scheduler_output
        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=results.pop(0))])

    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
    for _ in range(4):
        runner.execute_stepwise(scheduler_output())

    # Four denoise steps, one emitted AR block: one timing sample, matching the
    # one-sample-per-block meaning request mode already has.
    assert len(runner._perf_e2e_times) == 1
    assert len(platform_synchronize_calls) == 1
    assert "req-1" not in runner._stepwise_chunk_started


def test_execute_stepwise_synchronize_exception_fails_closed(monkeypatch):
    """The chunk barrier is where an asynchronous device error surfaces. It must
    clean up like an in-step failure instead of leaving the session, the cached
    step state and the timing entry to the scheduler's next cycle."""
    from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

    pipeline = StepCapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline, step_execution=True)
    runner.state_cache = {"req-1": object()}

    def fake_execute(self, scheduler_output):
        del self, scheduler_output
        return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=object())])

    def synchronize_boom():
        raise RuntimeError("asynchronous kernel failed")

    monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
    monkeypatch.setattr(
        "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
        SimpleNamespace(synchronize=synchronize_boom),
    )

    with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
        runner.execute_stepwise(scheduler_output())

    assert pipeline.bound_state is None
    assert pipeline.closes == ["req-1"]
    assert "req-1" not in runner._sessions
    assert "req-1" not in runner.state_cache
    assert "req-1" not in runner._stepwise_chunk_started
    assert not runner._perf_e2e_times


def test_model_specific_warmup_provider_is_consumed(monkeypatch):
    requests = [object(), object()]
    pipeline = WarmupPipeline(lingbot_like_spec(), requests)
    runner = make_runner(pipeline)
    seen: list[object] = []
    monkeypatch.setattr(runner, "execute_model", seen.append)

    runner._warmup_ar_rollout()

    assert seen == requests
    assert pipeline.closes == [runner._WARMUP_SID]


def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch):
    pipeline = CapablePipeline(lingbot_like_spec())
    runner = make_runner(pipeline)
    execute = SimpleNamespace(called=False)

    def fail_if_called(request):
        execute.called = True

    monkeypatch.setattr(runner, "execute_model", fail_if_called)
    runner._warmup_ar_rollout()
    assert execute.called is False


# ── paging unit vs eviction unit ────────────────────────────────────────────


def test_a_frame_that_is_already_a_legal_block_keeps_its_paging():
    """Every resolution that runs today must page exactly as it does today."""
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    for tokens_per_frame in (16, 320, 1440, 4096):
        assert tokens_per_frame % 16 == 0
        assert paging_block_size(tokens_per_frame) == tokens_per_frame


def test_a_frame_the_kernel_would_reject_pages_at_the_finest_legal_unit():
    """832x480 is the checkpoint's own default and the kernel rejects it.

    (480/16) x (832/16) = 1560 tokens per frame, and 1560 % 16 == 8, so a
    frame-sized block is not something FlashAttention's paged kernel accepts.
    """
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    assert 1560 % 16 == 8
    assert paging_block_size(1560) == 16
    # And whatever it returns is always something the kernel takes.
    for tokens_per_frame in range(1, 200):
        assert paging_block_size(tokens_per_frame) % 16 == 0


# ── the kernel's constraint is data, not a constant ─────────────────────────


def _multiple_of(base):
    from vllm.v1.attention.backend import MultipleOf

    return MultipleOf(base)


def test_the_dispatching_module_answers_in_vllms_own_vocabulary():
    """AR-Diffusion calls flash_attn itself, so no backend can be asked.

    The answer has to come from the module that picks the kernel, and it has to
    have the shape vLLM uses, so the policy that consumes it does not care
    where it came from.
    """
    from vllm.v1.attention.backend import MultipleOf

    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes

    advertised = supported_kernel_block_sizes()
    assert advertised, "a kernel that accepts nothing cannot be paged for"
    for entry in advertised:
        assert isinstance(entry, int | MultipleOf)


def test_a_kernel_that_takes_one_fixed_size_gets_that_size():
    """hpc_attn advertises [64] -- not a multiple, a single legal value.

    A frame of 1560 tokens is not 64, and 1560 % 64 == 24, so the frame cannot
    be the block and the only thing left to page at is 64. Assuming multiples
    of 16 here would hand the kernel a block size it rejects.
    """
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    assert 1560 % 64 == 24
    assert paging_block_size(1560, [64]) == 64
    assert paging_block_size(64, [64]) == 64
    assert paging_block_size(128, [64]) == 64  # a legal multiple is still not the advertised size


def test_a_kernel_with_large_pages_still_keeps_a_whole_frame_when_it_fits():
    """FlashInfer offers 128-and-up pages on Blackwell and not on Hopper.

    The same checkpoint on the two cards must therefore page differently, which
    is exactly what a constant cannot express.
    """
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    hopper = [_multiple_of(16)]
    blackwell = [_multiple_of(16), 128, 256]
    assert paging_block_size(1440, hopper) == 1440  # a legal multiple of 16
    assert paging_block_size(1440, blackwell) == 1440
    assert paging_block_size(1560, hopper) == 16
    assert paging_block_size(1560, blackwell) == 16  # finest legal unit, not the largest


def test_whatever_comes_back_is_always_something_the_kernel_accepts():
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    for advertised, check in (
        ([_multiple_of(16)], lambda n: n % 16 == 0),
        ([64], lambda n: n == 64),
        ([_multiple_of(32), 48], lambda n: n % 32 == 0 or n == 48),
    ):
        for tokens_per_frame in range(1, 300):
            assert check(paging_block_size(tokens_per_frame, advertised))


def test_a_kernel_that_advertises_nothing_is_an_error_not_a_guess():
    from vllm_omni.experimental.ar_diffusion.runner import paging_block_size

    with pytest.raises(ValueError, match="no legal block size"):
        paging_block_size(1560, [])