#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# gzf2_video.py - turn a GZCAPV01/GZF2 guest capture into a constant-rate mp4.
#
# Reads the append-only frame stream written by sdl_cap.c (LD_PRELOAD
# SDL_GL_SwapWindow shim) or xgrab.c, resamples by CLOCK_MONOTONIC timestamps
# (nearest frame), streams raw pixels straight into ffmpeg (no PPM
# intermediates), and validates the encoded file with ffprobe.
#
# Format: [0..8) = "GZCAPV01", zero pad to 4096; then records:
#   4B "GZF2" | u32 w | u32 h | u32 stride | u32 fmt | u64 ts_ns | u64 seq |
#   stride*h bytes, top-down.  fmt 32 = RGBA8 (GL swap path), 33 = BGRA8
#   (SDL_RenderReadPixels ARGB8888 / X surface).  A trailing
#   "GZF2" + u64 0xffffffffffffffff marks a clean stop.
#
# Usage:
#   gzf2_video.py CAPTURE OUT.mp4 [--fps 60] [--skip 5] [--duration 20]
#                 [--encoder libx264|nvenc] [--crf 18]
#                 [--pix-fmt rgba|bgra] [--info-only]
import argparse
import bisect
import hashlib
import json
import os
import struct
import subprocess
import sys

MAGIC = b"GZCAPV01"
SENTINEL = b"\xff" * 8


def die(msg):
    print(f"GZF2-ERROR {msg}", file=sys.stderr)
    sys.exit(2)


def index_capture(path):
    size = os.path.getsize(path)
    recs, clean, corrupt = [], False, 0
    with open(path, "rb") as f:
        hdr = f.read(4096)
        if len(hdr) < 4096 or hdr[:8] != MAGIC:
            die(f"{path}: not a GZCAPV01 file")
        off = 4096
        while off + 12 <= size:
            f.seek(off)
            head = f.read(36)
            if len(head) < 12 or head[:4] != b"GZF2":
                break
            if head[4:12] == SENTINEL:
                clean = True
                break
            if len(head) < 36:
                break
            w, h, stride, fmt = struct.unpack_from("<IIII", head, 4)
            ts, seq = struct.unpack_from("<QQ", head, 20)
            if not (0 < w <= 8192 and 0 < h <= 8192
                    and 0 < stride <= 4 * 8192) or seq == 0xFFFFFFFFFFFFFFFF:
                corrupt += 1
                break
            n = stride * h
            if off + 36 + n > size:
                corrupt += 1
                break
            f.seek(off + 36)
            px = f.read(n)
            if len(px) != n:
                corrupt += 1
                break
            recs.append({
                "off": off + 36, "ts": ts, "w": w, "h": h,
                "stride": stride, "fmt": fmt, "size": n,
                "md5": hashlib.md5(px).hexdigest(),
            })
            off += 36 + n
    if not recs:
        die(f"{path}: header only (no frames captured)")
    return recs, clean, corrupt


def sampled_black(px):
    s = px[:65536]
    return sum(s[0::4]) + sum(s[1::4]) + sum(s[2::4]) == 0


def pick_frames(recs, fps, skip_s, duration_s):
    t0 = recs[0]["ts"] + int(skip_s * 1e9)
    cand = [r for r in recs if r["ts"] >= t0] or recs
    ts = [r["ts"] for r in cand]
    span = (ts[-1] - ts[0]) / 1e9
    n = int(duration_s * fps) if duration_s else max(1, int(span * fps))
    picks, held = [], 0
    for i in range(n):
        want = ts[0] + int(i * 1e9 / fps)
        j = bisect.bisect_left(ts, want)
        if j == 0:
            k = 0
        elif j >= len(cand):
            k, held = len(cand) - 1, held + 1
        else:
            k = j if ts[j] - want < want - ts[j - 1] else j - 1
        picks.append(cand[k])
    return picks, span, held


def encoder_args(name, crf):
    if name == "nvenc":
        return ["-c:v", "h264_nvenc", "-preset", "p6", "-rc", "vbr",
                "-cq", str(crf), "-b:v", "12M"]
    return ["-c:v", "libx264", "-preset", "slow", "-crf", str(crf)]


def encode(cap, picks, args, encoder):
    first = picks[0]
    w, h, fmt = first["w"], first["h"], args.pix_fmt
    if any(p["w"] != w or p["h"] != h for p in picks):
        die("geometry changes inside the capture; resize before recording")
    if args.pix_fmt == "auto":
        fmt = "bgra" if first["fmt"] == 33 else "rgba"
    cmd = ["ffmpeg", "-hide_banner", "-v", "error", "-y",
           "-f", "rawvideo", "-pixel_format", fmt,
           "-video_size", f"{w}x{h}", "-framerate", str(args.fps),
           "-i", "pipe:0", "-an", *encoder_args(encoder, args.crf),
           "-frames:v", str(len(picks)), "-pix_fmt", "yuv420p",
           "-movflags", "+faststart", args.out]
    print(f"GZF2-ENCODE {encoder} {w}x{h} {fmt} {len(picks)}f @ {args.fps}fps")
    proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
    last, last_px = None, None
    try:
        with open(cap, "rb") as f:
            for p in picks:
                if p["off"] != last:
                    f.seek(p["off"])
                    last_px = f.read(p["size"])
                    if len(last_px) != p["size"]:
                        raise IOError(f"short read at {p['off']}")
                    last = p["off"]
                if p["stride"] == w * 4:
                    proc.stdin.write(last_px)
                else:
                    row = w * 4
                    proc.stdin.write(b"".join(
                        last_px[y * p["stride"]:y * p["stride"] + row]
                        for y in range(h)))
    except (BrokenPipeError, IOError) as e:
        print(f"GZF2-ENCODE write failed: {e}", file=sys.stderr)
    finally:
        if proc.stdin:
            try:
                proc.stdin.close()
            except BrokenPipeError:
                pass
    return proc.wait()


def ffprobe(path):
    r = subprocess.run(
        ["ffprobe", "-v", "error", "-count_frames", "-select_streams", "v:0",
         "-show_entries",
         "stream=codec_name,width,height,r_frame_rate,nb_read_frames",
         "-show_entries", "format=duration,size",
         "-of", "json", path], capture_output=True, text=True)
    if r.returncode:
        die(f"ffprobe failed: {r.stderr.strip()}")
    return json.loads(r.stdout)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("capture")
    ap.add_argument("out")
    ap.add_argument("--fps", type=int, default=60)
    ap.add_argument("--skip", type=float, default=0.0,
                    help="seconds of menu/load to drop from the start")
    ap.add_argument("--duration", type=float, default=0.0,
                    help="output seconds (0 = whole remaining span)")
    ap.add_argument("--encoder", choices=["libx264", "nvenc"],
                    default="libx264")
    ap.add_argument("--crf", type=int, default=18)
    ap.add_argument("--pix-fmt", choices=["auto", "rgba", "bgra"],
                    default="auto", help="auto uses record fmt (33=BGRA)")
    ap.add_argument("--info-only", action="store_true")
    args = ap.parse_args()

    recs, clean, corrupt = index_capture(args.capture)
    uniq = len({r["md5"] for r in recs})
    span = (recs[-1]["ts"] - recs[0]["ts"]) / 1e9
    rate = len(recs) / span if span else 0.0
    blacks = 0
    with open(args.capture, "rb") as f:
        for i, r in enumerate(recs):
            if i % 60 == 0:
                f.seek(r["off"])
                blacks += sampled_black(f.read(min(r["size"], 65536)))
    print(f"GZF2-CAPTURE frames={len(recs)} span={span:.2f}s rate={rate:.1f}/s "
          f"unique={uniq} ({100.0 * uniq / len(recs):.1f}%) "
          f"black_sampled={blacks} clean_end={clean} corrupt={corrupt} "
          f"size={os.path.getsize(args.capture)}B "
          f"geom={recs[0]['w']}x{recs[0]['h']} fmt={recs[0]['fmt']}")
    if args.info_only:
        return
    picks, span, held = pick_frames(recs, args.fps, args.skip,
                                    args.duration or None)
    md5_by_off = {r["off"]: r["md5"] for r in recs}
    pu = len({md5_by_off[p["off"]] for p in picks})
    print(f"GZF2-RESAMPLE out={len(picks)}f span={span:.2f}s "
          f"unique={pu}/{len(picks)} ({100.0 * pu / len(picks):.1f}%) "
          f"held_last={held}")
    if pu < len(picks) // 10:
        print("GZF2-WARN mostly static/black output", file=sys.stderr)

    rc = encode(args.capture, picks, args, args.encoder)
    if rc != 0 and args.encoder == "nvenc":
        print("GZF2-WARN nvenc failed; retrying libx264", file=sys.stderr)
        rc = encode(args.capture, picks, args, "libx264")
    if rc != 0:
        die(f"ffmpeg exited {rc}")
    info = ffprobe(args.out)
    st = info["streams"][0]
    print("GZF2-OUTPUT " + " ".join(
        f"{k}={st.get(k)}" for k in
        ("codec_name", "width", "height", "r_frame_rate", "nb_read_frames")))
    print(f"GZF2-OUTPUT duration={info['format'].get('duration')} "
          f"size={info['format'].get('size')}")
    if int(st.get("nb_read_frames", 0)) != len(picks):
        die("encoded frame count mismatch")
    print("GZF2-OK", args.out)


if __name__ == "__main__":
    main()