# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic.

The prompt transformations are not cosmetic: they decide token positions, and
token positions decide the audio. These tests pin the documented behaviours
that a well-meaning cleanup would otherwise "fix".
"""

from __future__ import annotations

import pytest

from vllm_omni.model_executor.models.minimax_music3.chunking import (
    chunk_windows,
    crop_sample_bounds,
    overlap_mel_length,
)
from vllm_omni.model_executor.models.minimax_music3.constants import (
    AR_CHUNK_FRAMES,
    AR_CHUNK_HOP_FRAMES,
    BLEND_MEL_FRAMES,
    DAV_HOP_SAMPLES,
    HOP_MEL_FRAMES,
)
from vllm_omni.model_executor.models.minimax_music3.prompt import (
    SPECIAL_TOKEN_IDS,
    build_cfg_null_token_ids,
    build_prompt,
    clean_caption,
    normalize_lyrics,
)

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


class TestNormalizeLyrics:
    def test_prepends_start_tag(self) -> None:
        assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n")

    def test_lowercases_structure_tags(self) -> None:
        assert "[verse]" in normalize_lyrics("[Verse]\nhello")
        assert "[Verse]" not in normalize_lyrics("[Verse]\nhello")

    def test_text_after_a_leading_tag_is_dropped(self) -> None:
        """Documented contract: a tag must sit on its own line.

        Lyrics written next to a tag are silently lost, which is why the
        serving docs tell callers to put tags on their own line.
        """
        out = normalize_lyrics("[Verse] Walking down the street")
        assert "Walking" not in out
        assert "[verse]" in out

    def test_text_on_its_own_line_survives(self) -> None:
        out = normalize_lyrics("[Verse]\nWalking down the street")
        assert "Walking down the street" in out

    def test_bracket_space_becomes_newline(self) -> None:
        assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]"


class TestCleanCaption:
    def test_special_token_form_is_rewritten_not_passed_through(self) -> None:
        """A caption must not be able to inject a real special token."""
        assert clean_caption("<|bpm 92|>") == "bpm is 92"
        assert "<|" not in clean_caption("<|bpm 92|>")

    def test_lone_tag_keeps_its_inner_text(self) -> None:
        assert clean_caption("<|instrumental|>") == "instrumental"

    def test_markdown_is_stripped(self) -> None:
        assert clean_caption("## Heading") == "Heading"
        assert clean_caption("**bold**") == "bold"
        assert clean_caption("- item") == "item"

    def test_blank_lines_collapse_at_two_not_three(self) -> None:
        """The source collapses runs of two or more, not three or more.

        This shifts blank-line token positions and therefore the audio, so it
        is pinned deliberately.
        """
        assert clean_caption("a\n\nb") == "a\nb"
        assert clean_caption("a\n\n\n\nb") == "a\nb"


class TestBuildPrompt:
    def test_framing_tokens_are_present_and_ordered(self) -> None:
        prompt = build_prompt("a bright pop song", "[Verse]\nhello")
        assert prompt.startswith("<|im_start|><|caption_start|>")
        assert prompt.endswith("<|im_end|><|audio_start|>")
        assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>")


class TestBuildCfgNullTokenIds:
    @staticmethod
    def _prompt_ids(body_len: int) -> list[int]:
        return (
            [SPECIAL_TOKEN_IDS["<|im_start|>"]]
            + list(range(1000, 1000 + body_len))
            + [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]]
        )

    def test_length_is_preserved(self) -> None:
        """Equal length is what lets the two rows decode in lockstep."""
        ids = self._prompt_ids(12)
        assert len(build_cfg_null_token_ids(ids)) == len(ids)

    def test_body_is_replaced_and_framing_survives(self) -> None:
        ids = self._prompt_ids(12)
        null = build_cfg_null_token_ids(ids)
        assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"]
        assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"]
        assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"]
        assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]}

    def test_input_is_not_mutated(self) -> None:
        ids = self._prompt_ids(6)
        before = list(ids)
        build_cfg_null_token_ids(ids)
        assert ids == before

    def test_minimal_prompt_with_empty_body(self) -> None:
        ids = [
            SPECIAL_TOKEN_IDS["<|im_start|>"],
            SPECIAL_TOKEN_IDS["<|audio_cfg|>"],
            SPECIAL_TOKEN_IDS["<|im_end|>"],
            SPECIAL_TOKEN_IDS["<|audio_start|>"],
        ]
        assert build_cfg_null_token_ids(ids) == ids

    def test_too_short_is_rejected(self) -> None:
        with pytest.raises(ValueError, match="too short"):
            build_cfg_null_token_ids([1, 2, 3])

    def test_wrong_framing_is_rejected(self) -> None:
        with pytest.raises(ValueError, match="framing"):
            build_cfg_null_token_ids([9, 9, 9, 9, 9])


class TestChunkWindows:
    def test_no_frames_yields_nothing(self) -> None:
        assert chunk_windows(0) == []

    def test_short_song_is_one_final_window(self) -> None:
        windows = chunk_windows(50)
        assert len(windows) == 1
        assert (windows[0].start, windows[0].end) == (0, 50)
        assert windows[0].is_first and windows[0].is_last

    def test_exactly_one_window(self) -> None:
        windows = chunk_windows(AR_CHUNK_FRAMES)
        assert len(windows) == 1
        assert windows[0].is_last

    def test_windows_advance_by_the_hop(self) -> None:
        windows = chunk_windows(450)
        starts = [w.start for w in windows]
        assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES))
        assert windows[-1].is_last
        assert not any(w.is_last for w in windows[:-1])

    def test_windows_cover_every_frame(self) -> None:
        frames = 437
        windows = chunk_windows(frames)
        covered: set[int] = set()
        for window in windows:
            covered.update(range(window.start, window.end))
        assert covered == set(range(frames))

    def test_final_window_is_clamped_to_the_song(self) -> None:
        windows = chunk_windows(437)
        assert windows[-1].end == 437

    def test_only_the_first_window_is_first(self) -> None:
        windows = chunk_windows(600)
        assert windows[0].is_first
        assert not any(w.is_first for w in windows[1:])


class TestCropBounds:
    def test_single_window_song_is_not_cropped(self) -> None:
        window = chunk_windows(50)[0]
        assert crop_sample_bounds(window) == (0, 0)

    def test_first_window_keeps_its_head(self) -> None:
        left, right = crop_sample_bounds(chunk_windows(600)[0])
        assert left == 0
        assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES

    def test_last_window_keeps_its_tail(self) -> None:
        left, right = crop_sample_bounds(chunk_windows(600)[-1])
        assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES
        assert right == 0

    def test_middle_window_is_cropped_on_both_sides(self) -> None:
        windows = chunk_windows(900)
        middle = windows[len(windows) // 2]
        left, right = crop_sample_bounds(middle)
        assert left > 0 and right > 0

    def test_overlap_is_half_a_window(self) -> None:
        assert overlap_mel_length() == HOP_MEL_FRAMES // 2