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

"""Synthetic media generation and media/text utilities for tests."""

import atexit
import base64
import concurrent.futures
import hashlib
import io
import logging
import math
import mimetypes
import multiprocessing
import os
import random
import re
import subprocess
import sys
import tempfile
import threading
import time
from concurrent.futures.process import BrokenProcessPool
from contextlib import contextmanager
from pathlib import Path
from typing import Any, Literal, TypedDict, overload

import numpy as np
import soundfile as sf
from PIL import Image

logger = logging.getLogger(__name__)

_MIN_FREE_VRAM = 16 * 1024**3


_synthetic_media_fallback_dir: Path | None = None


def _resolve_synthetic_media_cache_dir(cache_dir: Path | str | None) -> Path:
    if cache_dir is not None:
        return Path(cache_dir).expanduser().resolve()

    default = Path(tempfile.gettempdir()) / "vllm_omni_test_synthetic_media"
    try:
        default.mkdir(parents=True, exist_ok=True)
        # Verify write access: the directory may exist but belong to
        # another user (e.g. a previous CI job), causing PermissionError
        # later when individual files are saved.
        canary = default / ".write_test"
        canary.touch()
        canary.unlink()
        return default
    except (PermissionError, OSError):
        global _synthetic_media_fallback_dir
        if _synthetic_media_fallback_dir is None:
            _synthetic_media_fallback_dir = Path(tempfile.mkdtemp(prefix="vllm_omni_test_synthetic_media_"))
        return _synthetic_media_fallback_dir


def _np_array_from_mp4_bytes(video_bytes: bytes) -> np.ndarray:
    """Decode MP4 bytes to a (T, H, W, 3) uint8 RGB stack (matches in-memory synthetic frames)."""
    import cv2

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
        tmp.write(video_bytes)
        path = tmp.name
    cap = None
    try:
        cap = cv2.VideoCapture(path)
        if not cap.isOpened():
            raise RuntimeError("Failed to open cached synthetic video for decode")
        frames: list[np.ndarray] = []
        while True:
            ok, frame_bgr = cap.read()
            if not ok:
                break
            frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
        if not frames:
            raise RuntimeError("Cached synthetic video has no decodable frames")
        return np.stack(frames, axis=0)
    finally:
        if cap is not None:
            cap.release()
        try:
            os.unlink(path)
        except OSError:
            pass


def generate_synthetic_audio(
    duration: int,
    num_channels: int,
    sample_rate: int = 48000,
    *,
    phrase_text: str = "test",
    force_regenerate: bool = False,
    cache_dir: Path | str | None = None,
) -> dict[str, Any]:
    """
    Generate TTS speech with pyttsx3 and return base64 string.

    Caches the WAV under ``cache_dir`` when given, else under the default temp
    subdirectory. Reuses the file when the same
    ``duration`` / ``num_channels`` / ``sample_rate`` / ``phrase_text`` are
    requested unless ``force_regenerate`` is true.

    The cache filename includes a SHA-256 digest of ``phrase_text`` so different
    phrases never share a WAV cache entry.
    """
    root = _resolve_synthetic_media_cache_dir(cache_dir)
    root.mkdir(parents=True, exist_ok=True)
    phrase_key = hashlib.sha256(phrase_text.encode("utf-8")).hexdigest()
    cache_path = root / f"synth_audio_d{duration}_ch{num_channels}_sr{sample_rate}_pt{phrase_key}.wav"

    if not force_regenerate and cache_path.is_file():
        data, _sr = sf.read(str(cache_path), dtype="float32", always_2d=True)
        audio_bytes = cache_path.read_bytes()
        return {
            "np_array": np.asarray(data, dtype=np.float32),
            "base64": base64.b64encode(audio_bytes).decode("utf-8"),
            "file_path": str(cache_path.resolve()),
        }

    import pyttsx3

    def _pick_voice(engine: pyttsx3.Engine) -> str | None:
        voices = engine.getProperty("voices")
        if not voices:
            return None

        preferred_tokens = (
            "natural",
            "jenny",
            "sonia",
            "susan",
            "zira",
            "aria",
            "hazel",
            "samantha",
            "ava",
            "allison",
            "female",
            "woman",
            "english-us",
            "en-us",
            "english",
        )
        discouraged_tokens = (
            "espeak",
            "robot",
            "mbrola",
            "microsoft david",
            "male",
            "man",
        )

        best_voice = voices[0]
        best_score = float("-inf")
        for voice in voices:
            voice_text = f"{getattr(voice, 'id', '')} {getattr(voice, 'name', '')}".lower()
            voice_languages = " ".join(
                lang.decode(errors="ignore") if isinstance(lang, bytes) else str(lang)
                for lang in getattr(voice, "languages", [])
            ).lower()
            combined_text = f"{voice_text} {voice_languages}"
            score = 0
            for idx, token in enumerate(preferred_tokens):
                if token in combined_text:
                    score += 20 - idx
            for token in discouraged_tokens:
                if token in combined_text:
                    score -= 10
            if "english" in combined_text or "en_" in combined_text or "en-" in combined_text:
                score += 4
            if "en-us" in combined_text or "english-us" in combined_text:
                score += 4
            if score > best_score:
                best_score = score
                best_voice = voice

        return best_voice.id

    def _resample_audio(audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
        if src_sr == dst_sr or len(audio) == 0:
            return audio.astype(np.float32)
        src_len = audio.shape[0]
        dst_len = max(1, int(round(src_len * float(dst_sr) / float(src_sr))))
        src_idx = np.arange(src_len, dtype=np.float32)
        dst_idx = np.linspace(0, src_len - 1, dst_len, dtype=np.float32)
        resampled_channels: list[np.ndarray] = []
        for ch in range(audio.shape[1]):
            resampled_channels.append(np.interp(dst_idx, src_idx, audio[:, ch]).astype(np.float32))
        return np.stack(resampled_channels, axis=1)

    def _match_channels(audio: np.ndarray, target_channels: int) -> np.ndarray:
        current_channels = audio.shape[1]
        if current_channels == target_channels:
            return audio.astype(np.float32)
        if target_channels == 1:
            return np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
        if current_channels == 1:
            return np.repeat(audio, target_channels, axis=1).astype(np.float32)
        collapsed = np.mean(audio, axis=1, keepdims=True, dtype=np.float32)
        return np.repeat(collapsed, target_channels, axis=1).astype(np.float32)

    def _trim_silence(audio: np.ndarray, threshold: float = 0.01) -> np.ndarray:
        if len(audio) == 0:
            return audio
        energy = np.max(np.abs(audio), axis=1)
        voiced = np.where(energy > threshold)[0]
        if len(voiced) == 0:
            return audio
        start = max(0, int(voiced[0]) - int(sample_rate * 0.02))
        end = min(len(audio), int(voiced[-1]) + int(sample_rate * 0.04) + 1)
        return audio[start:end]

    def _enhance_speech(audio: np.ndarray) -> np.ndarray:
        if len(audio) == 0:
            return audio.astype(np.float32)
        enhanced = audio.astype(np.float32).copy()
        enhanced -= np.mean(enhanced, axis=0, keepdims=True, dtype=np.float32)
        if len(enhanced) > 1:
            preemphasis = enhanced.copy()
            preemphasis[1:] = enhanced[1:] - 0.94 * enhanced[:-1]
            enhanced = 0.7 * enhanced + 0.3 * preemphasis
        enhanced = np.sign(enhanced) * np.sqrt(np.abs(enhanced))
        fade = min(len(enhanced) // 4, max(1, int(sample_rate * 0.01)))
        if fade > 1:
            ramp_in = np.linspace(0.0, 1.0, fade, dtype=np.float32)
            ramp_out = np.linspace(1.0, 0.0, fade, dtype=np.float32)
            enhanced[:fade] *= ramp_in[:, None]
            enhanced[-fade:] *= ramp_out[:, None]
        peak = float(np.max(np.abs(enhanced)))
        if peak > 1e-8:
            enhanced = enhanced / peak * 0.95
        return enhanced.astype(np.float32)

    num_samples = int(sample_rate * max(1, duration))
    audio_data = np.zeros((num_samples, num_channels), dtype=np.float32)

    engine = pyttsx3.init()
    engine.setProperty("rate", 112)
    engine.setProperty("volume", 1.0)
    selected_voice = _pick_voice(engine)
    if selected_voice is not None:
        engine.setProperty("voice", selected_voice)

    temp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    temp_wav.close()
    try:
        engine.save_to_file(phrase_text, temp_wav.name)
        engine.runAndWait()
        engine.stop()

        ready = False
        for _ in range(50):
            if os.path.exists(temp_wav.name) and os.path.getsize(temp_wav.name) > 44:
                ready = True
                break
            time.sleep(0.1)
        if not ready:
            raise RuntimeError("pyttsx3 did not produce a WAV file in time.")

        tts_audio, tts_sr = sf.read(temp_wav.name, dtype="float32", always_2d=True)
    finally:
        if os.path.exists(temp_wav.name):
            os.unlink(temp_wav.name)

    if len(tts_audio) == 0:
        raise RuntimeError("pyttsx3 produced an empty WAV file.")

    tts_audio = _resample_audio(tts_audio, tts_sr, sample_rate)
    tts_audio = _match_channels(tts_audio, num_channels)
    tts_audio = _trim_silence(tts_audio, threshold=0.012)
    tts_audio = _enhance_speech(tts_audio)

    lead_silence = min(int(sample_rate * 0.02), num_samples // 8)
    pause_samples = int(sample_rate * 0.18)
    start = lead_silence
    phrase_len = tts_audio.shape[0]
    while start < num_samples:
        take = min(phrase_len, num_samples - start)
        audio_data[start : start + take] = tts_audio[:take]
        start += phrase_len + pause_samples

    max_amp = float(np.max(np.abs(audio_data)))
    if max_amp > 0:
        audio_data = audio_data / max_amp * 0.95

    sf.write(str(cache_path), audio_data, sample_rate, format="WAV", subtype="PCM_16")
    audio_bytes = cache_path.read_bytes()

    return {
        "np_array": audio_data.copy(),
        "base64": base64.b64encode(audio_bytes).decode("utf-8"),
        "file_path": str(cache_path.resolve()),
    }


def _mux_mp4_bytes_with_synthetic_audio(
    video_mp4_bytes: bytes,
    *,
    num_frames: int,
    fps: float = 30.0,
    sample_rate: int = 48000,
) -> bytes:
    duration_sec = num_frames / fps if fps > 0 else 0.0
    duration_int = max(1, int(math.ceil(duration_sec)))

    try:
        audio_result = generate_synthetic_audio(
            duration=duration_int,
            num_channels=1,
            sample_rate=sample_rate,
        )
        audio_pcm = audio_result["np_array"]
    except Exception as e:
        logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e)
        return video_mp4_bytes

    try:
        import imageio_ffmpeg

        ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
    except Exception:
        ffmpeg_exe = "ffmpeg"

    try:
        with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
            vid_path = os.path.join(tmp, "video.mp4")
            wav_path = os.path.join(tmp, "audio.wav")
            out_path = os.path.join(tmp, "out.mp4")
            with open(vid_path, "wb") as f:
                f.write(video_mp4_bytes)
            sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
            cmd = [
                ffmpeg_exe,
                "-y",
                "-nostdin",
                "-hide_banner",
                "-loglevel",
                "error",
                "-i",
                vid_path,
                "-i",
                wav_path,
                "-c:v",
                "copy",
                "-c:a",
                "aac",
                "-b:a",
                "128k",
                "-shortest",
                "-movflags",
                "+faststart",
                out_path,
            ]
            subprocess.run(cmd, check=True, stdin=subprocess.DEVNULL, timeout=300)
            with open(out_path, "rb") as f:
                return f.read()
    except (
        FileNotFoundError,
        subprocess.CalledProcessError,
        subprocess.TimeoutExpired,
        OSError,
    ) as e:
        logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e)
        return video_mp4_bytes


def generate_synthetic_video(
    width: int,
    height: int,
    num_frames: int,
    *,
    embed_audio: bool = False,
    force_regenerate: bool = False,
    cache_dir: Path | str | None = None,
) -> dict[str, Any]:
    """
    Generate synthetic MP4 (optional AAC audio). Caches final bytes by
    ``width`` / ``height`` / ``num_frames`` / ``embed_audio`` unless
    ``force_regenerate`` is true. Cache root: ``cache_dir`` if given, else the
    default temp subdirectory.
    """
    root = _resolve_synthetic_media_cache_dir(cache_dir)
    root.mkdir(parents=True, exist_ok=True)
    cache_path = root / f"synth_video_w{width}_h{height}_nf{num_frames}_ea{int(embed_audio)}.mp4"

    if not force_regenerate and cache_path.is_file():
        video_bytes = cache_path.read_bytes()
        return {
            "np_array": _np_array_from_mp4_bytes(video_bytes),
            "base64": base64.b64encode(video_bytes).decode("utf-8"),
            "file_path": str(cache_path.resolve()),
        }

    import cv2
    import imageio

    class _BouncingBall(TypedDict):
        x: float
        y: float
        vx: float
        vy: float
        radius: int
        color_bgr: tuple[int, int, int]

    num_balls = random.randint(3, 8)
    balls: list[_BouncingBall] = []
    for _ in range(num_balls):
        radius = min(width, height) // 8
        if radius < 1:
            raise ValueError(f"Video dimensions ({width}x{height}) too small")
        x = random.randint(radius, width - radius)
        y = random.randint(radius, height - radius)
        speed = random.uniform(3.0, 8.0)
        angle = random.uniform(0, 2 * math.pi)
        vx = speed * math.cos(angle)
        vy = speed * math.sin(angle)
        color_bgr = (random.randint(50, 255), random.randint(50, 255), random.randint(50, 255))
        balls.append(
            {
                "x": float(x),
                "y": float(y),
                "vx": float(vx),
                "vy": float(vy),
                "radius": radius,
                "color_bgr": color_bgr,
            }
        )

    video_frames = []
    for _ in range(num_frames):
        frame_bgr = np.zeros((height, width, 3), dtype=np.uint8)
        for ball in balls:
            ball["x"] += ball["vx"]
            ball["y"] += ball["vy"]
            if ball["x"] - ball["radius"] <= 0 or ball["x"] + ball["radius"] >= width:
                ball["vx"] = -ball["vx"]
                ball["x"] = max(ball["radius"], min(width - ball["radius"], ball["x"]))
            if ball["y"] - ball["radius"] <= 0 or ball["y"] + ball["radius"] >= height:
                ball["vy"] = -ball["vy"]
                ball["y"] = max(ball["radius"], min(height - ball["radius"], ball["y"]))
            x, y = int(ball["x"]), int(ball["y"])
            radius = int(ball["radius"])
            cv2.circle(frame_bgr, (x, y), radius, ball["color_bgr"], -1)
        frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        video_frames.append(frame_rgb)

    fps = 30
    buffer = io.BytesIO()
    writer_kwargs = {
        "format": "mp4",
        "fps": fps,
        "codec": "libx264",
        "quality": 7,
        "pixelformat": "yuv420p",
        "macro_block_size": 16,
        "ffmpeg_params": ["-preset", "medium", "-crf", "23", "-movflags", "+faststart", "-pix_fmt", "yuv420p"],
    }
    try:
        with imageio.get_writer(buffer, **writer_kwargs) as writer:
            for frame in video_frames:
                writer.append_data(frame)
        buffer.seek(0)
        video_only_bytes = buffer.read()
    except Exception as e:
        print(f"Warning: Failed to encode synthetic video: {e}")
        raise
    video_bytes = (
        _mux_mp4_bytes_with_synthetic_audio(video_only_bytes, num_frames=num_frames, fps=float(fps))
        if embed_audio
        else video_only_bytes
    )

    cache_path.write_bytes(video_bytes)

    return {
        "np_array": np.array(video_frames),
        "base64": base64.b64encode(video_bytes).decode("utf-8"),
        "file_path": str(cache_path.resolve()),
    }


def generate_synthetic_image(
    width: int,
    height: int,
    *,
    force_regenerate: bool = False,
    cache_dir: Path | str | None = None,
    seed: int | None = None,
) -> dict[str, Any]:
    """
    Random colored squares on white background. Caches JPEG by ``width`` /
    ``height`` unless ``force_regenerate`` is true. Cache root: ``cache_dir``
    if given, else the default temp subdirectory.
    """
    if seed is not None:
        random.seed(seed)

    root = _resolve_synthetic_media_cache_dir(cache_dir)
    root.mkdir(parents=True, exist_ok=True)
    cache_path = root / f"synth_image_w{width}_h{height}.jpg"

    if not force_regenerate and cache_path.is_file():
        from PIL import Image as PILImage

        image = PILImage.open(cache_path)
        image.load()
        image_bytes = cache_path.read_bytes()
        return {
            "np_array": np.array(image).copy(),
            "base64": base64.b64encode(image_bytes).decode("utf-8"),
            "file_path": str(cache_path.resolve()),
        }

    from PIL import ImageDraw

    image = Image.new("RGB", (width, height), (255, 255, 255))
    draw = ImageDraw.Draw(image)
    num_squares = random.randint(3, 8)
    for _ in range(num_squares):
        square_size = random.randint(max(1, min(width, height) // 8), max(2, min(width, height) // 4))
        x = random.randint(0, max(0, width - square_size - 1))
        y = random.randint(0, max(0, height - square_size - 1))
        color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
        border_width = random.randint(1, 5)
        draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width)

    image.save(str(cache_path), format="JPEG", quality=85, optimize=True)
    image_bytes = cache_path.read_bytes()

    return {
        "np_array": np.array(image).copy(),
        "base64": base64.b64encode(image_bytes).decode("utf-8"),
        "file_path": str(cache_path.resolve()),
    }


_TEST_ASSETS_ROOT = Path(__file__).resolve().parents[1] / "assets"

_AUDIO_MIME_BY_SUFFIX = {
    ".wav": "audio/wav",
    ".mp3": "audio/mpeg",
    ".flac": "audio/flac",
    ".ogg": "audio/ogg",
}


def _asset_mime_type(path: Path) -> str:
    mime = _AUDIO_MIME_BY_SUFFIX.get(path.suffix.lower())
    if mime is not None:
        return mime
    guessed, _ = mimetypes.guess_type(path.name)
    return guessed or "application/octet-stream"


@overload
def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: Literal[False] = False) -> Path: ...


@overload
def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: Literal[True]) -> str: ...


def get_asset_path(relative_path: str | os.PathLike, *, as_data_url: bool = False) -> Path | str:
    """Resolve a path under ``tests/assets/``.

    When ``as_data_url`` is true, read the file and return a base64 data URL
    suitable for embedding reference media in API requests without external URLs.
    """
    path = _TEST_ASSETS_ROOT / Path(relative_path)
    if not as_data_url:
        return path

    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return f"data:{_asset_mime_type(path)};base64,{encoded}"


def concat_audio(audio_val) -> np.ndarray:
    """Flatten a multimodal audio payload to mono float32 samples.

    Engines return ``multimodal_output["audio"]`` as a tensor, a list of
    per-chunk tensors (streaming decoders), or an array-like; concatenate
    in order and return a 1-D ``np.float32`` array (empty when a list has
    no tensors).
    """
    import torch

    if isinstance(audio_val, list):
        tensors = [t.detach().cpu().float().reshape(-1) for t in audio_val if isinstance(t, torch.Tensor)]
        if not tensors:
            return np.zeros((0,), dtype=np.float32)
        return torch.cat(tensors, dim=-1).numpy().astype(np.float32, copy=False)
    if isinstance(audio_val, torch.Tensor):
        return audio_val.detach().cpu().float().reshape(-1).numpy()
    return np.asarray(audio_val, dtype=np.float32).reshape(-1)


def preprocess_text(text):
    import opencc

    word_normalizations = {
        "zero": "0",
        "one": "1",
        "two": "2",
        "three": "3",
        "four": "4",
        "five": "5",
        "six": "6",
        "seven": "7",
        "eight": "8",
        "nine": "9",
        "ten": "10",
        "kilohertz": "khz",
    }
    for word, replacement in word_normalizations.items():
        pattern = r"\b" + re.escape(word) + r"\b"
        text = re.sub(pattern, replacement, text, flags=re.IGNORECASE)

    text = re.sub(r"[^\w\s]", "", text)
    text = re.sub(r"\s+", " ", text)
    cc = opencc.OpenCC("t2s")
    text = cc.convert(text)
    text = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff])", "", text)
    return text.lower().strip()


def cosine_similarity_text(text1, text2, n: int = 3):
    from collections import Counter

    if not text1 or not text2:
        return 0.0

    text1 = preprocess_text(text1)
    text2 = preprocess_text(text2)
    print(f"cosine similarity text1 is: {text1}, text2 is: {text2}")

    ngrams1 = [text1[i : i + n] for i in range(len(text1) - n + 1)]
    ngrams2 = [text2[i : i + n] for i in range(len(text2) - n + 1)]
    counter1 = Counter(ngrams1)
    counter2 = Counter(ngrams2)

    all_ngrams = set(counter1.keys()) | set(counter2.keys())
    vec1 = [counter1.get(ng, 0) for ng in all_ngrams]
    vec2 = [counter2.get(ng, 0) for ng in all_ngrams]
    dot_product = sum(a * b for a, b in zip(vec1, vec2))
    norm1 = sum(a * a for a in vec1) ** 0.5
    norm2 = sum(b * b for b in vec2) ** 0.5
    if norm1 == 0 or norm2 == 0:
        return 0.0
    cosine = dot_product / (norm1 * norm2)
    # Down-weight when lengths differ: repeated/hallucinated transcripts stay
    # high in bag-of-ngrams cosine (e.g. ABCABCABC vs ABC) but should score low.
    len1, len2 = len(text1), len(text2)
    length_harmony = (2.0 * min(len1, len2)) / (len1 + len2)
    return cosine * length_harmony


class _AudioBuffer:
    """Minimal replacement for pydub.AudioSegment used by test helpers."""

    def __init__(self, data: np.ndarray, sample_rate: int):
        self.data = data
        self.sample_rate = sample_rate

    def export(self, buf: io.BytesIO, format: str = "wav"):
        sf.write(buf, self.data, self.sample_rate, format=format.upper())
        buf.seek(0)


def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer:
    from vllm.multimodal.media.audio import AudioMediaIO

    io_ = AudioMediaIO()
    chunks: list[np.ndarray] = []
    sample_rate: int | None = None
    for b64 in base64_list:
        raw = base64.b64decode(b64.split(",", 1)[-1])
        waveform, sr = io_.load_bytes(raw)
        if sample_rate is None:
            sample_rate = int(sr)
        chunks.append(waveform)
    merged = np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
    return _AudioBuffer(merged, sample_rate or 16000)


@contextmanager
def _serialize_whisper_model_download(model_size: str = "small"):
    """Serialize Whisper cache writes across processes (Linux/Unix), per model."""
    if sys.platform == "win32":
        yield
        return

    import fcntl

    lock_path = Path.home() / ".cache" / "whisper" / f".{model_size}_model_download.lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    f = open(lock_path, "a+b")
    try:
        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
        yield
    finally:
        fcntl.flock(f.fileno(), fcntl.LOCK_UN)
        f.close()


def _select_whisper_device() -> str:
    device_index = None
    from vllm_omni.platforms import current_omni_platform

    if current_omni_platform.is_available():
        n = current_omni_platform.get_device_count()
        # Prefer the highest-index device with enough room for Whisper. The
        # 16 GiB floor protects low-memory single-GPU runners such as the L4
        # nightly, where the model server already occupies device 0, from the
        # server-plus-Whisper OOM fixed in #3822. Keep the CPU fallback when no
        # device has enough room.
        for candidate_index in range(n - 1, -1, -1):
            candidate_device = current_omni_platform.get_torch_device(candidate_index)
            if current_omni_platform.get_free_memory(candidate_device) >= _MIN_FREE_VRAM:
                device_index = candidate_index
                break

    if device_index is None:
        return "cpu"

    torch_device = current_omni_platform.get_torch_device(device_index)
    current_omni_platform.set_device(torch_device)
    return str(torch_device)


# Populated in the transcription worker, not in the pytest process.
_WHISPER_MODELS: dict[str, Any] = {}
# Pinned on the first ``whisper.load_model`` in this worker. Later sizes reuse
# it so a CPU fallback cannot be mixed with a later GPU load in the same pool.
_WHISPER_LOADED_DEVICE: str | None = None


def _get_whisper_model(model_size: str) -> Any:
    global _WHISPER_LOADED_DEVICE
    model = _WHISPER_MODELS.get(model_size)
    if model is None:
        import whisper

        # The device is picked on first load and the model stays on it for the
        # worker's lifetime: the current server or runner fixture instance, or
        # the test module for callers that transcribe without those fixtures.
        device = _WHISPER_LOADED_DEVICE
        if device is None:
            device = _select_whisper_device()
            _WHISPER_LOADED_DEVICE = device
        with _serialize_whisper_model_download(model_size):
            model = whisper.load_model(model_size, device=device)
        _WHISPER_MODELS[model_size] = model
    return model


def _whisper_process_reserved_gib() -> float:
    """GiB held by this worker's caching allocator, or 0 on CPU / if unreadable.

    Parent-side ``nvidia-smi --query-compute-apps`` is empty on some CI
    drivers (pmon shows ``-`` while device-wide used is tens of GiB), so the
    pytest process cannot map the worker PID to VRAM. The allocator in *this*
    process still knows what it reserved.
    """
    index = _accelerator_index_from_device(_WHISPER_LOADED_DEVICE)
    if index is None:
        return 0.0
    try:
        from vllm_omni.platforms import current_omni_platform

        if not current_omni_platform.is_available():
            return 0.0
        device = current_omni_platform.get_torch_device(index)
        return current_omni_platform.memory_reserved(device) / 1024**3
    except Exception:
        return 0.0


def _whisper_transcribe_in_current_process(
    output_path: str, model_size: str = "small", language: str | None = None
) -> tuple[str, str, float]:
    model = _get_whisper_model(model_size)
    text = model.transcribe(
        output_path,
        temperature=0.0,
        word_timestamps=True,
        condition_on_previous_text=False,
        # None keeps whisper's auto-detection. Do not default this to a
        # language: callers include non-English audio tests.
        language=language,
    )["text"]
    device = _WHISPER_LOADED_DEVICE or "cpu"
    reserved = _whisper_process_reserved_gib() if _accelerator_index_from_device(device) is not None else 0.0
    return text or "", device, reserved


# Serializes a whole submit->result->cleanup on the parent side, so at most one
# call is ever in flight on the single-worker pool. The child already runs
# transcriptions serially, so this costs no throughput; it lets a failed call
# discard the worker without disrupting another caller's in-flight future.
_TRANSCRIBER_CALL_LOCK = threading.Lock()
# Guards the _TRANSCRIBER pointer itself.
_TRANSCRIBER_LOCK = threading.Lock()
_TRANSCRIBER: concurrent.futures.ProcessPoolExecutor | None = None
# Parent-side record of the device the child reported and its allocator
# reserved GiB. Credit only after a GPU result (not while the model might
# still be on CPU). Reserved is the child's caching-allocator reading, not
# nvidia-smi PID memory (N/A on some CI drivers) and not a size table.
_TRANSCRIBER_DEVICE: str | None = None
_TRANSCRIBER_RESERVED_GIB: float = 0.0


def _accelerator_index_from_device(device: str | None) -> int | None:
    """Logical device index for ``cuda:1`` / ``npu:0``; ``None`` for CPU or unknown."""
    if device is None:
        return None
    name = device.strip().lower()
    if not name or name == "cpu":
        return None
    if ":" in name:
        suffix = name.rsplit(":", 1)[1]
        if suffix.isdigit():
            return int(suffix)
    return None


def whisper_resident_device_index() -> int | None:
    """Logical accelerator index holding Whisper, or ``None`` if it is not on GPU."""
    with _TRANSCRIBER_LOCK:
        if _TRANSCRIBER is None:
            return None
        device = _TRANSCRIBER_DEVICE
    return _accelerator_index_from_device(device)


def whisper_resident_vram_gib() -> float:
    """VRAM (GiB) held by the living Whisper worker, or 0 if none / CPU.

    Only the child's caching-allocator reading. Unmeasured or CPU → 0 (fail
    closed: do not invent a size-table credit). This is not a substitute for
    engine PID reap or a raised 5% wait threshold (RFC #6851).
    """
    with _TRANSCRIBER_LOCK:
        if _TRANSCRIBER is None:
            return 0.0
        device = _TRANSCRIBER_DEVICE
        measured = _TRANSCRIBER_RESERVED_GIB
    if _accelerator_index_from_device(device) is None:
        return 0.0
    return measured if measured > 0.0 else 0.0


def _get_transcriber() -> concurrent.futures.ProcessPoolExecutor:
    global _TRANSCRIBER
    with _TRANSCRIBER_LOCK:
        if _TRANSCRIBER is None:
            ctx = multiprocessing.get_context("spawn")
            _TRANSCRIBER = concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx)
        return _TRANSCRIBER


def _discard_transcriber(executor: concurrent.futures.ProcessPoolExecutor) -> None:
    """Drop ``executor``, but only while it is still the current one.

    Identity-checked so a stale reference can never shut down a newer worker that
    was installed after ``executor`` was replaced.
    """
    global _TRANSCRIBER, _TRANSCRIBER_DEVICE, _TRANSCRIBER_RESERVED_GIB
    with _TRANSCRIBER_LOCK:
        if _TRANSCRIBER is not executor:
            return
        _TRANSCRIBER = None
        _TRANSCRIBER_DEVICE = None
        _TRANSCRIBER_RESERVED_GIB = 0.0
    # Joining the worker can block; do it outside the lock.
    executor.shutdown(wait=True)


def release_audio_transcriber() -> None:
    """Shut the transcription worker down, freeing the device memory its models hold.

    Called when a server or runner fixture instance tears down, so that the next
    one -- including the next parametrization inside the same test module -- does
    not initialize its model while Whisper still occupies the device. A module
    teardown fixture repeats it for tests that transcribe without those fixtures.

    Takes the call lock, so it waits for any in-flight transcription to finish
    rather than shutting the worker down underneath it.
    """
    global _TRANSCRIBER, _TRANSCRIBER_DEVICE, _TRANSCRIBER_RESERVED_GIB
    with _TRANSCRIBER_CALL_LOCK:
        with _TRANSCRIBER_LOCK:
            executor, _TRANSCRIBER = _TRANSCRIBER, None
            _TRANSCRIBER_DEVICE = None
            _TRANSCRIBER_RESERVED_GIB = 0.0
        if executor is not None:
            executor.shutdown(wait=True)


def _unpack_transcribe_worker_result(result: object) -> tuple[str, str, float]:
    """Accept ``(text, device)`` from tests or ``(text, device, reserved_gib)`` from the worker."""
    if not isinstance(result, tuple) or len(result) not in (2, 3):
        raise TypeError(f"whisper worker returned {type(result)!r}, expected a 2- or 3-tuple")
    text, device = result[0], result[1]
    reserved = float(result[2]) if len(result) == 3 else 0.0
    return str(text), str(device), reserved


def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str:
    """Convert an audio file to text in a reused, isolated subprocess.

    The worker outlives the call so its Whisper model is loaded once rather than
    once per transcription. The call lock serializes callers onto the single
    worker (the child already transcribes serially, so this costs no throughput),
    which lets a failed call tear the worker down without racing another caller.
    The worker caches one model per size it is asked for, so a run that escalates
    to a stronger ASR keeps both models resident until release.

    Any failure discards the worker, restoring the failure isolation of the old
    one-process-per-call design: a failure (a ``torch`` OOM, or a
    ``KeyboardInterrupt``/``SystemExit`` transported through the future or
    raised while ``result()`` blocks) propagates unchanged after the worker --
    and its resident model -- is torn down, and a dead worker
    (``BrokenProcessPool``) is additionally retried once.
    """
    global _TRANSCRIBER_DEVICE, _TRANSCRIBER_RESERVED_GIB
    with _TRANSCRIBER_CALL_LOCK:
        for attempt in range(2):
            executor = _get_transcriber()
            try:
                text, device, reserved_gib = _unpack_transcribe_worker_result(
                    executor.submit(_whisper_transcribe_in_current_process, output_path, model_size, language).result()
                )
                with _TRANSCRIBER_LOCK:
                    _TRANSCRIBER_DEVICE = device
                    if _accelerator_index_from_device(device) is None:
                        _TRANSCRIBER_RESERVED_GIB = 0.0
                    elif reserved_gib > 0.0:
                        _TRANSCRIBER_RESERVED_GIB = reserved_gib
                return text
            except BrokenProcessPool:
                _discard_transcriber(executor)
                if attempt == 1:
                    raise
            except BaseException:
                # Any other failure -- a task exception (a torch OOM included), or
                # a KeyboardInterrupt/SystemExit transported through the future or
                # interrupting result() -- leaves the worker and its model
                # resident; drop it so it cannot contaminate later calls, matching
                # the old per-call teardown. Do not retry.
                _discard_transcriber(executor)
                raise
    raise AssertionError("unreachable")


def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small", language: str | None = None) -> str:
    output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav")
    os.close(output_fd)
    if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"):
        atexit.register(Path(output_path).unlink, missing_ok=True)
    data, samplerate = sf.read(io.BytesIO(raw_bytes))
    sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16")
    print(f"audio data is saved: {output_path}")
    return convert_audio_file_to_text(output_path, model_size, language)


__all__ = [
    "_merge_base64_audio_to_segment",
    "concat_audio",
    "convert_audio_bytes_to_text",
    "convert_audio_file_to_text",
    "cosine_similarity_text",
    "generate_synthetic_audio",
    "generate_synthetic_image",
    "generate_synthetic_video",
    "get_asset_path",
    "preprocess_text",
    "release_audio_transcriber",
    "whisper_resident_device_index",
    "whisper_resident_vram_gib",
]