# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Startup warmup probe for duplex ``/v1/realtime``.

Moved out of ``api_server.py`` under P0.2 of #5227. The
``omni_run_server_worker`` scheduling and ``/v1/realtime`` hold-clients
preamble stay in ``api_server.py`` until P0.3.

Video-required models default to one short audio-plus-frame turn so the four
pipeline stages compile real shapes before a client is accepted. That is
separate from the per-stage JIT kernel registry, which can finish in ~0.01 s
without compiling anything.
"""

import asyncio
import json
import math
import struct
import time
import uuid
from io import BytesIO

import httpx
import pybase64 as base64
from vllm.logger import init_logger

logger = init_logger(__name__)

# First-request compiles (Triton mRoPE, vision pos-embed, Talker torch.compile,
# Code2Wav) landed ~30–40 s after the socket was already "ready" on a cold
# process, then the short reply still has to be spoken. Hold clients past that.
DUPLEX_WARMUP_CLIENT_WAIT_S = 180
DUPLEX_WARMUP_AUDIO_WAIT_S = 150
_WARMUP_SAMPLE_RATE_HZ = 16000
_WARMUP_AUDIO_SECONDS = 0.4
# 640x480 is divisible by Qwen3-VL's patch*merge (32) and above the
# processor's shortest_edge, so the vision tower sees a real camera-sized grid
# rather than an upscaled thumbnail.
_WARMUP_FRAME_WIDTH = 640
_WARMUP_FRAME_HEIGHT = 480
_WARMUP_INSTRUCTIONS = "启动暖机。请只用这四个字回答，不要输出 silent：暖机完成。"


def lookup_duplex_plugin(engine_client: object) -> object | None:
    engine = getattr(engine_client, "engine", None)
    plugin = getattr(engine, "plugin", None) if engine is not None else None
    if plugin is None:
        plugin = getattr(engine_client, "plugin", None)
    return plugin


def startup_warmup_kind(plugin: object | None, warmup_frames: int) -> str | None:
    """Which throwaway realtime session to run before admitting clients.

    ``video_turn`` is one short audio chunk plus one image. It is the default
    for models that require video, including when ``warmup_frames`` is 0.
    ``silent_frames`` is the older audio-only path and stays opt-in.
    A negative ``warmup_frames`` disables both. ``None`` means do not hold clients.
    """
    if warmup_frames < 0:
        return None
    required: object = frozenset()
    if plugin is not None and hasattr(plugin, "capabilities"):
        caps = plugin.capabilities(max_sessions=1)
        required = getattr(caps, "required_input_modalities", frozenset())
    if isinstance(required, frozenset | set) and "video" in required:
        return "video_turn"
    if warmup_frames > 0:
        return "silent_frames"
    return None


def _short_pcm_f32le_b64() -> str:
    """Non-silent PCM so Stage0 ASR actually runs.

    All-zero audio is either rejected or replaced before the ASR forward, which
    would leave the mRoPE kernel uncompiled.
    """
    n = int(_WARMUP_SAMPLE_RATE_HZ * _WARMUP_AUDIO_SECONDS)
    pcm = bytearray()
    for i in range(n):
        sample = 0.2 * math.sin(2.0 * math.pi * 220.0 * i / _WARMUP_SAMPLE_RATE_HZ)
        pcm += struct.pack("<f", sample)
    return base64.b64encode(bytes(pcm)).decode("ascii")


def _warmup_jpeg_b64() -> str:
    from PIL import Image

    image = Image.new("RGB", (_WARMUP_FRAME_WIDTH, _WARMUP_FRAME_HEIGHT), (255, 255, 255))
    image.paste((0, 0, 0), (_WARMUP_FRAME_WIDTH // 2, 0, _WARMUP_FRAME_WIDTH, _WARMUP_FRAME_HEIGHT))
    buf = BytesIO()
    image.save(buf, format="JPEG", quality=70)
    return base64.b64encode(buf.getvalue()).decode("ascii")


async def _warmup_duplex_realtime(app, args, warmup_frames: int) -> None:
    """Run one throwaway realtime session before real clients are admitted.

    Video-required models send one short non-silent audio chunk and one image
    so each stage compiles the shape a real turn uses. Audio-primary models
    keep the opt-in silent-frame path. ``/v1/realtime`` holds other connections
    until this finishes (see ``duplex_warmup_done``).
    """
    warmup_done = getattr(app.state, "duplex_warmup_done", None)
    try:
        try:
            import websockets
        except ImportError:
            logger.warning("Duplex warmup skipped: the 'websockets' package is not installed.")
            return
        engine_client = getattr(app.state, "engine_client", None)
        plugin = lookup_duplex_plugin(engine_client) if engine_client is not None else None
        kind = startup_warmup_kind(plugin, warmup_frames)
        if kind is None:
            return
        served = getattr(args, "served_model_name", None)
        if isinstance(served, list | tuple) and served:
            model_name = served[0]
        elif isinstance(served, str) and served:
            model_name = served
        else:
            model_name = args.model
        # One silence unit as the engine-side plugin defines it (DuplexOmniEngine.plugin).
        frame_samples = int(getattr(plugin, "silence_continuation_samples", 16000))
        native_append = plugin.capabilities(max_sessions=1).supports_core_resumable_request if plugin else True
        from vllm_omni.clients.duplex import build_realtime_url

        url = (
            build_realtime_url(f"ws://127.0.0.1:{args.port}/v1/realtime", model_name, autostart=False)
            + "&vllm_omni_warmup=1"
        )
        if kind == "video_turn":
            logger.info(
                "Duplex warmup starting: one %.2f s audio chunk and one %dx%d frame.",
                _WARMUP_AUDIO_SECONDS,
                _WARMUP_FRAME_WIDTH,
                _WARMUP_FRAME_HEIGHT,
            )
        else:
            logger.info("Duplex warmup starting: %d silent frames of %d samples.", warmup_frames, frame_samples)
        start = time.time()
        # The socket is bound before uvicorn finishes starting, so a plain
        # connect can be accepted and then reset. Wait for /health first.
        async with httpx.AsyncClient() as http:
            for _ in range(120):
                try:
                    if (await http.get(f"http://127.0.0.1:{args.port}/health", timeout=2)).status_code == 200:
                        break
                except httpx.HTTPError:
                    pass
                await asyncio.sleep(1)
            else:
                logger.warning("Duplex warmup: /health never became ready; skipping warmup.")
                return
        session: dict[str, object] = {
            "session_id": f"warmup-{uuid.uuid4().hex[:8]}",
            "model": model_name,
            "modalities": ["audio", "text"],
            "input_audio_format": "pcm_f32le",
            "output_audio_format": "pcm16",
            "idle_timeout_s": 60,
            "turn_detection": None,
            "extra_body": {"auto_response": native_append},
        }
        if kind == "video_turn":
            # The default AURA prompt answers with <|silent|> unless the user
            # asked for something. This session is thrown away; force one short
            # spoken reply so Talker and Code2Wav compile too. auto_response
            # matches the web demo: a commit must actually start a turn.
            session["instructions"] = _WARMUP_INSTRUCTIONS
            session["extra_body"] = {"auto_response": True}
        async with websockets.connect(url, max_size=64 * 1024 * 1024, open_timeout=10) as ws:
            await ws.send(json.dumps({"type": "session.update", "session": session}))
            saw_audio = False
            sent = 0

            async def _recv_until(predicate, timeout_s: float) -> bool:
                deadline = time.monotonic() + timeout_s
                while time.monotonic() < deadline:
                    try:
                        raw = await asyncio.wait_for(ws.recv(), timeout=max(0.05, deadline - time.monotonic()))
                    except (TimeoutError, asyncio.TimeoutError):
                        return False
                    event = json.loads(raw)
                    if predicate(event):
                        return True
                return False

            if not await _recv_until(lambda e: e.get("type") == "session.created", 30):
                logger.warning("Duplex warmup: no session.created within 30 s; aborting warmup.")
                return
            if kind == "video_turn":
                await ws.send(
                    json.dumps(
                        {
                            "type": "input_audio_buffer.append",
                            "audio": _short_pcm_f32le_b64(),
                            "format": "pcm_f32le",
                            "sample_rate_hz": _WARMUP_SAMPLE_RATE_HZ,
                            "is_speech": True,
                            "video_frames": [_warmup_jpeg_b64()],
                        }
                    )
                )
                sent = 1
            else:
                silence = base64.b64encode(bytes(frame_samples * 4)).decode("ascii")
                while sent < warmup_frames:
                    await ws.send(json.dumps({"type": "input_audio_buffer.append", "audio": silence}))
                    sent += 1
                    await asyncio.sleep(0.08)
            if not native_append:
                commit: dict[str, object] = {"type": "input_audio_buffer.commit", "final": True}
                if kind == "video_turn":
                    commit["create_response"] = True
                await ws.send(json.dumps(commit))
            # First output audio means Code2Wav ran, so the Talker compile
            # ahead of it has run too. Do not wait out the whole utterance.
            saw_audio = await _recv_until(
                lambda e: e.get("type") == "response.output_audio.delta",
                DUPLEX_WARMUP_AUDIO_WAIT_S,
            )
            # Close the session EXPLICITLY: a bare websocket close parks the
            # session in its disconnect grace window, where it would keep
            # occupying the max_sessions=1 slot and block the first client.
            await ws.send(json.dumps({"type": "session.close"}))
            await _recv_until(lambda e: e.get("type") == "session.closed", 10)
        if kind == "video_turn":
            logger.info(
                "Duplex warmup finished in %.1f s (one audio+frame turn, first audio %s).",
                time.time() - start,
                "received" if saw_audio else "NOT received",
            )
        else:
            logger.info(
                "Duplex warmup finished in %.1f s (%d silent frames, first audio %s).",
                time.time() - start,
                sent,
                "received" if saw_audio else "NOT received",
            )
    except Exception:
        logger.exception("Duplex warmup failed; continuing to serve (first session pays cold-start costs).")
    finally:
        if warmup_done is not None:
            warmup_done.set()


__all__ = [
    "DUPLEX_WARMUP_CLIENT_WAIT_S",
    "_warmup_duplex_realtime",
    "lookup_duplex_plugin",
    "startup_warmup_kind",
]