# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E Online tests for MiMo-Audio model with audio/text input and audio output.
"""

import os
from pathlib import Path

import pytest

from tests.conftest import (
    OmniServerParams,
    dummy_messages_from_mix_data,
    generate_synthetic_audio,
    modify_stage_config,
)
from tests.utils import hardware_test
from vllm_omni.model_executor.model_loader.weight_utils import (
    download_weights_from_hf_specific,
)

MIMO_AUDIO_TOKENIZER_REPO = "XiaomiMiMo/MiMo-Audio-Tokenizer"
CHAT_TEMPLATE_PATH = str(
    Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "mimo_audio" / "chat_template.jinja"
)
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

models = ["XiaomiMiMo/MiMo-Audio-7B-Instruct"]


def get_chunk_config():
    path = modify_stage_config(
        str(Path(__file__).parent.parent / "stage_configs" / "mimo_audio_ci.yaml"),
        updates={
            "async_chunk": True,
            "stage_args": {
                0: {
                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav_async_chunk"
                },
                1: {
                    "engine_args.max_model_len": 8192,
                    "engine_args.max_num_batched_tokens": 8192,
                },
            },
        },
        deletes={"stage_args": {1: ["custom_process_input_func"]}},
    )
    return path


def download_tokenizer():
    tokenizer_path = os.environ.get("MIMO_AUDIO_TOKENIZER_PATH", MIMO_AUDIO_TOKENIZER_REPO)
    if os.path.exists(tokenizer_path):
        return tokenizer_path
    local_path = download_weights_from_hf_specific(
        model_name_or_path=MIMO_AUDIO_TOKENIZER_REPO,
        cache_dir=None,
        allow_patterns=["*"],
        require_all=True,
    )
    return local_path


# CI stage config for H100 / MI325
# Guard module-level setup so test collection doesn't fail in environments
# where the model cache is read-only or models aren't available.
try:
    stage_configs = [get_chunk_config()]
    tokenizer_path = download_tokenizer()
    os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path

    test_params = [
        OmniServerParams(
            model=model,
            stage_config_path=stage_config,
            server_args=["--chat-template", CHAT_TEMPLATE_PATH],
        )
        for model in models
        for stage_config in stage_configs
    ]
except Exception as exc:
    pytest.skip(
        f"MiMo-Audio online serving tests skipped: module setup failed ({type(exc).__name__}: {exc})",
        allow_module_level=True,
    )


def get_prompt(prompt_type="text_only"):
    prompts = {
        "text_only": "What is the capital of China? Answer in 20 words.",
        "audio": "What is recited in the audio?",
    }
    return prompts.get(prompt_type, prompts["text_only"])


def get_max_batch_size(size_type="few"):
    batch_sizes = {"few": 5, "medium": 100, "large": 256}
    return batch_sizes.get(size_type, 5)


@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
    """
    Test audio and text input processing and text/audio output generation via OpenAI API.
    Deploy Setting: default yaml
    Input Modal: text + audio
    Output Modal: text + audio
    Input Setting: stream=True
    Datasets: single request
    """

    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1, sample_rate=24000)['base64']}"
    messages = dummy_messages_from_mix_data(
        audio_data_url=audio_data_url,
        content_text=get_prompt("audio"),
    )

    request_config = {
        "model": omni_server.model,
        "messages": messages,
        "stream": True,
        "key_words": {
            "audio": ["test"],
        },
    }

    # Test single completion
    openai_client.send_omni_request(request_config)


@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_to_text_001(omni_server, openai_client) -> None:
    """
    Test text input processing and text-only output generation via OpenAI API.
    Deploy Setting: default yaml
    Input Modal: text
    Output Modal: text
    Datasets: few requests
    """
    messages = dummy_messages_from_mix_data(content_text=get_prompt())

    request_config = {
        "model": omni_server.model,
        "messages": messages,
        "stream": False,
        "modalities": ["text"],
        "key_words": {"text": ["beijing"]},
    }

    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())