Blame for processor/app/browser_preview.py

"""Browser-compatible video preview cache.

The collector and LeRobot dataset may use HEVC (including 12-bit depth
streams), but the Review page must not depend on the browser being able to
decode those codecs. This module creates a local, derived H.264 preview on
first use and reuses it for subsequent range requests.

The source file is never changed. The preview is a disposable local cache;
the authoritative recording and the lossless metric-depth asset remain in
the configured storage directory.
"""

from __future__ import annotations

import hashlib
import json
import os
import shutil
import subprocess
import threading
import time
from pathlib import Path

from app.config import settings


_LOCKS: dict[str, threading.Lock] = {}
_LOCKS_GUARD = threading.Lock()

_FASTSTART_CACHE: dict[str, bool] = {}
_FASTSTART_CACHE_LOCK = threading.Lock()


def _cache_root() -> Path:
    root = settings.upload_staging_root / "browser-preview"
    root.mkdir(parents=True, exist_ok=True)
    return root


def _cache_key(source: Path) -> str:
    stat = source.stat()
    fingerprint = f"{source.resolve()}:{stat.st_size}:{stat.st_mtime_ns}"
    return hashlib.sha256(fingerprint.encode("utf-8")).hexdigest()


def _lock_for(key: str) -> threading.Lock:
    with _LOCKS_GUARD:
        return _LOCKS.setdefault(key, threading.Lock())


# ── Local source mirror (depth streams) ──────────────────────────────
# Depth videos live on the SSHFS mount, where random reads for seek/decode
# cost seconds.  A bit-identical local copy turns every later decode into a
# local-disk read.  This is a disposable transport cache of the SOURCE file
# (never a derived/colorized asset); the authoritative stream stays in the
# configured storage directory.

_MIRROR_MAX_BYTES = 8 * 1024 ** 3
_MIRROR_INFLIGHT: set[str] = set()
_MIRROR_INFLIGHT_GUARD = threading.Lock()


def _mirror_root() -> Path:
    root = settings.upload_staging_root / "depth-mirror"
    root.mkdir(parents=True, exist_ok=True)
    return root


def _prune_mirror(keep: Path) -> None:
    """Remove oldest mirror files beyond the byte budget (LRU by mtime)."""
    try:
        files = [p for p in _mirror_root().iterdir()
                 if p.is_file() and p.suffix == ".mp4"]
    except OSError:
        return
    total = sum(p.stat().st_size for p in files if p != keep) + (
        keep.stat().st_size if keep.is_file() else 0)
    if total <= _MIRROR_MAX_BYTES:
        return
    for path in sorted(files, key=lambda p: p.stat().st_mtime_ns):
        if path == keep or total <= _MIRROR_MAX_BYTES:
            continue
        try:
            total -= path.stat().st_size
            path.unlink(missing_ok=True)
        except OSError:
            pass


def _copy_mirror(source: Path, destination: Path, key: str) -> None:
    """Background copy to a temp file, then atomic rename."""
    temporary = destination.with_suffix(".part")
    temporary.unlink(missing_ok=True)
    try:
        shutil.copy2(source, temporary)
        if not temporary.is_file() or temporary.stat().st_size <= 0:
            raise RuntimeError("mirror copy produced an empty file")
        os.replace(temporary, destination)
        _prune_mirror(destination)
    except (OSError, RuntimeError):
        temporary.unlink(missing_ok=True)
    finally:
        with _MIRROR_INFLIGHT_GUARD:
            _MIRROR_INFLIGHT.discard(key)


def _write_turbo_cube(path: Path) -> None:
    """Write the static 4096-entry 1D Turbo cube used by the transcode.

    4096 entries keep the range mapping at 12-bit precision: a 256-entry LUT
    evaluated after an 8-bit intermediate caused visible banding (only ~58
    distinct color levels survived the chain instead of ~173).
    """
    if path.is_file():
        return
    from app.lerobot_v21 import turbo_lut

    lut = turbo_lut(4096)
    lines = ["TITLE Turbo", "LUT_1D_SIZE 4096", "",
             "DOMAIN_MIN 0 0 0", "DOMAIN_MAX 1 1 1", ""]
    for r, g, b in lut:
        lines.append(f"{r / 255.0:.6f} {g / 255.0:.6f} {b / 255.0:.6f}")
    path.write_text("\n".join(lines) + "\n", encoding="ascii")


def _depth_stats_path(source: Path) -> Path:
    key = _cache_key(source)
    return _cache_root() / f"depth-{key}.stats.json"


def read_depth_stats(source: Path) -> dict | None:
    """Disk-persisted q01/q99 bounds for ``source`` (same fingerprint key).

    Written once by :func:`_ensure_depth_stats` during preview generation;
    survives restarts so neither the stats endpoint nor a re-transcode ever
    pays a second full-stream decode for the same source file.
    """
    try:
        path = _depth_stats_path(source)
        if not path.is_file():
            return None
        data = json.loads(path.read_text(encoding="utf-8"))
        if (isinstance(data, dict)
                and isinstance(data.get("q01"), (int, float))
                and isinstance(data.get("q99"), (int, float))
                and float(data["q99"]) > float(data["q01"])):
            return data
    except (OSError, ValueError):
        pass
    return None


def _histogram_depth_stats(source: Path) -> dict | None:
    """Single-decode q01/q99 code bounds via a streaming 4096-bin histogram.

    The stream is read as raw gray12le codes and histogrammed in O(1)
    memory; percentiles are exact over every decoded pixel.  The previous
    ``select=`` sampling pass decoded the whole stream anyway while keeping
    only a fraction of the frames, so nothing is lost.
    """
    import numpy as np

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        return None
    hist = np.zeros(4096, dtype=np.int64)
    proc = subprocess.Popen(
        [ffmpeg, "-v", "error", "-i", str(source), "-map", "0:v:0",
         "-f", "rawvideo", "-pix_fmt", "gray12le", "-"],
        stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    try:
        while True:
            chunk = proc.stdout.read(8 * 1024 * 1024)
            if not chunk:
                break
            values = np.frombuffer(chunk, dtype="<u2")
            hist += np.bincount(values.astype(np.int64), minlength=4096)[:4096]
    finally:
        try:
            proc.stdout.close()
        except OSError:
            pass
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
    total = int(hist.sum())
    if total < 2:
        return None
    cumulative = np.cumsum(hist)
    q01 = float(np.searchsorted(cumulative, total * 0.01, side="right"))
    q99 = float(np.searchsorted(cumulative, total * 0.99, side="right"))
    return {"q01": q01, "q99": q99, "total_pixels": total}


def _ensure_depth_stats(source: Path, local: Path | None = None) -> dict | None:
    """Return q01/q99 for ``source``, computed once and persisted to disk.

    ``source`` is the authoritative path used for the cache fingerprint;
    ``local`` (the transport mirror) is decoded when it is ready so the pass
    never competes with an SSHFS random read.
    """
    cached = read_depth_stats(source)
    if cached is not None:
        return cached
    with _lock_for(f"stats:{_cache_key(source)}"):
        cached = read_depth_stats(source)
        if cached is not None:
            return cached
        stats = _histogram_depth_stats(local or source)
        if stats is None:
            return None
        try:
            _depth_stats_path(source).write_text(
                json.dumps(stats), encoding="utf-8")
        except OSError:
            pass
        return stats


def depth_preview_exists(source: Path) -> bool:
    """Cheap cache-hit check for prewarm scheduling (no transcode)."""
    try:
        destination = _cache_root() / f"depth-{_cache_key(Path(source))}.mp4"
        return destination.is_file() and destination.stat().st_size > 0
    except OSError:
        return False


def ensure_depth_h264_preview(source: Path, q01: float = 0.0,
                              q99: float = 4095.0) -> Path:
    """Generate (once) an 8-bit Turbo-colored H.264 preview of a depth stream.

    The browser display pipeline already quantizes codes into an 8-bit Turbo
    ramp; this bakes the identical mapping (same clamp/scale math, same
    colormap) into a disposable local MP4 so the review page plays it through
    a native ``<video>`` element — instant open, native seek, no per-open
    12-bit decode or 229MB raw-code transfer.  The preview never enters the
    dataset; the canonical gray12le stream stays authoritative.
    """
    source = Path(source)
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg is required to build depth previews")
    key = _cache_key(source)
    destination = _cache_root() / f"depth-{key}.mp4"
    if destination.is_file() and destination.stat().st_size > 0:
        return destination
    with _lock_for(key):
        if destination.is_file() and destination.stat().st_size > 0:
            return destination
        q01 = float(q01 if q01 is not None else 0.0)
        q99 = float(q99 if q99 is not None else 4095.0)
        span = max(q99 - q01, 1.0)
        # Decode from the local mirror when it is ready; SSHFS random reads
        # during the stats/transcode passes otherwise dominate generation.
        # Wait briefly for the background copy — a ~6s copy is cheaper than
        # decoding the whole stream over the mount.
        original = source
        try:
            source = ensure_local_mirror(source, wait_seconds=45)
        except Exception:
            pass
        if q01 <= 0 and q99 >= 4095:
            # No cached adaptive bounds: compute them once from a single
            # decode (persisted next to the preview, keyed on the
            # authoritative source path) so the ramp matches the former
            # q01/q99 colorization and never re-decodes for the same file.
            stats = _ensure_depth_stats(original, source)
            if stats and float(stats.get("q99") or 0) > float(stats.get("q01") or 0):
                q01 = float(stats["q01"])
                q99 = float(stats["q99"])
            span = max(q99 - q01, 1.0)
        # Map the range in the native 12-bit domain (the lut filter evaluates
        # ``val`` at the input depth), then apply a 4096-entry Turbo cube:
        # quantizing through an 8-bit intermediate produced visible banding.
        cube = _cache_root() / "turbo.cube"
        _write_turbo_cube(cube)
        temporary = destination.with_name(
            f".{destination.stem}.{os.getpid()}.part.mp4")
        temporary.unlink(missing_ok=True)
        try:
            subprocess.run(
                [ffmpeg, "-y", "-loglevel", "error", "-i", str(source),
                 "-vf", (f"lut=y='clip((val-{q01:.2f})*4095/{span:.2f},0,4095)',"
                         f"lut1d=file={cube},format=yuv420p"),
                 "-an", "-c:v", "libx264", "-preset", "veryfast",
                 "-crf", "18", "-g", "25", "-keyint_min", "25",
                 "-sc_threshold", "0", "-movflags", "+faststart",
                 str(temporary)],
                check=True, capture_output=True, timeout=3600,
            )
            if not temporary.is_file() or temporary.stat().st_size == 0:
                raise RuntimeError("ffmpeg produced an empty depth preview")
            os.replace(temporary, destination)
        except subprocess.TimeoutExpired as exc:
            raise RuntimeError("depth preview transcode timed out") from exc
        except subprocess.CalledProcessError as exc:
            detail = (exc.stderr or b"").decode("utf-8", errors="replace")[-500:]
            raise RuntimeError(f"depth preview transcode failed: {detail}") from exc
        finally:
            temporary.unlink(missing_ok=True)
    return destination


def ensure_local_mirror(source: Path, wait_seconds: float = 0.0) -> Path:
    """Return a local copy of ``source`` when ready, else ``source``.

    The first request for a video schedules an asynchronous copy and keeps
    serving from the original path; once the mirror lands (atomic rename),
    every later call — including the next frame of the same request stream —
    returns the local path.  The key includes size+mtime, so a reprocessed
    MP4 silently invalidates the old mirror.

    ``wait_seconds > 0`` polls for the copy to land before falling back to
    the remote path (used by preview generation, whose decode pass is much
    faster from local disk than over SSHFS).
    """
    source = Path(source)
    key = _cache_key(source)
    destination = _mirror_root() / f"{key}.mp4"
    if destination.is_file() and destination.stat().st_size > 0:
        return destination
    with _MIRROR_INFLIGHT_GUARD:
        if key not in _MIRROR_INFLIGHT:
            _MIRROR_INFLIGHT.add(key)
            threading.Thread(
                target=_copy_mirror, args=(source, destination, key),
                daemon=True,
            ).start()
    if wait_seconds > 0:
        deadline = time.monotonic() + wait_seconds
        while time.monotonic() < deadline:
            if destination.is_file() and destination.stat().st_size > 0:
                return destination
            time.sleep(0.25)
        return destination if destination.is_file() else source
    return source


def _is_h264_yuv420(source: Path) -> bool:
    """Avoid re-encoding an already browser-compatible source."""
    ffprobe = shutil.which("ffprobe")
    if not ffprobe:
        return False
    try:
        result = subprocess.run(
            [ffprobe, "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=codec_name,pix_fmt",
             "-of", "csv=p=0", str(source)],
            check=True, capture_output=True, text=True, timeout=15,
        )
        codec, _, pix_fmt = result.stdout.strip().partition(",")
        return codec.lower() == "h264" and pix_fmt in {"yuv420p", "yuvj420p"}
    except (OSError, subprocess.SubprocessError, ValueError):
        return False


def _probe_faststart(source: Path) -> bool:
    """moov 在 mdat 之前 = faststart(浏览器可流式播放,否则播放会早停)。

    只读前几个顶层 box 的头部,moov/mdat 先后顺序即可判定;探测失败
    保守返回 True(透传原文件,维持旧行为)。结果按源文件指纹记忆。
    """
    key = _cache_key(source)
    with _FASTSTART_CACHE_LOCK:
        cached = _FASTSTART_CACHE.get(key)
        if cached is not None:
            return cached
    result = True
    try:
        with open(source, "rb") as handle:
            handle.seek(0, os.SEEK_END)
            size = handle.tell()
            handle.seek(0)
            for _ in range(8):
                if handle.tell() > size - 8:
                    break
                pos = handle.tell()
                header = handle.read(8)
                if len(header) < 8:
                    break
                length = int.from_bytes(header[:4], "big")
                box_type = header[4:8]
                if length == 1:
                    ext = handle.read(8)
                    length = int.from_bytes(ext, "big") if len(ext) == 8 else 0
                if length == 0:
                    length = size - pos
                if length < 8:
                    break
                if box_type == b"moov":
                    result = True
                    break
                if box_type == b"mdat":
                    result = False
                    break
                handle.seek(pos + length)
    except OSError:
        result = True
    with _FASTSTART_CACHE_LOCK:
        _FASTSTART_CACHE[key] = result
    return result


def _remux_faststart(source: Path, destination: Path) -> None:
    """``-c copy -movflags +faststart`` 容器重排(不重编码,字节级无损)。

    moov 索引移到文件头,浏览器无需跳到文件尾即可开始流式播放。
    """
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg is required for faststart remux")
    temporary = destination.with_name(
        f".{destination.stem}.{os.getpid()}.part.mp4")
    temporary.unlink(missing_ok=True)
    try:
        subprocess.run(
            [ffmpeg, "-y", "-loglevel", "error", "-i", str(source),
             "-map", "0:v:0", "-an", "-c:v", "copy",
             "-movflags", "+faststart", str(temporary)],
            check=True, capture_output=True, timeout=600,
        )
        if not temporary.is_file() or temporary.stat().st_size == 0:
            raise RuntimeError("faststart remux produced an empty file")
        os.replace(temporary, destination)
    except subprocess.TimeoutExpired as exc:
        raise RuntimeError("faststart remux timed out") from exc
    except subprocess.CalledProcessError as exc:
        detail = (exc.stderr or b"").decode("utf-8", errors="replace")[-300:]
        raise RuntimeError(f"faststart remux failed: {detail}") from exc
    finally:
        temporary.unlink(missing_ok=True)


def ensure_h264_preview(source: Path) -> Path:
    """Return an H.264/yuv420p MP4 suitable for HTML5 playback.

    Browser-decodable sources are passed through directly when their moov
    index sits at the front (faststart); non-faststart files are remuxed
    once into the local cache (``-c copy``, byte-lossless, only the
    container box order changes) so HTML5 streaming does not stall a few
    seconds into playback.  Remux failure falls back to the original file.

    Callers should run this function via ``asyncio.to_thread`` so a cold
    transcode never blocks FastAPI's event loop.
    """
    source = Path(source)
    if not source.is_file():
        raise FileNotFoundError(f"Video source not found: {source}")

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg is required to build browser video previews")

    key = _cache_key(source)
    destination = _cache_root() / f"{key}.mp4"
    if destination.is_file() and destination.stat().st_size > 0:
        return destination
    if _is_h264_yuv420(source) and _probe_faststart(source):
        return source

    with _lock_for(key):
        if destination.is_file() and destination.stat().st_size > 0:
            return destination
        if _is_h264_yuv420(source):
            # 非 faststart 的 h264:一次性无损重排 moov 到文件头。
            # remux 失败回退原文件,绝不阻断播放。
            try:
                _remux_faststart(source, destination)
                if destination.is_file() and destination.stat().st_size > 0:
                    return destination
            except Exception:
                pass
            return source

        # Keep the final .mp4 suffix so ffmpeg can infer the output muxer.
        temporary = destination.with_name(
            f".{destination.stem}.{os.getpid()}.part.mp4")
        temporary.unlink(missing_ok=True)
        try:
            subprocess.run(
                [ffmpeg, "-y", "-loglevel", "error", "-i", str(source),
                 "-map", "0:v:0", "-an",
                 "-c:v", "libx264", "-preset", "veryfast",
                 "-pix_fmt", "yuv420p",
                 # Short GOP keeps frame stepping and range seeks responsive.
                 "-g", "25", "-keyint_min", "25", "-sc_threshold", "0",
                 "-movflags", "+faststart", str(temporary)],
                check=True, capture_output=True, timeout=3600,
            )
            if not temporary.is_file() or temporary.stat().st_size == 0:
                raise RuntimeError("ffmpeg produced an empty browser preview")
            os.replace(temporary, destination)
        except subprocess.TimeoutExpired as exc:
            raise RuntimeError("H.264 preview transcode timed out") from exc
        except subprocess.CalledProcessError as exc:
            detail = (exc.stderr or b"").decode("utf-8", errors="replace")[-500:]
            raise RuntimeError(f"H.264 preview transcode failed: {detail}") from exc
        finally:
            temporary.unlink(missing_ok=True)
    return destination