# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Drive complete native duplex turns while Stage-0 repeatedly rebuilds KV."""

from __future__ import annotations

import asyncio
import base64
from pathlib import Path

from vllm_omni.clients.duplex import DuplexClient, EventCollector, read_pcm16_wav, wait_for_condition
from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config


async def run_window_turn(
    *,
    url: str,
    model: str,
    input_wav: Path,
    mode: str,
    ref_audio: Path | None = None,
    repeats: int = 2,
    video_frames: list[str] | None = None,
    buffered_flush: bool = False,
    timeout_s: float = 240.0,
) -> dict[str, object]:
    """Use aggressive limits to exercise multiple windows in one audio turn."""
    pcm = read_pcm16_wav(input_wav) * repeats
    collector = EventCollector()
    reference = (
        "data:audio/wav;base64," + base64.b64encode(ref_audio.read_bytes()).decode("ascii")
        if ref_audio is not None
        else None
    )
    config = create_duplex_session_config(
        modalities=("text", "audio") if ref_audio is not None else ("text",),
        ref_audio=reference,
        temperature=0.0,
        extra_body={
            "sliding_window_mode": mode,
            # Keep a complete ~5.5 s fixture utterance available to the
            # native listen/speak policy. These limits still rebuild KV
            # during the two-repeat input and its response.
            "context_max_units": 8,
            "context_previous_max_tokens": 16,
            "basic_window_high_tokens": 192,
            "basic_window_low_tokens": 128,
        },
    )
    client = DuplexClient(
        url, model=model, config=config, reconnect=None, heartbeat_interval_s=None, handshake_timeout_s=timeout_s
    )
    async with client:
        reader = asyncio.create_task(collector.consume(client))
        collector.add({"type": "session.created", "session": client.session_info})
        try:
            frames_sent = await client.stream_pcm(
                pcm,
                chunk_ms=2500 if buffered_flush else 200,
                realtime=not buffered_flush,
                is_speech=True,
                video_frames=video_frames,
            )
            # The fixture ends in speech. A native duplex commit seals input
            # but does not force an answer; keep the microphone clock running
            # through a pause so the model can choose to speak.
            await client.stream_pcm(bytes(5 * 32000), chunk_ms=200, realtime=True, is_speech=False)
            await client.commit(final=True)
            await wait_for_condition(
                lambda: collector.count("input_audio_buffer.committed") > 0 or bool(collector.errors()),
                timeout_s=timeout_s,
                label="final input commit",
            )
            await wait_for_condition(
                lambda: collector.count("response.done") > 0 or bool(collector.errors()),
                timeout_s=timeout_s,
                label="completed windowed response",
            )
            # Final execution follows the commit acknowledgement asynchronously.
            await asyncio.sleep(3.0)
            await client.close(timeout_s=timeout_s)
            await asyncio.wait_for(reader, timeout=5.0)
        finally:
            if not reader.done():
                reader.cancel()
                await asyncio.gather(reader, return_exceptions=True)
    return {
        "created": collector.count("session.created"),
        "committed": collector.count("input_audio_buffer.committed"),
        "done": collector.count("response.done"),
        "closed": collector.count("session.closed"),
        "errors": collector.errors(),
        "audio_bytes": len(collector.audio_bytes()),
        "frames_sent": frames_sent,
        "input_seconds": len(pcm) / 32000,
    }