# SPDX-License-Identifier: Apache-2.0
"""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.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]


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


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._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

    for _ in range(8):
        state = commit_one_frame(runner, "s1", "main")

    table = kv.block_table(state.adapter("main"))
    assert table[0] != kv.null_block_id
    assert table[1] == kv.null_block_id

    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 128 bytes; two require 192.
    runner = make_runner(
        pipeline,
        available_bytes=128,
        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(device):
        raise RuntimeError("asynchronous kernel failed")

    monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
    runner.device = torch.device("cuda")
    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_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_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