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

# imports for structured outputs tests
import asyncio
import io
import json

import numpy as np
import openai
import pytest
import pytest_asyncio
import soundfile as sf

from tests.utils import RemoteOpenAIServer
from vllm.multimodal.media.audio import load_audio
from vllm.platforms import current_platform

MODEL_NAME = "openai/whisper-large-v3-turbo"

# Disable prefix caching on ROCm to reduce non-determinism in
# streaming-vs-non-streaming comparisons.
_ROCM_ARGS = ["--no-enable-prefix-caching"] if current_platform.is_rocm() else []


def _get_attention_backend_params() -> list[str | None]:
    """Return attention backends to parametrize the server fixture with.

    On ROCm, we test multiple backends explicitly:
      - None: default auto-selection (ROCM_ATTN for decoder self-attention,
              falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for
              cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER)
      - TRITON_ATTN: always available on ROCm
      - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950

    On non-ROCm platforms, we just run with the default backend.
    """
    try:
        from vllm.platforms import current_platform

        if current_platform.is_rocm():
            backends: list[str | None] = [None, "TRITON_ATTN"]
            from vllm.platforms.rocm import _ON_MI3XX

            if _ON_MI3XX:
                backends.append("ROCM_AITER_UNIFIED_ATTN")
            return backends
    except Exception:
        pass
    return [None]


# Aiter backends need VLLM_ROCM_USE_AITER=1 (and MHA=1 for ROCM_AITER_FA)
# to be enabled in the server subprocess.
_AITER_ENV = {
    "VLLM_ROCM_USE_AITER": "1",
    "VLLM_ROCM_USE_AITER_MHA": "1",
}

_ATTN_BACKENDS = _get_attention_backend_params()
_ATTN_IDS = [b or "default" for b in _ATTN_BACKENDS]


@pytest.fixture(scope="module", params=_ATTN_BACKENDS, ids=_ATTN_IDS)
def server(request):
    args = [*_ROCM_ARGS]
    env_dict = None
    if request.param is not None:
        args += ["--attention-backend", request.param]
        if "AITER" in request.param:
            # TODO: re-enable once AITER reenables fp16 unified attention.
            pytest.skip("ROCM_AITER_UNIFIED_ATTN does not support fp16")
    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
        yield remote_server


@pytest_asyncio.fixture
async def whisper_client(server):
    async with server.get_async_client() as async_client:
        yield async_client


@pytest.mark.asyncio
async def test_basic_audio(whisper_client, mary_had_lamb):
    # Based on https://github.com/openai/openai-cookbook/blob/main/examples/Whisper_prompting_guide.ipynb.
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        temperature=0.0,
    )
    out = json.loads(transcription)
    out_text = out["text"]
    out_usage = out["usage"]
    assert "Mary had a little lamb," in out_text
    assert out_usage["seconds"] == 16, out_usage["seconds"]


@pytest.mark.asyncio
async def test_basic_audio_batched(mary_had_lamb, winning_call, whisper_client):
    transcription = whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        temperature=0.0,
    )
    transcription2 = whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=winning_call,
        language="en",
        response_format="text",
        temperature=0.0,
    )
    # Await both transcriptions by scheduling coroutines together
    transcription, transcription2 = await asyncio.gather(transcription, transcription2)
    out = json.loads(transcription)
    out_text = out["text"]
    assert "Mary had a little lamb," in out_text
    out2 = json.loads(transcription2)
    out_text2 = out2["text"]
    assert "Edgar Martinez" in out_text2


@pytest.mark.asyncio
async def test_bad_requests(mary_had_lamb, whisper_client):
    # invalid language
    with pytest.raises(openai.BadRequestError):
        await whisper_client.audio.transcriptions.create(
            model=MODEL_NAME, file=mary_had_lamb, language="hh", temperature=0.0
        )


@pytest.mark.asyncio
async def test_long_audio_request(mary_had_lamb, whisper_client):
    mary_had_lamb.seek(0)
    audio, sr = load_audio(mary_had_lamb)
    # Add small silence after each audio for repeatability in the split process
    audio = np.pad(audio, (0, 1600))
    repeated_audio = np.tile(audio, 10)
    # Repeated audio to buffer
    buffer = io.BytesIO()
    sf.write(buffer, repeated_audio, sr, format="WAV")
    buffer.seek(0)
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=buffer,
        language="en",
        response_format="text",
        temperature=0.0,
    )
    out = json.loads(transcription)
    out_text = out["text"]
    out_usage = out["usage"]
    counts = out_text.count("Mary had a little lamb")
    assert counts == 10, counts
    assert out_usage["seconds"] == 161, out_usage["seconds"]


@pytest.mark.asyncio
async def test_invalid_audio_file(whisper_client):
    """Corrupted audio should surface as HTTP 400."""
    invalid_audio = io.BytesIO(b"not a valid audio file")
    invalid_audio.name = "invalid.wav"

    with pytest.raises(openai.BadRequestError) as exc_info:
        await whisper_client.audio.transcriptions.create(
            model=MODEL_NAME,
            file=invalid_audio,
            language="en",
        )

    assert exc_info.value.status_code == 400
    assert "Invalid or unsupported audio file" in exc_info.value.message


@pytest.mark.asyncio
async def test_completion_endpoints(whisper_client):
    # text to text model
    with pytest.raises(openai.NotFoundError):
        await whisper_client.chat.completions.create(
            model=MODEL_NAME,
            messages=[{"role": "system", "content": "You are a helpful assistant."}],
        )

    with pytest.raises(openai.NotFoundError):
        await whisper_client.completions.create(model=MODEL_NAME, prompt="Hello")


@pytest.mark.asyncio
async def test_streaming_response(winning_call, whisper_client):
    transcription = ""
    res_no_stream = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=winning_call,
        response_format="json",
        language="en",
        temperature=0.0,
    )
    res = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=winning_call,
        language="en",
        temperature=0.0,
        stream=True,
        timeout=30,
    )
    # Reconstruct from chunks and validate
    async for chunk in res:
        text = chunk.choices[0]["delta"]["content"]
        transcription += text

    assert transcription == res_no_stream.text


@pytest.mark.asyncio
async def test_stream_options(winning_call, whisper_client):
    res = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=winning_call,
        language="en",
        temperature=0.0,
        stream=True,
        extra_body=dict(stream_include_usage=True, stream_continuous_usage_stats=True),
        timeout=30,
    )
    final = False
    continuous = True
    async for chunk in res:
        if not len(chunk.choices):
            # final usage sent
            final = True
        else:
            continuous = continuous and hasattr(chunk, "usage")
    assert final and continuous


@pytest.mark.asyncio
async def test_sampling_params(mary_had_lamb, whisper_client):
    """
    Compare sampling with params and greedy sampling to assert results
    are different when extreme sampling parameters values are picked.
    """
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        temperature=0.8,
        extra_body=dict(
            seed=42,
            repetition_penalty=1.9,
            top_k=12,
            top_p=0.4,
            min_p=0.5,
            frequency_penalty=1.8,
            presence_penalty=2.0,
        ),
    )

    greedy_transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        temperature=0.0,
        extra_body=dict(seed=42),
    )

    assert greedy_transcription.text != transcription.text


@pytest.mark.asyncio
async def test_audio_prompt(mary_had_lamb, whisper_client):
    prompt = "This is a speech, recorded in a phonograph."
    # Prompts should not omit the part of original prompt while transcribing.
    prefix = "The first words I spoke in the original phonograph"
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        temperature=0.0,
    )
    out = json.loads(transcription)["text"]
    assert prefix in out
    transcription_wprompt = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        prompt=prompt,
        temperature=0.0,
    )
    out_prompt = json.loads(transcription_wprompt)["text"]
    assert prefix in out_prompt


@pytest.mark.asyncio
async def test_audio_with_timestamp(mary_had_lamb, whisper_client):
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="verbose_json",
        temperature=0.0,
    )
    assert transcription.segments is not None
    assert len(transcription.segments) > 0
    assert transcription.segments[0].avg_logprob is not None
    assert transcription.segments[0].compression_ratio is not None


@pytest.mark.asyncio
async def test_audio_with_max_tokens(whisper_client, mary_had_lamb):
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        temperature=0.0,
        extra_body={"max_completion_tokens": 1},
    )
    out = json.loads(transcription)
    out_text = out["text"]
    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(MODEL_NAME)
    out_tokens = tok(out_text, add_special_tokens=False)["input_ids"]
    assert len(out_tokens) == 1
    # max_completion_tokens > max_model_len
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        temperature=0.0,
        extra_body={"max_completion_tokens": int(1e6)},
    )
    out = json.loads(transcription)
    out_text = out["text"]
    out_tokens = tok(out_text, add_special_tokens=False)["input_ids"]
    assert len(out_tokens) < 450  # ~Whisper max output len


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("fixture_name", "expected_lang", "expected_text"),
    [
        ("mary_had_lamb", "en", ["Mary had a little lamb"]),
        ("foscolo", "it", ["zacinto", "sacre"]),
    ],
    ids=["english", "italian"],
)
async def test_language_auto_detect(
    whisper_client, fixture_name, expected_lang, expected_text, request
):
    """Auto-detect language when no language param is provided."""
    audio_file = request.getfixturevalue(fixture_name)
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=audio_file,
        response_format="verbose_json",
        temperature=0.0,
    )
    assert transcription.language == expected_lang
    text_lower = transcription.text.lower()
    assert any(word.lower() in text_lower for word in expected_text), (
        f"Expected {expected_lang} text but got: {transcription.text}"
    )


@pytest.mark.asyncio
async def test_whisper_beam_search_single_beam(mary_had_lamb, whisper_client):
    """Test beam search with encoder-decoder model (Whisper) on transcriptions with
    one beam aligns with greedy decoding.
    """
    beam_transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        temperature=0.0,
        extra_body=dict(
            use_beam_search=True,
            n=1,
        ),
    )

    greedy_transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        response_format="text",
        temperature=0.0,
    )

    greedy_res = json.loads(greedy_transcription)["text"]
    beam_res = json.loads(beam_transcription)["text"]
    assert greedy_res == beam_res


@pytest.mark.asyncio
async def test_whisper_beam_search_multibeam(mary_had_lamb, whisper_client):
    """Test n>1 for beam search returns one transcription (best beam)."""
    transcription = await whisper_client.audio.transcriptions.create(
        model=MODEL_NAME,
        file=mary_had_lamb,
        language="en",
        response_format="text",
        temperature=0.0,
        extra_body=dict(
            use_beam_search=True,
            n=2,
        ),
    )

    result = json.loads(transcription)

    text = result["text"]

    assert text is not None
    assert len(text) > 0
    assert "mary had a little lamb" in text.lower()


@pytest.mark.asyncio
async def test_stream_with_beams_raises(winning_call, whisper_client):
    """Test that stream=True + beam search raises bad request for now."""
    with pytest.raises(openai.BadRequestError):
        await whisper_client.audio.transcriptions.create(
            model=MODEL_NAME,
            file=winning_call,
            language="en",
            stream=True,
            extra_body=dict(
                use_beam_search=True,
                n=2,
            ),
        )
