# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Preserve every sample from the engine's fresh DELTA audio payloads."""

import asyncio
import base64
import io
import wave
from types import SimpleNamespace

import numpy as np
import pytest
import torch
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.v1.engine import FinishReason

from tests.helpers.serving_chat import build_serving_chat
from vllm_omni.entrypoints.omni_base import OmniBase
from vllm_omni.entrypoints.openai.serving_video_stream import QwenOmniStreamingVideoHandler
from vllm_omni.entrypoints.openai.video_stream_base import (
    _CODEC_FRAME_SAMPLES,
    OmniStreamingVideoHandler,
    StreamingVideoSessionConfig,
)
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput, MultimodalPayload
from vllm_omni.outputs.output_processor import OmniRequestState

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


def _audio_result(payload):
    completion = MultimodalCompletionOutput(
        index=0,
        text="",
        token_ids=[],
        cumulative_logprob=None,
        logprobs=None,
        multimodal_output=MultimodalPayload.from_dict({"audio": payload}),
    )
    return OmniRequestOutput(
        outputs=[completion],
        final_output_type="audio",
    )


def _pcm(encoded):
    assert encoded is not None
    with wave.open(io.BytesIO(base64.b64decode(encoded)), "rb") as wav:
        assert (wav.getframerate(), wav.getnchannels(), wav.getsampwidth()) == (24000, 1, 2)
        return np.frombuffer(wav.readframes(wav.getnframes()), dtype="<i2")


def _chunks():
    # Exactly representable PCM16 values; distinct ordered samples catch loss,
    # repetition, reordering, and trimming more than once.
    return [torch.arange(start, start + 5000, dtype=torch.float32) / 32768 for start in (0, 5000, 10000)]


@pytest.mark.parametrize("mode", ["fast", "slow"])
@pytest.mark.parametrize("shape", ["tensor", "singleton_list", "batched_list"])
def test_fresh_delta_samples_are_forwarded_once(monkeypatch, mode, shape):
    monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", mode)
    chunks = _chunks()
    payloads = chunks if shape == "tensor" else [[chunk] for chunk in chunks]
    if shape == "batched_list":
        payloads = [[chunks[0]], chunks[1:]]
    seen = 0
    pcm = []
    for payload in payloads:
        encoded, seen = OmniStreamingVideoHandler._extract_audio_delta_b64(_audio_result(payload), seen)
        pcm.append(_pcm(encoded))
    np.testing.assert_array_equal(np.concatenate(pcm), np.arange(_CODEC_FRAME_SAMPLES, 15000))
    assert seen == 3


@pytest.mark.parametrize("mode", ["fast", "slow"])
@pytest.mark.parametrize("empty", [None, [], torch.empty(0), [torch.empty(0)]])
def test_empty_delta_does_not_consume_first_emission(monkeypatch, mode, empty):
    monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", mode)
    encoded, seen = OmniStreamingVideoHandler._extract_audio_delta_b64(_audio_result(empty), 0)
    assert encoded is None
    assert seen == 0
    encoded, seen = OmniStreamingVideoHandler._extract_audio_delta_b64(_audio_result(_chunks()[0]), seen)
    np.testing.assert_array_equal(_pcm(encoded), np.arange(_CODEC_FRAME_SAMPLES, 5000))
    assert seen == 1
    encoded, after = OmniStreamingVideoHandler._extract_audio_delta_b64(_audio_result(empty), seen)
    assert encoded is None
    assert after == seen


@pytest.mark.parametrize("first_length", [1, 3840, 3841])
def test_first_chunk_trim_threshold_is_preserved(first_length):
    first = torch.arange(first_length, dtype=torch.float32) / 32768
    later = torch.arange(5000, 10000, dtype=torch.float32).reshape(1, -1) / 32768
    encoded, seen = OmniStreamingVideoHandler._delta_fast(first, 0)
    start = _CODEC_FRAME_SAMPLES if first_length > 2 * _CODEC_FRAME_SAMPLES else 0
    np.testing.assert_array_equal(_pcm(encoded), np.arange(start, first_length))
    encoded, seen = OmniStreamingVideoHandler._delta_fast(later, seen)
    np.testing.assert_array_equal(_pcm(encoded), np.arange(5000, 10000))
    assert seen == 2


@pytest.mark.parametrize("explicit", [False, True])
def test_only_omitted_sampling_parameters_are_coerced(explicit):
    base = object.__new__(OmniBase)
    base.engine = SimpleNamespace(num_stages=1)
    params = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
    base.default_sampling_params_list = [params]
    base.sampling_constraints_list = [{}]
    resolved = base.resolve_sampling_params_list([params] if explicit else None, allow_delta_coercion=True)
    expected = RequestOutputKind.CUMULATIVE if explicit else RequestOutputKind.DELTA
    assert resolved[0].output_kind == expected
    assert params.output_kind == RequestOutputKind.CUMULATIVE


def _state(output_kind):
    return OmniRequestState(
        request_id="audio-delta",
        external_req_id="audio-delta",
        parent_req=None,
        request_index=0,
        lora_request=None,
        prompt=None,
        prompt_token_ids=[0],
        prompt_embeds=None,
        logprobs_processor=None,
        detokenizer=None,
        max_tokens_param=None,
        arrival_time=0.0,
        queue=None,
        log_stats=False,
        stream_interval=1,
        output_kind=output_kind,
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("wire_mode", ["on", "off"])
@pytest.mark.parametrize("delta_mode", ["fast", "slow"])
@pytest.mark.parametrize("batch_later_chunks", [False, True])
@pytest.mark.parametrize("finish_with_payload", [False, True])
@pytest.mark.parametrize("explicit_sampling", [False, True])
@pytest.mark.parametrize("lengths", [(5000, 5000, 5000), (0, 5000, 0, 2000, 5000)])
async def test_real_producer_deltas_reach_video_client(
    monkeypatch, wire_mode, delta_mode, batch_later_chunks, finish_with_payload, explicit_sampling, lengths
):
    monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", wire_mode)
    monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", delta_mode)
    chunks = []
    offset = 0
    for length in lengths:
        # Qwen's generation runner hands off [1, samples] CPU tensors.
        chunks.append(torch.arange(offset, offset + length, dtype=torch.float32).reshape(1, -1) / 32768)
        offset += length
    groups = [[chunks[0]], chunks[1:]] if batch_later_chunks else [[chunk] for chunk in chunks]
    state = _state(RequestOutputKind.DELTA)
    snapshots = []
    produced = []
    # Complete production before consumption: retained DELTA snapshots must
    # survive subsequent accumulation and draining under consumer backpressure.
    for index, group in enumerate(groups):
        for chunk in group:
            state.add_multimodal_tensor(chunk, mm_type="audio")
        finish = FinishReason.STOP if finish_with_payload and index == len(groups) - 1 else None
        result = state.make_request_output([], None, finish, None)
        assert result is not None
        assert "audio" not in state.mm_accumulated
        snapshots.append(result.outputs[0].multimodal_output["audio"])
        produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio"))
    if not finish_with_payload:
        result = state.make_request_output([], None, FinishReason.STOP, None)
        assert result is not None
        produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio"))

    class Engine:
        async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None):
            if explicit_sampling:
                assert sampling_params_list is not None
                assert sampling_params_list[0].output_kind == RequestOutputKind.DELTA
            else:
                assert sampling_params_list is None
            # AsyncOmni coerces omitted parameters to DELTA as well.
            for result in produced:
                yield result

    class Handler(QwenOmniStreamingVideoHandler):
        async def _preprocess_to_engine_prompt(self, request):
            return {"prompt": "describe"}

    class WebSocket:
        def __init__(self):
            self.sent = []

        async def send_json(self, message):
            self.sent.append(message)

    ws = WebSocket()
    handler = Handler(chat_service=build_serving_chat(), engine_client=Engine())
    config = StreamingVideoSessionConfig(
        model="test",
        modalities=["text", "audio"],
        sampling_params_list=[{"temperature": 0, "output_kind": 0}] if explicit_sampling else None,
    )
    await handler._process_query_engine(
        ws,
        config,
        [],
        bytearray(),
        [],
        "describe",
        "audio-delta",
        asyncio.Event(),
        {},
    )
    assert not any(message["type"] == "error" for message in ws.sent)
    audio = [message for message in ws.sent if message["type"] == "response.output_audio.delta"]
    nonempty_groups = sum(any(chunk.numel() for chunk in group) for group in groups)
    assert len(audio) == (nonempty_groups if wire_mode == "on" else 1)
    np.testing.assert_array_equal(
        np.concatenate([_pcm(message["data"]) for message in audio]),
        np.arange(_CODEC_FRAME_SAMPLES, sum(lengths)),
    )
    assert sum(message["type"] == "response.output_audio.done" for message in ws.sent) == 1
    assert torch.equal(snapshots[0], chunks[0])
    last = torch.cat(snapshots[-1], dim=-1) if isinstance(snapshots[-1], list) else snapshots[-1]
    assert torch.equal(last, torch.cat(groups[-1], dim=-1))


def test_cumulative_producer_is_a_different_contract():
    # The fix must not change the producer globally: explicit CUMULATIVE users
    # still receive accumulated history, not fresh deltas.
    state = _state(RequestOutputKind.CUMULATIVE)
    chunks = _chunks()
    for chunk in chunks:
        state.add_multimodal_tensor(chunk, mm_type="audio")
        result = state.make_request_output([], None, None, None)
        assert result is not None
        assert "audio" in state.mm_accumulated
    payload = result.outputs[0].multimodal_output["audio"]
    assert isinstance(payload, torch.Tensor)
    assert torch.equal(payload, torch.cat(chunks))