# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CPU coverage for Gepard's prompt layout, streaming-decode window and config.

The e2e test needs a GPU, NeMo and the checkpoint, so it only runs nightly.
These pin the parts that are pure logic and break silently: which frames each
decode covers, which samples survive the lookback trim, the prompt layout, and
that the packaged deploy config sits where the example looks for it.

The codec is replaced by a shift-invariant stand-in whose output depends only
on the codes, which makes "chunked emissions concatenate to one whole-history
decode" an exact assertion — the contract lookback exists to provide.
"""

from __future__ import annotations

from pathlib import Path

import pytest
import torch
import torch.nn as nn

import vllm_omni
from vllm_omni.config.stage_config import load_deploy_config
from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
from vllm_omni.model_executor.models.gepard.gepard_talker import (
    CHUNK_FRAMES,
    FIRST_CHUNK_FRAMES,
    GepardTalkerForConditionalGeneration,
    _GepardState,
)
from vllm_omni.model_executor.models.gepard.nanocodec import (
    LOOKBACK_FRAMES,
    NanoCodec,
    apply_end_of_speech_tail,
)
from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids

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

SAMPLES_PER_FRAME = 1024
NUM_HEADS = 32
SAMPLE_RATE = 22050
# Text-token count at or above which the layout is left unrepeated, read from
# the checkpoint defaults so the boundary cases follow the config.
APPLY_BELOW = GepardConfig().text_repetition_apply_below


class _ShiftInvariantCodec:
    """Frame tagged ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits."""

    def __init__(self) -> None:
        self.windows: list[int] = []  # frame count of every decode call

    def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor):
        assert codes.shape[0] == 1, "streaming decode is per-request"
        assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels"
        tags = codes[0, 0, :].to(torch.int64)  # (T,)
        self.windows.append(int(tags.shape[0]))
        offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
        audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1)
        return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME


def _make_codec() -> tuple[NanoCodec, _ShiftInvariantCodec]:
    codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE)
    fake = _ShiftInvariantCodec()
    codec._codec = fake  # bypasses load(); NeMo is never imported
    assert codec.is_loaded
    return codec, fake


def _make_talker(codec: NanoCodec) -> GepardTalkerForConditionalGeneration:
    """A real talker with only the audio-path attributes populated.

    ``__init__`` needs a VllmConfig and builds the backbone; the methods under
    test touch only these four attributes.
    """
    talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration)
    nn.Module.__init__(talker)  # sets up _parameters/_buffers/_modules only
    talker._codec = codec
    talker._audio_queue = []
    talker._active_states = {}
    talker._deferred_cleanup_ids = set()
    return talker


def _frame(tag: int) -> torch.Tensor:
    """One committed frame: 32 per-dimension codes, all carrying the tag."""
    return torch.full((NUM_HEADS,), tag, dtype=torch.long)


def _expected_samples(num_frames: int) -> torch.Tensor:
    offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
    tags = torch.arange(num_frames, dtype=torch.int64).unsqueeze(1)
    return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32)


def _drain(talker) -> torch.Tensor:
    """Concatenate and clear the queued deltas, as make_omni_output does."""
    if not talker._audio_queue:
        return torch.empty(0)
    out = torch.cat([a.reshape(-1) for _, a in talker._audio_queue])
    talker._audio_queue.clear()
    return out


def _generate(talker, state, num_frames: int, *, stop_at_end: bool) -> None:
    """Drive the commit loop the way forward() does, one frame per step."""
    for tag in range(num_frames):
        state.frames.append(_frame(tag))
        state.frame_count += 1
        state.past_first_step = True
        talker._emit_audio(state, is_final=stop_at_end and tag == num_frames - 1)


# Counts straddling both cadence thresholds, where an off-by-one hides.
@pytest.mark.parametrize(
    "num_frames",
    [
        1,
        2,
        FIRST_CHUNK_FRAMES - 1,
        FIRST_CHUNK_FRAMES,
        FIRST_CHUNK_FRAMES + 1,
        CHUNK_FRAMES - 1,
        CHUNK_FRAMES,
        CHUNK_FRAMES + 1,
        FIRST_CHUNK_FRAMES + CHUNK_FRAMES,
        79,
        200,
    ],
)
def test_chunked_emissions_equal_a_whole_history_decode(num_frames: int) -> None:
    codec, _fake = _make_codec()
    talker = _make_talker(codec)
    state = _GepardState(request_id="req-0")

    _generate(talker, state, num_frames, stop_at_end=True)

    torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
    assert state.emitted_frames == num_frames


def test_every_frame_is_emitted_exactly_once() -> None:
    """The trim must not drop or duplicate samples at a chunk boundary."""
    codec, _fake = _make_codec()
    talker = _make_talker(codec)
    state = _GepardState(request_id="req-0")
    num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES + 3

    _generate(talker, state, num_frames, stop_at_end=True)
    audio = _drain(talker)

    assert audio.numel() == num_frames * SAMPLES_PER_FRAME
    # Tags are globally unique, so a duplicated or dropped frame shows up here.
    assert len(torch.unique(audio)) == audio.numel()


def test_decode_windows_carry_lookback_context() -> None:
    """Every decode after the first must re-decode LOOKBACK_FRAMES of context."""
    codec, fake = _make_codec()
    talker = _make_talker(codec)
    state = _GepardState(request_id="req-0")
    num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES

    _generate(talker, state, num_frames, stop_at_end=True)

    assert len(fake.windows) >= 3, "expected a first chunk plus steady-state chunks"
    assert fake.windows[0] == FIRST_CHUNK_FRAMES, "the first chunk has no left context"
    for window in fake.windows[1:]:
        assert window == CHUNK_FRAMES + LOOKBACK_FRAMES


def test_codec_runs_once_per_chunk_not_once_per_frame() -> None:
    codec, fake = _make_codec()
    talker = _make_talker(codec)
    state = _GepardState(request_id="req-0")
    num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES

    _generate(talker, state, num_frames, stop_at_end=True)

    assert len(fake.windows) == 4
    assert sum(fake.windows) < num_frames * 2  # far below one decode per frame


def test_final_emit_drains_a_partial_chunk() -> None:
    """A request can end mid-chunk (max_tokens, STOP); the frames committed
    since the last cadence boundary must still reach the queue.

    Arithmetic only: this drives ``_emit_audio`` directly, so it says nothing
    about *whether* the engine reaches it on the request's last step. That
    wiring lives in ``test_gepard_wiring.py`` — a bug there leaves this green.
    """
    codec, _fake = _make_codec()
    talker = _make_talker(codec)
    state = _GepardState(request_id="req-0")
    num_frames = FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5

    _generate(talker, state, num_frames, stop_at_end=False)
    assert state.emitted_frames < num_frames, "a partial chunk should still be pending"

    talker._emit_audio(state, is_final=True)

    torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
    assert state.emitted_frames == num_frames


def test_emitting_with_nothing_pending_is_a_no_op() -> None:
    codec, fake = _make_codec()
    talker = _make_talker(codec)
    state = _GepardState(request_id="req-0")

    _generate(talker, state, 4, stop_at_end=True)
    calls_after_flush = len(fake.windows)
    talker._audio_queue.clear()

    talker._emit_audio(state, is_final=True)

    assert len(fake.windows) == calls_after_flush, "codec ran with no new frames"
    assert talker._audio_queue == []


def test_concurrent_requests_keep_separate_windows() -> None:
    codec, _fake = _make_codec()
    talker = _make_talker(codec)
    states = [_GepardState(request_id=f"req-{i}") for i in range(2)]

    for state in states:
        _generate(talker, state, FIRST_CHUNK_FRAMES, stop_at_end=True)

    by_req: dict[str, list[torch.Tensor]] = {}
    for req_id, audio in talker._audio_queue:
        by_req.setdefault(req_id, []).append(audio.reshape(-1))
    assert set(by_req) == {"req-0", "req-1"}
    for chunks in by_req.values():
        torch.testing.assert_close(torch.cat(chunks), _expected_samples(FIRST_CHUNK_FRAMES))


def test_end_of_speech_tail_fades_and_pads() -> None:
    audio = torch.ones(SAMPLE_RATE)  # 1 s of full-scale signal
    out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=10.0, silence_ms=20.0)

    fade = int(SAMPLE_RATE * 10.0 / 1000.0)
    pad = int(SAMPLE_RATE * 20.0 / 1000.0)
    assert out.numel() == SAMPLE_RATE + pad
    assert torch.all(out[-pad:] == 0.0), "trailing pad must be silent"
    assert out[SAMPLE_RATE - fade - 1] == 1.0, "the fade must not reach back past fade_ms"
    # Monotonically ramped to zero — the pad alone would leave the click.
    ramp = out[SAMPLE_RATE - fade : SAMPLE_RATE]
    assert torch.all(ramp[:-1] >= ramp[1:])
    assert ramp[-1] == 0.0


def test_end_of_speech_tail_is_a_no_op_by_default() -> None:
    audio = torch.ones(128)
    out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=0.0, silence_ms=0.0)
    assert out is audio


def _repeat_count(ids: list[int], cfg: GepardConfig) -> int:
    """How many ``[SOT ... EOT]`` copies the layout carries."""
    return ids.count(cfg.start_of_text)


def test_gepard_prompt_layout() -> None:
    """The prompt carries the speaker slots and exactly one SOS, at the end."""
    cfg = GepardConfig()
    ids = build_gepard_prompt_ids([1, 2, 3], config=cfg)

    slots = [cfg.speaker_token_base + i for i in range(cfg.num_speaker_prefix)]
    assert ids[: cfg.num_speaker_prefix] == slots
    assert ids.count(cfg.start_of_speech) == 1, "only the canonical copy may carry SOS"
    assert ids[-1] == cfg.start_of_speech, "the first frame is sampled from the SOS position"


def test_short_text_repeats_to_reach_the_target_token_mass() -> None:
    """A short text is repeated until its region holds ~target_text_tokens."""
    cfg = GepardConfig()
    text = [1, 2, 3, 4]  # target 16 / 4 -> 4 copies
    ids = build_gepard_prompt_ids(text, config=cfg)

    assert _repeat_count(ids, cfg) == 4
    assert ids.count(cfg.start_of_speech) == 1, "only the last copy opens the speech region"
    assert ids[-1] == cfg.start_of_speech


@pytest.mark.parametrize(
    ("num_text_tokens", "repeated"),
    [
        (APPLY_BELOW - 1, True),
        (APPLY_BELOW, False),
        (APPLY_BELOW + 1, False),
    ],
)
def test_apply_below_is_the_repetition_boundary(num_text_tokens: int, repeated: bool) -> None:
    """Both sides of the threshold: at apply_below the text is left alone.

    The layout has to match training here or WER collapses, and the boundary
    is where an off-by-one would be invisible in a single happy-path clip.
    """
    cfg = GepardConfig()
    ids = build_gepard_prompt_ids(list(range(1, num_text_tokens + 1)), config=cfg)

    assert (_repeat_count(ids, cfg) > 1) is repeated


def test_max_repeats_caps_a_one_token_text() -> None:
    """target/1 would ask for 16 copies; the cap keeps the prefill bounded."""
    cfg = GepardConfig()
    ids = build_gepard_prompt_ids([1], config=cfg)

    assert _repeat_count(ids, cfg) == cfg.text_repetition_max_repeats


def test_repetition_thresholds_are_read_from_the_config() -> None:
    """The numbers come from the checkpoint sidecar, not from literals.

    A checkpoint whose text_repetition block differs must change the layout;
    duplicating the values in the builder would silently ignore it.
    """
    cfg = GepardConfig(text_repetition={"enabled": True, "target_text_tokens": 9, "apply_below": 8, "max_repeats": 2})
    assert _repeat_count(build_gepard_prompt_ids([1, 2, 3], config=cfg), cfg) == 2  # ceil(9/3)=3, capped at 2
    assert _repeat_count(build_gepard_prompt_ids(list(range(8)), config=cfg), cfg) == 1  # at apply_below


def test_repetition_can_be_disabled() -> None:
    cfg = GepardConfig(text_repetition={"enabled": False})
    ids = build_gepard_prompt_ids([1], config=cfg)

    assert _repeat_count(ids, cfg) == 1
    assert ids[-1] == cfg.start_of_speech


def test_empty_text_is_rejected() -> None:
    """An empty text would still assemble a valid layout and be voiced.

    ``[slots, SOT, EOT, SOS]`` parses, prefills and generates: the failure is
    audible, not raised, which is the worst shape for a caller to debug.
    """
    with pytest.raises(ValueError, match="at least one text token"):
        build_gepard_prompt_ids([], config=GepardConfig())


def test_packaged_deploy_config_is_where_the_example_looks_for_it() -> None:
    """The example resolves its default off ``vllm_omni.__file__``; tests reach
    the same YAML through ``tests.helpers.stage_config``, so nothing covered the
    packaged location until this.

    ``pipeline: gepard`` is the load-bearing key: the checkpoint self-identifies
    as qwen3_5_text, so without the pin the architectures fallback routes it to
    the diffusion registry.

    This asserts the location, not the example's constant — importing the
    example would pull ``Omni`` and NeMo into a CPU test. Running the quick
    start verbatim stays a release check.
    """
    path = Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml"
    assert path.is_file(), f"packaged deploy config missing at {path}"

    deploy = load_deploy_config(path)
    assert deploy.pipeline == "gepard"
    assert deploy.stages, "deploy config declares no stages"