# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import numpy as np
import pytest

from vllm.config import ModelConfig
from vllm.multimodal import MULTIMODAL_REGISTRY

AUDIO_MODEL_SETTINGS = {
    "ibm-granite/granite-speech-3.3-2b": {
        "prompt": (
            "<|start_of_role|>system<|end_of_role|>"
            "You are a helpful AI assistant<|end_of_text|>\n"
            "<|start_of_role|>user<|end_of_role|>"
            "<|audio|>can you transcribe the speech into a written format?"
            "<|end_of_text|>\n"
            "<|start_of_role|>assistant<|end_of_role|>"
        ),
    },
    "nvidia/audio-flamingo-3-hf": {
        "prompt": (
            "<|im_start|>system\n"
            "You are a helpful assistant.<|im_end|>\n"
            "<|im_start|>user\n"
            "<sound>Transcribe the input speech.<|im_end|>\n"
            "<|im_start|>assistant\n"
        ),
    },
    "mistralai/Voxtral-Mini-3B-2507": {
        "prompt": ("[INST][AUDIO]What can you tell me about this audio?[/INST]"),
    },
    "microsoft/VibeVoice-ASR-HF": {
        "prompt": (
            "<|im_start|>system\n"
            "You are a helpful assistant that transcribes audio input "
            "into text output in JSON format.<|im_end|>\n"
            "<|im_start|>user\n"
            "<|object_ref_start|><|box_start|><|object_ref_end|>\n"
            "This is a 1.0 seconds audio, please transcribe it with "
            "these keys: Start time, End time, Speaker ID, Content"
            "<|im_end|>\n"
            "<|im_start|>assistant\n"
        ),
    },
    "zai-org/GLM-ASR-Nano-2512": {
        "prompt": (
            "<|user|>\n"
            "<|begin_of_audio|><|pad|><|end_of_audio|><|user|>\n"
            "Please transcribe this audio into text"
            "<|assistant|>\n"
        ),
    },
}


@pytest.mark.parametrize(
    "model_id",
    [
        "ibm-granite/granite-speech-3.3-2b",
        "nvidia/audio-flamingo-3-hf",
        pytest.param(
            "mistralai/Voxtral-Mini-3B-2507",
            marks=pytest.mark.xfail(
                reason="MistralCommonBackend.encode does not produce the audio "
                "placeholder token (ID 24) from raw text. apply_chat_template "
                "yields token IDs with placeholders, but MultiModalProcessor."
                "apply() decodes the prompt back to text and re-tokenizes, at "
                "which point the placeholders are lost. Fix belongs in "
                "mistral_common or in the Voxtral-specific path.",
                strict=False,
            ),
        ),
        "microsoft/VibeVoice-ASR-HF",
        "zai-org/GLM-ASR-Nano-2512",
    ],
)
def test_audio_multimodal_processor(model_id):
    settings = AUDIO_MODEL_SETTINGS[model_id]

    model_config = ModelConfig(
        model=model_id,
        model_impl="transformers",
    )

    mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)

    audio = np.zeros(16000, dtype=np.float32)
    mm_data = {"audio": (audio, 16000)}

    result = mm_processor(
        prompt=settings["prompt"],
        mm_items=mm_processor.info.parse_mm_data(mm_data),
        hf_processor_mm_kwargs={},
    )

    assert "prompt_token_ids" in result
    assert len(result["prompt_token_ids"]) > 0

    mm_placeholders = result.get("mm_placeholders", {})
    assert "audio" in mm_placeholders, f"No audio placeholders found for {model_id}"
    assert len(mm_placeholders["audio"]) == 1

    placeholder = mm_placeholders["audio"][0]
    assert placeholder.length > 0
    assert placeholder.offset >= 0

    audio_items = result.get("mm_kwargs", {}).get("audio", [])
    assert len(audio_items) == 1, f"Expected 1 audio item, got {len(audio_items)}"
    item_keys = list(audio_items[0].keys())
    has_features = "input_features" in item_keys or "input_values" in item_keys
    assert has_features, (
        f"No audio features (input_features/input_values) in {item_keys} for {model_id}"
    )


def _process_granite_speech(separator: str):
    model_id = "ibm-granite/granite-speech-3.3-2b"
    model_config = ModelConfig(model=model_id, model_impl="transformers")
    mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)

    audio_token = mm_processor.info.get_hf_processor().audio_token
    # One token per audio; the processor expands each to its placeholder run.
    prompt = (
        "<|start_of_role|>user<|end_of_role|>"
        f"{audio_token}{separator}{audio_token} transcribe<|end_of_text|>\n"
    )
    audios = [np.zeros(16000, dtype=np.float32), np.zeros(24000, dtype=np.float32)]

    return mm_processor(
        prompt=prompt,
        mm_items=mm_processor.info.parse_mm_data({"audio": audios}),
        hf_processor_mm_kwargs={},
    )


def test_audio_multiple_inputs():
    """Multiple audios per prompt are each detected as a separate placeholder
    and multi-modal item by the Transformers modelling backend."""
    result = _process_granite_speech(separator=" and ")

    assert len(result["mm_placeholders"]["audio"]) == 2
    assert len(result["mm_kwargs"]["audio"]) == 2


def test_audio_fields_not_claimed_by_image():
    """Audio fields survive when the image branch is also active."""
    model_id = "ibm-granite/granite-speech-3.3-2b"
    model_config = ModelConfig(model=model_id, model_impl="transformers")
    mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)

    audio_keys = ["input_features", "input_features_mask"]
    owned = mm_processor._partition_keys_by_modality(audio_keys, ["audio", "image"])

    assert owned["audio"] == audio_keys
    assert owned["image"] == []


def test_unclaimed_fields_warn_rather_than_raise():
    """Keys no sub-processor declares are dropped with a warning, not an error."""
    model_id = "ibm-granite/granite-speech-3.3-2b"
    model_config = ModelConfig(model=model_id, model_impl="transformers")
    mm_processor = MULTIMODAL_REGISTRY.create_processor(model_config)

    owned = mm_processor._partition_keys_by_modality(
        ["input_features", "surprise_field"], ["audio", "image"]
    )

    assert owned["audio"] == ["input_features"]
    assert owned["image"] == []


def test_audio_adjacent_inputs():
    """Adjacent audios are rejected rather than silently merged into one placeholder."""
    with pytest.raises(ValueError, match="told apart"):
        _process_granite_speech(separator="")