# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

"""
E2E expansion tests for Voxtral TTS offline inference (nightly CI).

Tests both synchronous (Omni) and async streaming (AsyncOmni) paths.
"""

import asyncio
import os
import uuid

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

from pathlib import Path

import numpy as np
import pytest
import torch
from mistral_common.protocol.speech.request import SpeechRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from vllm import SamplingParams

from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniRunner
from tests.helpers.stage_config import get_deploy_config_path
from vllm_omni.entrypoints.async_omni import AsyncOmni

MODEL = "mistralai/Voxtral-4B-TTS-2603"
STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml")
MIN_AUDIO_SAMPLES = 1000
VOICE = "casual_female"
TEST_TEXT = "Hello, how are you?"

pytestmark = [pytest.mark.slow, pytest.mark.tts]

_OMNI_RUNNER_PARAM = (
    MODEL,
    STAGE_CONFIG,
    {"enforce_eager": True},
)


def _compose_request(model_name: str, text: str, voice: str) -> dict:
    """Build the TTS input dict using mistral tokenizer."""
    if Path(model_name).is_dir():
        mistral_tokenizer = MistralTokenizer.from_file(str(Path(model_name) / "tekken.json"))
    else:
        mistral_tokenizer = MistralTokenizer.from_hf_hub(model_name)
    instruct_tokenizer = mistral_tokenizer.instruct_tokenizer

    tokenized = instruct_tokenizer.encode_speech_request(SpeechRequest(input=text, voice=voice))
    return {
        "prompt_token_ids": tokenized.tokens,
        "additional_information": {"voice": [voice]},
    }


@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None:
    """Offline sync path; function-scoped runner so AsyncOmni test does not overlap a live Omni."""
    omni = omni_runner_function.omni
    inputs = _compose_request(MODEL, TEST_TEXT, VOICE)

    sampling_params = SamplingParams(max_tokens=2500)
    sampling_params_list = [sampling_params, sampling_params]

    outputs = list(omni.generate([inputs], sampling_params_list))

    assert len(outputs) > 0, "No outputs generated"

    audio_data = None
    for o in outputs:
        mm = getattr(o, "multimodal_output", None)
        if mm and "audio" in mm:
            audio_data = mm["audio"]
            break

    assert audio_data is not None, "No audio output found in any stage output"

    if isinstance(audio_data, list):
        audio_tensor = torch.cat(audio_data)
    elif isinstance(audio_data, torch.Tensor):
        audio_tensor = audio_data
    else:
        audio_tensor = torch.tensor(audio_data)

    audio_array = audio_tensor.float().cpu().numpy()

    assert len(audio_array) > MIN_AUDIO_SAMPLES, (
        f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}"
    )
    assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence"


@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_voxtral_tts_offline_streaming():
    """Test AsyncOmni streaming inference for Voxtral TTS."""

    async def _run():
        async_omni = AsyncOmni(
            model=MODEL,
            deploy_config=STAGE_CONFIG,
            stage_init_timeout=300,
            enforce_eager=True,
        )

        try:
            inputs = _compose_request(MODEL, TEST_TEXT, VOICE)
            sampling_params = SamplingParams(max_tokens=2500)

            all_audio_chunks = []
            accumulated_samples = 0
            chunk_idx = 0
            async for stage_output in async_omni.generate(
                prompt=inputs,
                request_id=str(uuid.uuid4()),
                sampling_params_list=[sampling_params, sampling_params],
            ):
                mm = getattr(stage_output, "multimodal_output", None)
                if not mm or "audio" not in mm:
                    continue

                audio_chunk = mm["audio"]
                finished = stage_output.finished

                if isinstance(audio_chunk, torch.Tensor):
                    if finished:
                        audio_np = audio_chunk[accumulated_samples:].float().detach().cpu().numpy()
                    else:
                        audio_np = audio_chunk.float().detach().cpu().numpy()
                elif isinstance(audio_chunk, list):
                    audio_np = audio_chunk[chunk_idx].float().detach().cpu().numpy()
                else:
                    audio_np = audio_chunk

                accumulated_samples += len(audio_np)
                all_audio_chunks.append(audio_np)
                chunk_idx += 1

            assert len(all_audio_chunks) > 0, "No audio chunks received from streaming"

            audio_array = np.concatenate(all_audio_chunks)

            assert len(audio_array) > MIN_AUDIO_SAMPLES, (
                f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}"
            )
            assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence"

        finally:
            async_omni.shutdown()

    asyncio.run(_run())