# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""A declaration must match the cache that actually gets allocated.

Two earlier versions of this file could not fail. The first asserted hand-typed
constants against each other. The second built real `transformers` caches but
compared them to specs written in the test, so no production declaration was
ever executed. What matters here is that a real declarer is called and its
numbers are checked against a real cache built the way the model builds it.
"""

import ast
import pathlib
from types import SimpleNamespace

import pytest
import torch
from transformers import Qwen2Config
from transformers.cache_utils import DynamicCache, StaticCache

from vllm_omni.model_executor.models.model_local_kv import (
    HasModelLocalKV,
    ModelLocalKVScope,
    ModelLocalKVSpec,
    RowDriver,
    collect_model_local_kv_specs,
    spec_from_hf_config,
)

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

MODELS_DIR = pathlib.Path(__file__).resolve().parents[3] / "vllm_omni" / "model_executor" / "models"


def _total_bytes(model: object, max_num_seqs: int = 1) -> int:
    """Sum declared peak bytes across a model's declarations."""
    return sum(spec.peak_bytes(max_num_seqs) for _, spec in collect_model_local_kv_specs(model))


def _cache_bytes(cache) -> int:
    total = 0
    for layer in cache.layers:
        for tensor in (layer.keys, layer.values):
            if tensor is not None:
                total += tensor.numel() * tensor.element_size()
    return total


def _qwen2_config(layers=4, kv_heads=2, heads=8, hidden=256) -> Qwen2Config:
    return Qwen2Config(
        num_hidden_layers=layers,
        num_attention_heads=heads,
        num_key_value_heads=kv_heads,
        hidden_size=hidden,
        intermediate_size=hidden * 2,
        vocab_size=1000,
    )


# --------------------------------------------------------------------------
# A real declarer, called, checked against a real cache.
# --------------------------------------------------------------------------


def test_ming_declaration_matches_the_cache_it_describes():
    """Calls MingAudioGenerator.model_local_kv_specs and checks the bytes.

    Builds the declarer without its weights -- only `_llm_config` and `_model`
    are read -- then builds the `StaticCache` that `_init_kv_cache` builds from
    the same config and compares. A wrong layer count, kv-head count, dtype or
    capacity in the production declaration fails here.
    """
    from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator

    config = _qwen2_config()
    generator = object.__new__(MingAudioGenerator)
    generator._llm_config = config
    generator._model = torch.nn.Linear(1, 1).to(torch.float16)

    (spec,) = generator.model_local_kv_specs()

    cache = StaticCache(
        config=config,
        max_batch_size=1,
        max_cache_len=MingAudioGenerator._STATIC_CACHE_LEN,
        device="cpu",
        dtype=torch.float16,
    )
    head_dim = config.hidden_size // config.num_attention_heads
    for layer_idx in range(config.num_hidden_layers):
        keys = torch.zeros(1, config.num_key_value_heads, 1, head_dim, dtype=torch.float16)
        cache.update(keys, keys.clone(), layer_idx)

    assert spec.peak_bytes(1) == _cache_bytes(cache)


def test_ming_does_not_scale_with_max_num_seqs():
    """forward() takes runtime_additional_information[0] and loops inline.

    One worker holds one of these no matter how many sequences the scheduler
    admits. A previous revision derived the multiplier from scope and reported
    `max_num_seqs` times the real figure.
    """
    from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator

    generator = object.__new__(MingAudioGenerator)
    generator._llm_config = _qwen2_config()
    generator._model = torch.nn.Linear(1, 1).to(torch.float16)

    (spec,) = generator.model_local_kv_specs()
    assert spec.peak_bytes(64) == spec.peak_bytes(1)


def test_growing_cache_declaration_matches_a_real_dynamic_cache():
    """The declared positions are the extent the tensor reaches when full."""
    layers, kv_heads, head_dim, positions = 4, 2, 32, 11
    config = _qwen2_config(layers=layers, kv_heads=kv_heads, heads=kv_heads, hidden=kv_heads * head_dim)
    spec = spec_from_hf_config(
        config,
        name="local_transformer",
        dtype=torch.bfloat16,
        physical_capacity_positions=positions,
        capacity_source="group_size + max(delay_pattern)",
        scope=ModelLocalKVScope.INVOCATION,
        rows=RowDriver.FIXED,
        rows_reason="test",
    )

    cache = DynamicCache()
    for layer_idx in range(layers):
        keys = torch.zeros(1, kv_heads, positions, head_dim, dtype=torch.bfloat16)
        cache.update(keys, torch.zeros_like(keys), layer_idx)

    assert spec.peak_bytes() == _cache_bytes(cache)


def test_multiplicity_does_not_claim_an_object_layout():
    """One allocation B rows wide costs the same as B one-row allocations.

    The protocol reports bytes and deliberately does not distinguish the two.
    An earlier revision carried a separate `allocations` count and got the
    object topology wrong in three of four declarers, which is worse than not
    claiming it: the layout belongs in `allocation_note`, where it cannot be
    mistaken for arithmetic.
    """
    config = _qwen2_config(layers=4, kv_heads=2, heads=2, hidden=64)
    common = dict(
        config=config,
        name="c",
        dtype=torch.bfloat16,
        physical_capacity_positions=11,
        capacity_source="test",
        scope=ModelLocalKVScope.INVOCATION,
    )
    batched = spec_from_hf_config(rows=RowDriver.MAX_NUM_SEQS, **common)

    assert batched.row_count(64) == 64
    assert batched.peak_bytes(64) == batched.bytes_per_row * 64
    assert not hasattr(batched, "allocations"), "the object-count axis was removed on purpose"


def test_a_fixed_driver_ignores_concurrency():
    """A serialized or capture-pinned cache does not widen with max_num_seqs."""
    spec = spec_from_hf_config(
        _qwen2_config(),
        name="graph_pool",
        dtype=torch.bfloat16,
        physical_capacity_positions=11,
        capacity_source="test",
        scope=ModelLocalKVScope.MODEL,
        rows=RowDriver.FIXED,
        rows_fixed=261,
        rows_reason="sum of captured buckets",
    )
    assert spec.row_count(1) == 261
    assert spec.row_count(64) == 261
    assert spec.peak_bytes(64) == spec.bytes_per_row * 261


def test_qwen3_tts_declares_nothing_on_the_stateless_path():
    """Calls the real decoder declarer with no cache-holding mode configured.

    `_forward_exact` runs `pre_transformer` without `use_cache`, so stateless
    decoding allocates no KV. A revision that declared a per-request cache
    unconditionally invented over a gigabyte at max_num_seqs=256, and the
    revision before it declared nothing at all under enforce_eager, where the
    eager async-chunk path does allocate. Both are regressions this pins.
    """
    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import (
        modeling_qwen3_tts_tokenizer_v2 as mod,
    )

    decoder = object.__new__(mod.Qwen3TTSTokenizerV2Decoder)
    # nn.Module.__setattr__ needs the machinery __init__ would have built, and
    # this instance deliberately has none of it.
    setattr_ = object.__setattr__
    setattr_(
        decoder,
        "config",
        SimpleNamespace(
            sliding_window=72,
            num_hidden_layers=8,
            num_key_value_heads=16,
            head_dim=64,
            hidden_size=1024,
            num_attention_heads=16,
        ),
    )
    setattr_(decoder, "parameters", lambda: iter([torch.zeros(1, dtype=torch.float32)]))

    setattr_(decoder, "_cudagraph_wrapper", None)
    assert decoder.model_local_kv_specs() == []

    setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=False, model_local_kv_specs=lambda: []))
    assert decoder.model_local_kv_specs() == []

    setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=True, model_local_kv_specs=lambda: []))
    names = [s.name for s in decoder.model_local_kv_specs()]
    assert names == ["codec_decoder_request", "codec_decoder_working_copy"], names


def test_fixed_rejects_unreviewable_values():
    """A bare or unexplained constant cannot be reviewed, so it is rejected."""
    common = dict(
        name="x",
        layers=1,
        kv_heads=1,
        head_dim=1,
        dtype=torch.float16,
        physical_capacity_positions=1,
        capacity_source="",
        scope=ModelLocalKVScope.REQUEST,
    )
    with pytest.raises(ValueError, match="rows_reason"):
        ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=1, rows_reason="   ", **common)
    with pytest.raises(ValueError, match=">= 1"):
        ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=0, rows_reason="test", **common)


def test_rows_is_required():
    """Defaulting the batch extent is how the previous guarantee died.

    The dataclass declared it required while the helper every declarer uses
    defaulted it to 1, so nothing enforced it where it mattered.
    """
    with pytest.raises(TypeError):
        spec_from_hf_config(  # type: ignore[call-arg]
            _qwen2_config(),
            name="x",
            dtype=torch.float16,
            physical_capacity_positions=1,
            capacity_source="",
            scope=ModelLocalKVScope.REQUEST,
        )


def test_spec_from_hf_config_raises_rather_than_guessing():
    class Empty:
        pass

    with pytest.raises(ValueError, match="layer count"):
        spec_from_hf_config(
            Empty(),
            name="x",
            dtype=torch.float16,
            physical_capacity_positions=1,
            capacity_source="",
            scope=ModelLocalKVScope.REQUEST,
            rows=RowDriver.FIXED,
            rows_reason="test",
        )


# --------------------------------------------------------------------------
# Collection
# --------------------------------------------------------------------------


def _spec(**overrides) -> ModelLocalKVSpec:
    kwargs = dict(
        name="c",
        layers=2,
        kv_heads=2,
        head_dim=4,
        dtype=torch.float16,
        physical_capacity_positions=8,
        capacity_source="test",
        scope=ModelLocalKVScope.REQUEST,
        rows=RowDriver.FIXED,
        rows_reason="test",
    )
    kwargs.update(overrides)
    return ModelLocalKVSpec(**kwargs)


class _Owner(torch.nn.Module):
    def model_local_kv_specs(self):
        return [_spec()]


class _CUDAGraphWrapperLike:
    """Mirrors vllm.compilation.cuda_graph.CUDAGraphWrapper.

    Not an nn.Module, and its __getattr__ forwards everything the runnable
    has. That forwarding is why a previous revision's "unwrap fix" was
    unnecessary, and why a walk that does not unwrap can enter the tree twice.
    """

    def __init__(self, runnable):
        self.runnable = runnable

    def unwrap(self):
        return self.runnable

    def __getattr__(self, key):
        return getattr(self.runnable, key)


def test_collector_finds_owners_nested_in_the_module_tree():
    class Mid(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.inner = _Owner()

    class Top(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.mid = Mid()

    model = Top()
    assert [path for path, _ in collect_model_local_kv_specs(model)] == ["mid.inner"]
    # 2 layers * 2 kv * 4 head_dim * 8 positions * 2 bytes * 2 (K and V) = 512
    assert _total_bytes(model) == 512


def test_a_root_declarer_behind_a_wrapper_is_counted_once():
    """The wrapper forwards attributes, so a naive walk sees it twice."""

    class Root(_Owner):
        pass

    wrapped = _CUDAGraphWrapperLike(Root())
    collected = collect_model_local_kv_specs(wrapped)

    assert len(collected) == 1, collected
    assert _total_bytes(wrapped) == 512


def test_collector_reaches_owners_through_a_wrapper():
    class Top(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.owner = _Owner()

    wrapped = _CUDAGraphWrapperLike(Top())
    assert [path for path, _ in collect_model_local_kv_specs(wrapped)] == ["owner"]


def test_collector_survives_an_owner_that_raises():
    class Broken(torch.nn.Module):
        def model_local_kv_specs(self):
            raise RuntimeError("checkpoint field missing")

    class Top(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.broken = Broken()
            self.ok = _Owner()

    assert _total_bytes(Top()) == 512


def test_undeclared_models_report_zero_not_an_error():
    assert _total_bytes(object()) == 0
    assert _total_bytes(torch.nn.Linear(2, 2)) == 0


def test_protocol_is_structural():
    assert isinstance(_Owner(), HasModelLocalKV)
    assert not isinstance(torch.nn.Linear(2, 2), HasModelLocalKV)


@pytest.mark.parametrize(
    ("relative_path", "owner"),
    [
        ("qwen3_tts/segmented_graph_wrapper.py", "CUDAGraphDecoderWrapper"),
        ("qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py", "Qwen3TTSTokenizerV2Decoder"),
        ("mimo_audio/mimo_audio_llm.py", "MiMoAudioLLMForConditionalGeneration"),
        ("minicpmo_4_5/minicpmo_4_5_omni_llm.py", "MiniCPMWhisperEncoder"),
        ("ming_flash_omni/talker_module.py", "MingAudioGenerator"),
        ("ming_flash_omni/ming_flash_omni_talker.py", "MingFlashOmniTalkerForConditionalGeneration"),
    ],
)
def test_known_cache_owners_still_declare(relative_path: str, owner: str):
    """Parsed rather than imported: these modules need CUDA and real weights.

    Weak by construction -- it only proves the method exists. The ming tests
    above are the ones that check a declaration's contents.
    """
    tree = ast.parse((MODELS_DIR / relative_path).read_text())
    classes = [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef) and node.name == owner]
    assert classes, f"{owner} not found in {relative_path}"
    methods = {child.name for cls in classes for child in cls.body if isinstance(child, ast.FunctionDef)}
    assert "model_local_kv_specs" in methods, f"{owner} no longer declares its model-local KV"