path: root/packages/meshbay-node/src/meshbay_node/transport/webrtc/media_tools.py

"""One-shot ffmpeg jobs: whole-file audio transcode, subtitle extraction, seek probe."""

import asyncio
import logging
import os
import tempfile
from pathlib import Path

from meshbay_node import platform

log = logging.getLogger("meshbay_node.transport.webrtc_server")


# A whole audio file is small enough to transcode in one shot rather than
# live-piped like video's fMP4 segments — a few seconds of ffmpeg at most,
# bounded generously so one slow/huge outlier can't pin a transcode slot
# (shared with video, MAX_CONCURRENT_TRANSCODES in webrtc_server.py) indefinitely.
AUDIO_TRANSCODE_TIMEOUT_SECS = 120
# Extracting one subtitle track is a demux and a text conversion, not an
# encode: measured at ~1.2 s for a full film. The bound is generous against a
# pathological container rather than against the work itself, and it is short
# next to the audio one because nothing here decodes a media stream.
SUBTITLE_EXTRACT_TIMEOUT_SECS = 60
# Extracting a subtitle demuxes the whole container, so the cost is set by the
# file and not by the subtitle: measured at **9.8 s per GB** on a library held
# on an external disk — 36 s for a 3.9 GB title, 71 s for a 7.3 GB one. A flat
# 60 s therefore worked on most of a library and failed on the big films, with
# nothing to distinguish that from a broken feature. The allowance is three
# times the measured rate so a slower disk, or one being read by a stream at
# the same time, still finishes.
SUBTITLE_EXTRACT_SECS_PER_GB = 30
SUBTITLE_EXTRACT_TIMEOUT_MAX_SECS = 900


def _subtitle_timeout_for(size_bytes: int) -> float:
    """How long this file is allowed to take. See the constants above."""
    gb = max(0.0, size_bytes) / 1_000_000_000
    return min(SUBTITLE_EXTRACT_TIMEOUT_MAX_SECS,
               max(SUBTITLE_EXTRACT_TIMEOUT_SECS,
                   SUBTITLE_EXTRACT_SECS_PER_GB * gb))


# Probing where a copied seek lands costs 0.06–0.07 s on a real title, so this
# bounds a pathological container rather than the work itself.
SEEK_PROBE_TIMEOUT_SECS = 10
# An index seek backing off further than this did not find a keyframe gap; it
# measured something other than the stream about to be served, and the old
# label beats a fabricated one. The largest real gap seen was under 10 s.
SEEK_PROBE_MAX_BACKOFF_SECS = 60
# A subtitle file is text; a film's is ~96 KB. Anything past this is not a
# subtitle track, it is an ffmpeg that found something else to write, and it
# would sit in the media cache for ever.
SUBTITLE_MAX_BYTES = 8 * 1024 * 1024

# What a whole-file audio transcode may produce. The output is AAC at 192 kbit/s,
# so this is about forty-five minutes of source — past any track, any single
# piece, most sets.
#
# The bound is the media cache's, not memory's. `put_thumb` writes one SQLite row
# and the store is 512 MB with least-recently-used eviction, sized for what it
# holds: thumbnails, posters, subtitle tracks, short transcodes. A three-hour
# audiobook at this bitrate is ~260 MB — a single row that would evict most of
# the cache to make room for itself, and be evicted in turn by the next few
# thumbnails. It is not a size this store can hold usefully.
#
# It does not take away something that worked: `AUDIO_TRANSCODE_TIMEOUT_SECS` is
# 120, so a source long enough to reach this cap was already liable to be killed
# mid-transcode. What changes is that the refusal now says which limit was met.
# Serving audio of that length properly is streaming the transcode rather than
# buffering it, which is a different feature from this one.
AUDIO_TRANSCODE_MAX_BYTES = 64 * 1024 * 1024


def _read_scratch_capped(tmp_path: Path, cap: int, what: str) -> bytes:
    """
    Stat ffmpeg's output, refuse it if it is too big, read it. Blocking.

    Run through `asyncio.to_thread` and not `off_disk`: this file is ffmpeg's
    own, under `tempfile.mkstemp` on the system disk, so it is not a group root
    and there is no spun-down platter to serialise against — it only has to be
    off the event loop. A whole transcode read inline is tens of megabytes of
    blocking read while nothing else in the node is served.

    The size is checked before the bytes are asked for, so an oversized result
    costs a stat rather than the read *and* the memory.
    """
    size = tmp_path.stat().st_size
    if size > cap:
        raise RuntimeError(f"{what} is {size} bytes, over the {cap} cap")
    return tmp_path.read_bytes()


async def _discard_scratch(tmp_path: Path) -> None:
    """Remove one of ffmpeg's temp files, off the loop like the read of it."""
    await asyncio.to_thread(tmp_path.unlink, True)


async def _transcode_audio_to_aac(file_path: Path) -> bytes:
    """
    One-shot, whole-file transcode to AAC in an M4A container — no live
    piping, no seeking, unlike `_stream_video_inner`'s fMP4 segments: a
    WMA/Musepack source here is a few MB at most, so there is nothing to
    gain from streaming it and a real cost to the added complexity
    (fragmented output needs `-movflags empty_moov` and its own
    client-side reassembly). A plain temp file lets ffmpeg write a normal,
    fully-seekable M4A container instead. `-vn` drops any attached-picture
    "video" stream some taggers embed as cover art — without it, ffmpeg's
    mp4 muxer has been seen treating that picture as a video track to
    encode, which is not what this is for; cover art still comes from the
    ordinary embedded/sibling-file path (enrich_audio.py), never from here.
    """
    fd, tmp_name = tempfile.mkstemp(suffix=".m4a")
    os.close(fd)
    tmp_path = Path(tmp_name)
    try:
        proc = await asyncio.create_subprocess_exec(
            platform.ffmpeg_cmd(), "-hide_banner", "-loglevel", "error", "-y",
            "-i", str(file_path),
            "-vn", "-c:a", "aac", "-ac", "2", "-b:a", "192k",
            "-f", "ipod", str(tmp_path),
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.PIPE,
        )
        try:
            _, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=AUDIO_TRANSCODE_TIMEOUT_SECS)
        except TimeoutError:
            proc.kill()
            await proc.wait()
            raise RuntimeError(f"ffmpeg timed out after {AUDIO_TRANSCODE_TIMEOUT_SECS}s")
        if proc.returncode != 0:
            raise RuntimeError(
                f"ffmpeg exited {proc.returncode}: {stderr.decode(errors='replace')[:300]}")
        return await asyncio.to_thread(
            _read_scratch_capped, tmp_path, AUDIO_TRANSCODE_MAX_BYTES,
            "transcoded audio")
    finally:
        await _discard_scratch(tmp_path)


async def _seek_lands_at(file_path: Path, t: float, map_args: list[str]) -> float | None:
    """Where an index seek to `t` actually puts this stream, in source time.

    Measured, not predicted. Copied video can only begin on a keyframe, and
    the obvious way to find that keyframe — scan ffprobe's key frames and take
    the last one at or before `t` — is wrong twice over. Matroska's Cues index
    only some keyframes, so the seek backs off to an indexed one that can be
    much earlier; and the landing point depends on **which streams are
    mapped**, because the container is positioned where every mapped stream
    has data. Measured on a real title: a seek to 4913.7 s landed at 4909.863
    with video alone and at 4907.236 with the second audio track mapped
    alongside it. A prediction from the frame list gave the first number and
    the stream delivered the second, which is 2.65 s of subtitles standing
    away from the voice.

    So ffmpeg is asked instead: the same seek, the same mapping, one copied
    frame, `-copyts` to keep the source's own timestamps, and the answer read
    back off the result. Measured at 0.06–0.07 s, which is cheaper than the
    frame scan it replaces.

    Returns None if anything about the probe fails, and the caller then keeps
    the old label: a number that is wrong by a few seconds is worth much less
    than a stream that does not start.
    """
    fd, tmp_name = tempfile.mkstemp(suffix=".mp4")
    os.close(fd)
    tmp_path = Path(tmp_name)
    try:
        proc = await asyncio.create_subprocess_exec(
            platform.ffmpeg_cmd(), "-hide_banner", "-loglevel", "error", "-y",
            "-copyts", "-noaccurate_seek", "-ss", f"{t:.3f}",
            "-i", str(file_path),
            *map_args, "-c", "copy", "-frames:v", "1",
            "-f", "mp4", str(tmp_path),
            stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL,
        )
        await asyncio.wait_for(proc.wait(), timeout=SEEK_PROBE_TIMEOUT_SECS)
        if proc.returncode != 0:
            return None
        probe = await asyncio.create_subprocess_exec(
            platform.ffprobe_cmd(), "-v", "error",
            "-select_streams", "v:0", "-show_entries", "stream=start_time",
            "-of", "csv=p=0", str(tmp_path),
            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL,
        )
        stdout, _ = await asyncio.wait_for(
            probe.communicate(), timeout=SEEK_PROBE_TIMEOUT_SECS)
    except (TimeoutError, OSError) as e:
        log.warning("stream: seek probe failed at %.1fs: %r", t, e)
        return None
    finally:
        await _discard_scratch(tmp_path)
    text = stdout.decode(errors="replace").strip().rstrip(",")
    try:
        landed = float(text)
    except ValueError:
        return None
    # A seek never lands after what was asked for, and a landing point wildly
    # before it is a probe that measured something else — a chapter track, an
    # attachment. Either way the old label beats a fabricated one.
    if not 0 <= landed <= t + 0.5 or t - landed > SEEK_PROBE_MAX_BACKOFF_SECS:
        log.warning("stream: seek probe at %.1fs answered %.3f — ignored", t, landed)
        return None
    return landed


async def _extract_subtitle_to_webvtt(file_path: Path, ordinal: int,
                                      timeout: float = SUBTITLE_EXTRACT_TIMEOUT_SECS) -> bytes:
    """
    One subtitle track out of a container, whole, as WebVTT.

    Whole-file rather than following the stream, which is what makes the
    result reusable: the cues carry the source's own absolute timestamps, so
    the same extraction serves every seek, every audio-language change and
    every later viewing, and the `<track>` the client attaches never has to be
    rebuilt. It is also the only shape the cache makes sense in — a segment
    keyed on a seek position would be a different blob every time.

    `-map 0:s:<ordinal>` counts subtitle streams (see media_probe.py), and
    `-c:s webvtt` converts subrip/ass to text; a bitmap codec reaching here
    would produce an empty file rather than an error, which is why the caller
    checks membership of the probed text list first and never a range.

    Written to a temp file rather than read off a pipe: the caller wants one
    complete blob to hash and cache, and there is nothing to gain from
    streaming a hundred kilobytes.
    """
    fd, tmp_name = tempfile.mkstemp(suffix=".vtt")
    os.close(fd)
    tmp_path = Path(tmp_name)
    try:
        proc = await asyncio.create_subprocess_exec(
            platform.ffmpeg_cmd(), "-hide_banner", "-loglevel", "error", "-y",
            "-i", str(file_path),
            "-map", f"0:s:{ordinal}", "-c:s", "webvtt",
            "-f", "webvtt", str(tmp_path),
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.PIPE,
        )
        try:
            _, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=timeout)
        except TimeoutError:
            proc.kill()
            await proc.wait()
            raise RuntimeError(f"ffmpeg timed out after {timeout:.0f}s")
        if proc.returncode != 0:
            raise RuntimeError(
                f"ffmpeg exited {proc.returncode}: {stderr.decode(errors='replace')[:300]}")
        blob = await asyncio.to_thread(
            _read_scratch_capped, tmp_path, SUBTITLE_MAX_BYTES, "subtitle track")
        # A WebVTT file that is only its header has no cues in it. That is what
        # a bitmap track extracted by mistake produces, and what a text track
        # whose stream is empty produces; either way there is nothing to show,
        # and an empty track attached to the player is worse than none — it
        # appears in the menu and does nothing when picked.
        if len(blob.strip()) <= len(b"WEBVTT"):
            raise RuntimeError("extracted subtitle contains no cues")
        return blob
    finally:
        await _discard_scratch(tmp_path)