"""Sparse per-substep MP4 recorder for ``deploy.py`` evals.

One MP4 per rollout (rollout boundary = ``h`` key in deploy.py); all
the MP4s for one eval session live in a single directory::

    <eval_videos_root>/<eval_name>_<timestamp>/
    ├── vid_001.mp4
    ├── vid_002.mp4
    └── ...

The recorder is just an MP4 writer with a per-rollout file counter.
**Sparse**: the caller hands it BGR frames at meaningful moments
(start of each rollout + after each ``robot.move_to`` substep), and
the recorder writes them straight to disk. No background thread, no
shared ZED state.

Default ``fps=5`` makes the result look like a stop-motion of the
robot's trajectory; raise it if you want slower playback.
"""
from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np


class VideoRecorder:
    # `mp4v` writes a DivX-era codec that OpenCV always has via FFmpeg, but
    # browsers cannot play it. `avc1` (H.264) plays in browsers but OpenCV
    # often picks `h264_v4l2m2m` (hardware encoder) under FFmpeg, which fails
    # to init on most machines (no V4L2-M2M). We keep mp4v at write-time for
    # reliability and run a `transcode_to_h264()` finalize pass at end of
    # session (uses the host's `ffmpeg` binary with libx264 + faststart so
    # the MP4 is web-streamable). See ``transcode_session_to_h264`` below.
    def __init__(self, out_dir: Path, W: int, H: int, fps: int = 5,
                 fourcc: str = "mp4v"):
        self.out_dir = Path(out_dir)
        self.out_dir.mkdir(parents=True, exist_ok=True)
        self.W = int(W); self.H = int(H)
        self.fps = int(fps)
        self.fourcc = cv2.VideoWriter_fourcc(*fourcc)
        self._idx = 0
        self._writer: cv2.VideoWriter | None = None
        self._frames_in_current = 0
        self.rotate()                                    # opens vid_001 immediately

    @property
    def current_path(self) -> Path:
        return self.out_dir / f"vid_{self._idx:03d}.mp4"

    def append(self, bgr: np.ndarray) -> None:
        """Write one frame to the current rollout's MP4."""
        if self._writer is None:
            return
        if bgr.shape[1] != self.W or bgr.shape[0] != self.H:
            bgr = cv2.resize(bgr, (self.W, self.H), interpolation=cv2.INTER_AREA)
        self._writer.write(bgr)
        self._frames_in_current += 1

    def _close_and_maybe_drop(self) -> None:
        """Release the current writer. If the rollout only has the start
        frame (no ``y`` press → no substeps), drop the MP4 — the rollout
        wasn't a real deployment, just a re-home."""
        if self._writer is None:
            return
        self._writer.release()
        path = self.current_path
        if self._frames_in_current <= 1:
            try:
                path.unlink()
            except FileNotFoundError:
                pass
            print(f"  vid: dropped {path.name} (only {self._frames_in_current} frame — "
                  f"no deploy)")
        else:
            print(f"  vid: closed {path.name} ({self._frames_in_current} frames)")

    def rotate(self) -> None:
        """Close current ``vid_N`` and open ``vid_{N+1}``.

        Caller is expected to ``append`` the new rollout's start frame
        right after; the recorder doesn't grab the camera itself.
        """
        self._close_and_maybe_drop()
        self._idx += 1
        self._frames_in_current = 0
        self._writer = cv2.VideoWriter(
            str(self.current_path), self.fourcc, self.fps, (self.W, self.H))
        if not self._writer.isOpened():
            raise RuntimeError(f"could not open VideoWriter for {self.current_path}")
        print(f"  vid: opened {self.current_path.name}")

    def stop(self) -> None:
        """Close the current writer; the final MP4 stays where it is
        unless it had ≤1 frame (no deploy). Then transcode all session
        MP4s from mp4v → H.264 so they're directly playable in browsers
        (mp4v is DivX-era and Chrome/Firefox/Safari refuse to render it)."""
        if self._writer is not None:
            self._close_and_maybe_drop()
            self._writer = None
            print(f"  → {self.out_dir}/")
        transcode_session_to_h264(self.out_dir)


def transcode_session_to_h264(session_dir: Path) -> int:
    """Transcode every ``vid_*.mp4`` in ``session_dir`` to H.264 in place.

    OpenCV's ``mp4v`` (MPEG-4 Part 2) writes reliably across OpenCV builds
    but isn't browser-playable. We shell out to ffmpeg (libx264 +
    movflags +faststart) and atomically replace each original. Originals
    are kept under ``session_dir/_mp4v_orig/`` in case the transcode
    botches anything.

    Returns the number of successful transcodes. Silent if ffmpeg isn't
    on PATH (deploy still works; videos just stay mp4v).
    """
    import shutil
    import subprocess
    if shutil.which("ffmpeg") is None:
        print(f"  vid: ffmpeg not on PATH — leaving {session_dir.name}/ as mp4v "
              f"(not browser-playable). `apt-get install ffmpeg` to fix.")
        return 0
    backup_dir = session_dir / "_mp4v_orig"
    backup_dir.mkdir(exist_ok=True)
    n = 0
    for src in sorted(session_dir.glob("vid_*.mp4")):
        if src.parent.name == "_mp4v_orig":
            continue
        out_tmp = src.with_suffix(".h264.mp4")
        try:
            r = subprocess.run(
                ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
                 "-i", str(src),
                 "-c:v", "libx264", "-preset", "veryfast",
                 "-movflags", "+faststart", "-pix_fmt", "yuv420p",
                 str(out_tmp)],
                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, timeout=120)
            if r.returncode != 0:
                print(f"  vid: transcode failed {src.name}: "
                      f"{r.stderr.decode(errors='replace')[:200]}")
                out_tmp.unlink(missing_ok=True)
                continue
            shutil.move(str(src), str(backup_dir / src.name))
            shutil.move(str(out_tmp), str(src))
            n += 1
        except subprocess.TimeoutExpired:
            print(f"  vid: ffmpeg timeout on {src.name}")
            out_tmp.unlink(missing_ok=True)
        except Exception as e:
            print(f"  vid: transcode error on {src.name}: {e}")
            out_tmp.unlink(missing_ok=True)
    if n:
        print(f"  ✓ transcoded {n} vid_*.mp4 to H.264 "
              f"(originals in {backup_dir.name}/)")
    return n