# SPDX-License-Identifier: Apache-2.0
"""ARDiffusionKVCache — the engine-level KV cache orchestrator for one AR-Diffusion model.

This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a
single chunk-window group) and the per-request adapter lifecycle, and exposes the
per-chunk operations a rollout needs — allocate, slot mapping, commit, window
lookup, free. It lives in the model runner (worker / GPU side), co-located with
the model and the KV tensors; the DreamZero pipeline calls these methods during a
rollout. The main-process ``ARDiffusionEngine`` only selects the engine and is otherwise
thin.
"""

from __future__ import annotations

import inspect
import os
from collections.abc import Sequence

import torch
from vllm.logger import init_logger
from vllm.v1.core.kv_cache_manager import KVCacheManager
from vllm.v1.kv_cache_interface import (
    KVCacheConfig,
    KVCacheGroupSpec,
    KVCacheSpec,
    KVCacheTensor,
)
from vllm.v1.request import RequestStatus

from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
    ChunkWindowSpec,
    allocate_kv_pool_with_views,
    chunk_slot_mapping,
    pool_write_chunk,
    resident_block_ids,
)

_log = init_logger(__name__)


class ARDiffusionRequestAdapter:
    """Duck-types the subset of ``vllm.v1.request.Request`` that the
    ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` /
    ``free`` and the coordinator they call into).

    It is intentionally NOT a full ``Request``. The conformance test exercises a
    real ``KVCacheManager`` against this adapter so the surface cannot silently
    drift across vLLM versions.

    A AR-Diffusion request advances one *chunk* at a time: ``allocate_slots`` is called
    once per chunk and ``num_computed_tokens`` advances only when a chunk is
    committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk
    reuse the same slots.
    """

    def __init__(
        self,
        request_id: str,
        *,
        chunk_size: int,
        prefill_prefix_tokens: int = 0,
    ) -> None:
        self.request_id = request_id
        self._chunk_size = chunk_size
        self._prefill = prefill_prefix_tokens
        self._completed_chunks = 0
        # Filled only when cross-request prefix reuse is enabled (Phase 3).
        self.block_hashes: list = []
        self.skip_reading_prefix_cache = True
        self.num_preemptions = 0
        # vLLM watermark gate reads this; map the request lifecycle onto it.
        self.status = RequestStatus.WAITING

    @property
    def num_computed_tokens(self) -> int:
        """Persistent KV already materialized (committed chunks + prefill)."""
        return self._prefill + self._completed_chunks * self._chunk_size

    @property
    def num_tokens(self) -> int:
        """Total tokens once the in-flight chunk is committed."""
        return self._prefill + (self._completed_chunks + 1) * self._chunk_size

    @property
    def num_prompt_tokens(self) -> int:
        """The prefill prefix length (read by ``cache_blocks`` when caching)."""
        return self._prefill

    @property
    def completed_chunks(self) -> int:
        return self._completed_chunks

    def on_chunk_committed(self) -> None:
        """Advance by one chunk. Call once per chunk, not per denoise step."""
        self._completed_chunks += 1


def compute_num_blocks(
    available_bytes: int,
    gpu_memory_fraction: float,
    page_size_bytes: int,
) -> int:
    """Number of KV blocks that fit in ``fraction`` of the memory budget."""
    if page_size_bytes <= 0:
        raise ValueError(f"page_size_bytes must be positive, got {page_size_bytes}")
    if not 0.0 < gpu_memory_fraction <= 1.0:
        raise ValueError(f"gpu_memory_fraction must be in (0, 1], got {gpu_memory_fraction}")
    budget = int(available_bytes * gpu_memory_fraction)
    return max(0, budget // page_size_bytes)


def build_kv_manager(
    spec: KVCacheSpec,
    layer_names: Sequence[str],
    num_blocks: int,
    max_model_len: int,
    *,
    enable_caching: bool = False,
) -> KVCacheManager:
    """Build a ``KVCacheManager`` with a single KV cache group for ``spec``.

    Args:
        spec: The KV cache spec for the group (e.g. a ``ChunkWindowSpec``).
        layer_names: Attention layers sharing this group's block table.
        num_blocks: Total physical blocks in the pool.
        max_model_len: Upper bound on a request's sequence length.
        enable_caching: Cross-request prefix caching (Phase 3); off in Phase 1.
    """
    layer_names = list(layer_names)
    group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
    config = KVCacheConfig(
        num_blocks=num_blocks,
        kv_cache_tensors=tensors,
        kv_cache_groups=[group],
    )
    kwargs = dict(max_model_len=max_model_len, hash_block_size=spec.block_size, enable_caching=enable_caching)
    params = inspect.signature(KVCacheManager).parameters
    if "scheduler_block_size" in params:
        kwargs["scheduler_block_size"] = spec.block_size
    if "max_num_batched_tokens" in params:
        kwargs["max_num_batched_tokens"] = max_model_len
    return KVCacheManager(config, **kwargs)


class ARDiffusionKVCache:
    """Owns the paged KV pool + per-request lifecycle for a AR-Diffusion model.

    Build once per loaded model (dimensions known); then per request:
    ``begin_request`` → per chunk (``allocate_chunk`` → ``chunk_write_slots`` →
    [model writes K/V] → ``commit_chunk``) → ``end_request``.
    """

    def __init__(
        self,
        config: ARDiffusionKVConfig,
        *,
        num_layers: int,
        num_kv_heads: int,
        head_size: int,
        dtype: torch.dtype,
        block_size: int,
        max_model_len: int,
        available_bytes: int,
        cross_attn_length: int = 0,
        cross_attn_img_length: int = 0,
        device: torch.device | None = None,
        local_branches: int = 2,
        num_frame_per_block: int = 1,
    ) -> None:
        if not config.enable:
            raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig")
        if config.window_chunks is None:
            raise ValueError("Phase 1 requires a bounded window (window_chunks)")
        if config.chunk_size <= 0:
            raise ValueError("ARDiffusionKVConfig.chunk_size must be set (> 0)")
        if local_branches not in (1, 2):
            raise ValueError(f"local_branches must be 1 or 2, got {local_branches}")

        self.config = config
        # How many CFG branches allocate from THIS rank's pool. Under
        # CFG-parallel x2 each rank executes exactly one branch (rank0 pos,
        # rank1 neg) and the other branch's lazy contexts never allocate, so
        # sizing for both would leave ~half the pool as idle capacity. A
        # single-process run (TP1 / offline) executes both branches -> 2.
        self.local_branches = local_branches
        self.num_frame_per_block = max(1, int(num_frame_per_block))
        self.block_size = block_size
        self.num_layers = num_layers
        self.num_kv_heads = num_kv_heads
        self.head_size = head_size
        self.dtype = dtype
        self.cross_attn_length = cross_attn_length
        self.device = device or torch.device("cpu")
        self._adapters: dict[str, ARDiffusionRequestAdapter] = {}

        # -- cross-attention pool (static, one-time-fill) -----------------------
        # Cross-attn KV is computed once from the text encoder and never changes.
        # Allocate a per-layer contiguous tensor for each branch: (2, text_len,
        # kv_heads, head_dim) where dim-0 = [pos, neg].  Allocated separately
        # from the self-attn block pool (static, small — ~400 MiB for
        # DreamZero I2V).  Both pools draw from the same GPU free memory but
        # the cross-attn pool is sized directly rather than via the block pool
        # budget (no eviction/chunk lifecycle needed).
        self.cross_attn_img_length = cross_attn_img_length
        self._cross_k: list[torch.Tensor] = []
        self._cross_v: list[torch.Tensor] = []
        # I2V image-token cross-attn pool. Like the text k/v, the image-token
        # k_img/v_img are session-invariant (the conditioning image doesn't change),
        # so they are cached once and read every denoise step — see
        # WanI2VCrossAttention (#4154 caches these model-side too). Empty for T2V.
        self._cross_k_img: list[torch.Tensor] = []
        self._cross_v_img: list[torch.Tensor] = []

        # Bytes the cross-attn pools consume directly (K+V, pos+neg, all layers).
        # Deducted from the self-attn paged-pool budget below so the two
        # allocations together stay within the GPU free-memory budget.
        def _cross_pool_bytes(length: int) -> int:
            return 2 * 2 * length * num_kv_heads * head_size * dtype.itemsize * num_layers

        cross_total_bytes = 0
        if device is not None and cross_attn_length > 0:
            cross_shape = (2, cross_attn_length, num_kv_heads, head_size)
            cross_bytes = _cross_pool_bytes(cross_attn_length)
            cross_total_bytes += cross_bytes
            for _ in range(num_layers):
                self._cross_k.append(torch.empty(cross_shape, dtype=dtype, device=device))
                self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device))
            _log.info(
                "AR-Diffusion cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB",
                num_layers,
                cross_attn_length,
                num_kv_heads,
                head_size,
                cross_bytes / (1024 * 1024),
            )
            if cross_attn_img_length > 0:
                img_shape = (2, cross_attn_img_length, num_kv_heads, head_size)
                cross_total_bytes += _cross_pool_bytes(cross_attn_img_length)
                for _ in range(num_layers):
                    self._cross_k_img.append(torch.empty(img_shape, dtype=dtype, device=device))
                    self._cross_v_img.append(torch.empty(img_shape, dtype=dtype, device=device))
                _log.info(
                    "AR-Diffusion cross-attn IMG pool: %d layers × %d img-tok (I2V)",
                    num_layers,
                    cross_attn_img_length,
                )

        self.spec = ChunkWindowSpec(
            block_size=block_size,
            num_kv_heads=num_kv_heads,
            head_size=head_size,
            dtype=dtype,
            sliding_window=config.window_chunks * config.chunk_size,
            chunk_size=config.chunk_size,
            window_chunks=config.window_chunks,
            sink_chunks=config.sink_chunks,
            reset_at_boundary=config.reset_at_boundary,
        )
        # Each pool block spans all layers' K/V, so size against the per-layer
        # page size times the layer count.
        # Size the self-attn pool against the memory left after the cross-attn
        # pools (allocated above), so cross + self-attn stays within the budget.
        num_blocks = compute_num_blocks(
            max(0, available_bytes - cross_total_bytes),
            config.gpu_memory_fraction,
            self.spec.page_size_bytes * num_layers,
        )
        # Floor: one forward needs the resident window plus the in-flight chunk
        # (num_frame_per_block frame-blocks) for every branch THIS rank runs,
        # with a little eviction-transient headroom. The memory-fraction
        # heuristic can under-size this once block_size grows — e.g. frame-granular
        # paging at the true frame_seqlen makes each block larger and the pool
        # fewer-blocks — so guarantee the minimum the rollout cannot run without,
        # otherwise allocate_chunk hits an exhausted pool mid-forward.
        min_blocks = self.local_branches * (config.window_chunks + self.num_frame_per_block) + 2
        if num_blocks < min_blocks:
            _log.warning(
                "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
                "floor (%d local CFG branch(es) x (window_chunks=%d + num_frame_per_block=%d) + 2 headroom)",
                num_blocks,
                min_blocks,
                self.local_branches,
                config.window_chunks,
                self.num_frame_per_block,
            )
            num_blocks = min_blocks
        layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
        self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
        self.managed_num_blocks = num_blocks
        self.num_blocks = num_blocks
        # Scratch blocks are outside KVCacheManager ownership. They hold current
        # denoise-step video KV when update_kv_cache=False and action/state KV in
        # both modes; they are reused every forward and never committed.
        # Default 4: measured high-water is num_frame_per_block video blocks (2)
        # + 1 action block per branch per forward, +1 margin.
        scratch_per_branch = int(os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "4"))
        if scratch_per_branch <= 0:
            raise ValueError("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH must be positive")
        self.scratch_blocks_per_branch = scratch_per_branch
        self.scratch_num_blocks = self.local_branches * scratch_per_branch
        self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks
        self.null_block_id = self.manager.block_pool.null_block.block_id

        # Allocate the per-layer paged K/V pools on the given device.
        self._kv_pools: list[torch.Tensor] = []
        self._k_pools: list[torch.Tensor] = []
        self._v_pools: list[torch.Tensor] = []
        if device is not None:
            self._kv_pools, self._k_pools, self._v_pools = allocate_kv_pool_with_views(
                self.num_blocks_total,
                block_size,
                num_layers,
                num_kv_heads,
                head_size,
                dtype,
                device,
            )

    # -- cross-attention pool access -------------------------------------------
    # Cross-attn KV is static once populated — write once (from text encoder),
    # read many (every denoising step). Not managed through the paged block pool.

    def write_cross_kv(
        self,
        layer_idx: int,
        is_negative: bool,
        k: torch.Tensor | None,
        v: torch.Tensor | None,
        k_img: torch.Tensor | None = None,
        v_img: torch.Tensor | None = None,
    ) -> None:
        """Write one layer's cross-attn K/V into the pool.

        ``k`` / ``v``: ``(B, text_len, tp_num_heads, head_dim)``; only batch-0
        is copied (B=1 for inference). ``None`` skips the text half (window
        restarts keep the still-valid text K/V and rewrite only the image half).
        The ``is_negative`` flag selects the correct CFG branch slot. ``k_img`` /
        ``v_img`` (I2V image tokens, ``(B, 257, ...)``) are written when the
        image pool is allocated.
        """
        branch = 1 if is_negative else 0
        if k is not None:
            self._cross_k[layer_idx][branch].copy_(k[0])
            self._cross_v[layer_idx][branch].copy_(v[0])
        if k_img is not None and self._cross_k_img:
            self._cross_k_img[layer_idx][branch].copy_(k_img[0])
            self._cross_v_img[layer_idx][branch].copy_(v_img[0])

    def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict:
        """Return a pool-backed cross-attn cache dict for one layer.

        The dict matches the ``{"is_init": True, "k": Tensor, "v": Tensor}``
        convention the cross-attention module expects — it reads from the pool
        slice rather than from the lazy-initialised model-local dict. For I2V,
        ``k_img`` / ``v_img`` are added so the image-token cache (added by #4154)
        reads from the pool too.
        """
        branch = 1 if is_negative else 0
        k = self._cross_k[layer_idx][branch].unsqueeze(0)  # (1, L, heads, dim)
        v = self._cross_v[layer_idx][branch].unsqueeze(0)
        cache = {"is_init": True, "k": k, "v": v}
        if self._cross_k_img:
            cache["k_img"] = self._cross_k_img[layer_idx][branch].unsqueeze(0)
            cache["v_img"] = self._cross_v_img[layer_idx][branch].unsqueeze(0)
        return cache

    # -- request lifecycle ---------------------------------------------------

    def begin_request(self, request_id: str, *, prefill_prefix_tokens: int = 0) -> ARDiffusionRequestAdapter:
        adapter = ARDiffusionRequestAdapter(
            request_id,
            chunk_size=self.spec.chunk_size,
            prefill_prefix_tokens=prefill_prefix_tokens,
        )
        self._adapters[request_id] = adapter
        _log.debug("AR-Diffusion begin_request: req=%s prefill=%d", request_id, prefill_prefix_tokens)
        return adapter

    def end_request(self, adapter: ARDiffusionRequestAdapter) -> None:
        _log.debug(
            "AR-Diffusion end_request: req=%s chunks=%d free=%d",
            adapter.request_id,
            adapter.completed_chunks,
            self.manager.block_pool.get_num_free_blocks(),
        )
        self.manager.free(adapter)
        self._adapters.pop(adapter.request_id, None)

    # -- per-chunk operations ------------------------------------------------

    def allocate_chunk(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
        """Allocate a chunk's blocks (evicting out-of-window blocks first).

        Returns the request's full block table (incl. null_block placeholders).
        """
        blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size)
        if blocks is None:
            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating a chunk")
        table = self.block_table(adapter)
        resident = resident_block_ids(table, self.null_block_id)
        _log.debug(
            "AR-Diffusion allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
            adapter.request_id,
            adapter.completed_chunks,
            len(table),
            len(resident),
            self.manager.block_pool.get_num_free_blocks(),
        )
        return table

    def allocate_token_slots(self, adapter: ARDiffusionRequestAdapter, num_tokens: int) -> list[int]:
        """Allocate managed blocks for an in-flight video span without committing it."""
        if num_tokens <= 0:
            raise ValueError(f"num_tokens must be positive, got {num_tokens}")
        blocks = self.manager.allocate_slots(adapter, num_new_tokens=num_tokens)
        if blocks is None:
            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating paged attention slots")
        return self.block_table(adapter)

    def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
        return list(self.manager.get_block_ids(adapter.request_id)[0])

    def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
        """Slot mapping for the in-flight chunk — the K/V write target."""
        return chunk_slot_mapping(
            self.block_table(adapter),
            adapter.num_computed_tokens,
            self.spec.chunk_size,
            self.block_size,
        )

    def scratch_block_ids(self, is_negative: bool, start: int, count: int) -> list[int]:
        """Return branch-local scratch block ids outside manager ownership."""
        if count < 0 or start < 0:
            raise ValueError(f"scratch start/count must be non-negative, got start={start}, count={count}")
        if start + count > self.scratch_blocks_per_branch:
            raise RuntimeError(
                "AR-Diffusion paged attention scratch blocks exhausted: "
                f"need [{start}, {start + count}) of {self.scratch_blocks_per_branch}. "
                "Increase AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH."
            )
        # With one local branch (CFG-parallel x2) the rank's only branch maps to
        # slot 0 whichever CFG side it is; the other branch never allocates here.
        branch_offset = self.scratch_blocks_per_branch if (is_negative and self.local_branches == 2) else 0
        base = self.managed_num_blocks + branch_offset + start
        return list(range(base, base + count))

    def key_cache(self, layer_idx: int) -> torch.Tensor:
        return self._kv_pools[layer_idx][0]

    def value_cache(self, layer_idx: int) -> torch.Tensor:
        return self._kv_pools[layer_idx][1]

    def window_block_ids(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
        """Resident (non-null) managed blocks visible to paged attention."""
        return resident_block_ids(self.block_table(adapter), self.null_block_id)

    def commit_chunk(self, adapter: ARDiffusionRequestAdapter) -> None:
        """Advance the adapter by one chunk after its K/V is written.

        This standalone primitive is used by low-level manager tests. In the
        DreamZero paged-attention path, :meth:`ARDiffusionKVState.commit_paged_context`
        advances the adapter only after the forward succeeds. Call once per
        committed chunk, not per denoise step.
        """
        _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
        adapter.on_chunk_committed()
        _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)

    # -- pool-backed K/V access --------------------------------------------

    def write_chunk_kv(
        self,
        layer_index: int,
        new_k: torch.Tensor,
        new_v: torch.Tensor,
        adapter: ARDiffusionRequestAdapter,
    ) -> None:
        """Write one layer's committed-chunk K/V into the pool."""
        slots = self.chunk_write_slots(adapter)
        _log.debug(
            "AR-Diffusion write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
            adapter.request_id,
            layer_index,
            adapter.completed_chunks,
            (tuple(new_k.shape), tuple(new_v.shape)),
            slots.device,
        )
        pool_write_chunk(
            self._k_pools[layer_index],
            self._v_pools[layer_index],
            new_k,
            new_v,
            slots,
        )