# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Breeze design, reference conditioning and CFG through the shared speech client."""

import pytest

from tests.helpers.mark import hardware_test
from tests.helpers.media import get_asset_path
from tests.helpers.runtime import OmniServerParams
from tests.helpers.stage_config import get_deploy_config_path

MODEL = "BreezeBlue/Breeze-TTS-2"

pytestmark = [
    pytest.mark.slow,
    pytest.mark.tts,
    pytest.mark.parametrize(
        "omni_server",
        [OmniServerParams(model=MODEL, stage_config_path=get_deploy_config_path("breeze_tts_2.yaml"))],
        indirect=True,
    ),
]


@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_chinese_voice_design(omni_server, online_client, run_level) -> None:
    online_client.send_audio_speech_request(
        {
            "model": omni_server.model,
            "input": "你好，欢迎使用语音合成系统。今天的天气很好。",
            "instructions": "年轻女性，使用标准普通话，语速适中。",
            "extra_params": {"guidance_scale": 4.0},
            "transcript_language": "zh",
            # Whisper-small transcribes "合成" as "和程" on this fixed clip.
            # Use the stronger primary ASR without changing the expected text
            # or similarity threshold, or retrying after a failed assertion.
            "transcript_model": "large-v3",
            "response_format": "wav",
            "sample_rate": 24000,
            "min_audio_bytes": 24000,
            "seed": 42,
            "max_new_tokens": 8 if run_level == "core_model" else 250,
        }
    )


@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_english_streaming(omni_server, online_client, run_level) -> None:
    online_client.send_audio_speech_request(
        {
            "model": omni_server.model,
            "input": "Hello, this is a demonstration of natural speech synthesis.",
            "instructions": "A warm, clear female voice speaking English at a moderate pace.",
            "transcript_language": "en",
            "transcript_pcm_sample_rate": 24000,
            "transcript_escalation_model": "large-v3",
            "response_format": "pcm",
            "sample_rate": 24000,
            "min_audio_bytes": 24000,
            "stream": True,
            "stream_format": "audio",
            # This fixed prompt scores -0.11 dB on L4 and 0.56 dB on H100.
            # Keep a small HNR margin and also verify the spoken words above.
            "min_hnr_db": -1.0,
            "seed": 42,
            "max_new_tokens": 8 if run_level == "core_model" else 250,
        }
    )


@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_two_concurrent_requests(omni_server, online_client, run_level) -> None:
    online_client.send_audio_speech_request(
        {
            "model": omni_server.model,
            "input": "Each request keeps its own audio generation state.",
            "instructions": "A clear English voice.",
            "extra_params": {"guidance_scale": 4.0},
            "transcript_language": "en",
            "response_format": "wav",
            "sample_rate": 24000,
            "min_audio_bytes": 24000,
            "seed": 123,
            "max_new_tokens": 8 if run_level == "core_model" else 250,
        },
        request_num=2,
    )


@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("guidance_scale", [1.0, 4.0])
def test_reference_cloning_and_direction(omni_server, online_client, run_level, guidance_scale) -> None:
    online_client.send_audio_speech_request(
        {
            "model": omni_server.model,
            "input": "Please read this sentence clearly and naturally.",
            "instructions": "Keep the reference voice. Speak gently and calmly.",
            "ref_audio": get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True),
            "ref_text": "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you.",
            "task_type": "Base",
            "extra_params": {"guidance_scale": guidance_scale},
            "transcript_language": "en",
            "response_format": "wav",
            "sample_rate": 24000,
            "min_audio_bytes": 24000,
            "seed": 42,
            "max_new_tokens": 8 if run_level == "core_model" else 250,
        }
    )