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

"""Small media helpers used by the duplex runner and judge.

Video and audio decoding is delegated to PyAV (``av``), which bundles a pinned
FFmpeg version, so results do not depend on whatever ffmpeg binary the host
provides (see #7364: host ffmpeg 4.4.2 could hang forever on some HEVC clips).
"""

from __future__ import annotations

import io
import logging
import os
import tempfile
import wave
from collections.abc import Iterator
from pathlib import Path
from typing import Any

from vllm_omni.experimental.fullduplex.client import PCM16_BYTES_PER_SAMPLE, PCM16_SAMPLE_RATE

logger = logging.getLogger(__name__)


def materialize_media(value: Any, output_dir: str | Path, stem: str, suffix: str) -> Path:
    """Resolve common Hugging Face media values to a local file.

    ``str``/``Path`` values are returned as-is. Inline byte payloads (a
    ``bytes``/``bytearray`` or a Hugging Face ``Audio``/``Video``/``Image``
    feature mapping ``{"bytes": ..., "path": ...}``) are written to a real file
    under ``output_dir`` and that existing path is returned. For such mappings
    the ``path`` key is the original artifact name inside the dataset, which is
    usually *not* present on the local disk, so it is only used to pick a
    content extension and never treated as an existing file. A mapping or
    object carrying only a ``path`` (no inline ``bytes``) is returned as an
    already-materialized file reference.
    """
    if isinstance(value, str | Path):
        return Path(value)
    path = value.get("path") if isinstance(value, dict) else getattr(value, "path", None)
    payload = value.get("bytes") if isinstance(value, dict) else getattr(value, "bytes", None)
    if payload is None and isinstance(value, bytes | bytearray):
        payload = value
    if payload is not None:
        return _write_media_bytes(payload, output_dir, stem, _media_extension(path, suffix))
    if path:
        return Path(path)
    raise ValueError(f"cannot materialize media value for {stem!r}")


def _media_extension(path: str | Path | None, fallback: str) -> str:
    """Pick a safe lower-case file extension from ``path`` when available."""
    extension = Path(str(path)).suffix.lower() if path else ""
    if extension.startswith(".") and len(extension) <= 8 and extension[1:].isalnum():
        return extension
    return fallback


def _write_media_bytes(payload: bytes | bytearray, output_dir: str | Path, stem: str, extension: str) -> Path:
    data = bytes(payload)
    root = Path(output_dir)
    root.mkdir(parents=True, exist_ok=True)
    destination = root / f"{stem}{extension}"
    if destination.exists() and destination.stat().st_size == len(data) and destination.read_bytes() == data:
        return destination
    fd, temporary = tempfile.mkstemp(dir=root, prefix=f".{destination.name}.", suffix=".part")
    try:
        with os.fdopen(fd, "wb") as handle:
            handle.write(data)
        os.replace(temporary, destination)
    except OSError:
        try:
            os.unlink(temporary)
        except OSError:
            pass
        raise
    return destination


def video_duration(path: str | Path) -> float:
    """Return the video duration in seconds using PyAV (no ffprobe subprocess).

    Raises:
        OSError: If the file cannot be opened (missing file, permission denied, etc.).
        ValueError: If the file is corrupt or cannot be parsed (e.g.
            :class:`av.error.InvalidDataError`).
    """
    import av

    try:
        with av.open(str(path)) as container:
            # Container-level duration is the same source ffprobe's format=duration reads.
            if container.duration and container.duration > 0:
                return max(0.0, container.duration / av.time_base)
            if container.streams.video:
                stream = container.streams.video[0]
                if stream.duration is not None and stream.time_base is not None:
                    return max(0.0, float(stream.duration * stream.time_base))
            # Last resort: count decodable frames.
            if container.streams.video:
                stream = container.streams.video[0]
                stream.thread_type = "AUTO"
                frame_count = sum(1 for _ in container.decode(stream))
                if stream.average_rate and frame_count > 0:
                    return frame_count / float(stream.average_rate)
            return 0.0
    except (OSError, ValueError) as exc:
        logger.warning("Failed to read video duration from %s: %s", path, exc)
        raise


def extract_jpeg(path: str | Path, *, timestamp: float, quality: int = 3) -> bytes:
    """Extract a single JPEG frame at ``timestamp`` using PyAV.

    ``quality`` follows the ffmpeg ``-q:v`` convention (lower = better, 1..31)
    and is mapped to the PIL JPEG quality scale (higher = better).

    After ``seek`` the function continues decoding until a frame with PTS >=
    target_pts is found, matching the ``ffmpeg -ss`` semantic of returning the
    first decodable frame *at or after* the requested timestamp.

    Raises:
        ValueError: If the file contains no video stream or no frame could be decoded.
    """
    import av

    with av.open(str(path)) as container:
        if not container.streams.video:
            raise ValueError(f"no video stream in {path}")
        stream = container.streams.video[0]
        stream.thread_type = "AUTO"
        time_base = stream.time_base
        if time_base is None:
            time_base = 1 / 90000
        target_pts = int(max(0.0, timestamp) / time_base)
        container.seek(target_pts, any_frame=False, backward=True, stream=stream)
        closest_frame = None
        for frame in container.decode(stream):
            closest_frame = frame
            if frame.pts is not None and frame.pts >= target_pts:
                break
        if closest_frame is not None:
            image = closest_frame.to_image()
            buf = io.BytesIO()
            image.save(buf, format="JPEG", quality=_ffmpeg_q_to_pil_quality(quality))
            return buf.getvalue()
    raise ValueError(f"no decodable video frames in {path} at timestamp {timestamp}")


def iter_jpegs(
    path: str | Path, *, fps: float = 1.0, duration: float | None = None, quality: int = 3
) -> Iterator[tuple[float, bytes]]:
    end = duration if duration is not None else video_duration(path)
    step = 1.0 / fps
    timestamp = 0.0
    while timestamp < end:
        try:
            frame = extract_jpeg(path, timestamp=timestamp, quality=quality)
        except (OSError, ValueError) as exc:
            logger.warning("Frame extraction failed at %.3fs for %s: %s", timestamp, path, exc)
            break
        if frame:
            yield timestamp, frame
        timestamp += step


def read_audio_pcm16(path: str | Path) -> bytes:
    source = Path(path)
    if source.suffix.lower() == ".wav":
        try:
            with wave.open(str(source), "rb") as wav_file:
                if (
                    wav_file.getnchannels() == 1
                    and wav_file.getsampwidth() == 2
                    and wav_file.getframerate() == PCM16_SAMPLE_RATE
                    and wav_file.getcomptype() == "NONE"
                ):
                    return wav_file.readframes(wav_file.getnframes())
        except (wave.Error, OSError):
            pass
    # Non-WAV audio: decode and resample to mono s16 PCM via PyAV.
    import av
    import numpy as np

    try:
        resampler = av.AudioResampler(format="s16", layout="mono", rate=PCM16_SAMPLE_RATE)
        chunks: list[Any] = []
        with av.open(str(source)) as container:
            for frame in container.decode(audio=0):
                for resampled in resampler.resample(frame):
                    chunks.append(resampled.to_ndarray())
            # Flush the resampler to drain buffered tail samples (P3).
            for resampled in resampler.resample(None):
                chunks.append(resampled.to_ndarray())
        if not chunks:
            return b""
        return np.concatenate([chunk.reshape(-1) for chunk in chunks]).astype(np.int16).tobytes()
    except (OSError, ValueError) as exc:
        logger.warning("Failed to decode audio in %s: %s", source, exc)
        raise


def iter_av_units(
    audio_pcm16: bytes, frames: Iterator[tuple[float, bytes]], *, unit_ms: int = 1000
) -> Iterator[tuple[bytes, bytes | None]]:
    unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * unit_ms // 1000
    frame_iter = iter(frames)
    next_frame = next(frame_iter, None)
    for offset in range(0, len(audio_pcm16), unit_bytes):
        timestamp = offset / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
        frame = None
        if next_frame is not None and next_frame[0] <= timestamp + unit_ms / 1000:
            frame = next_frame[1]
            next_frame = next(frame_iter, None)
        yield audio_pcm16[offset : offset + unit_bytes], frame


def _ffmpeg_q_to_pil_quality(quality: int) -> int:
    """Map ffmpeg ``-q:v`` (1..31, lower = better) to PIL JPEG quality (1..100, higher = better)."""
    ffmpeg_q = max(1, min(31, int(quality)))
    return max(1, min(100, round(95 - (ffmpeg_q - 1) * (93 / 30))))