"""Decode a ZED SVO2 file → mp4 (H.264) + per-frame timestamps json.

Reads the LEFT eye stream (RGB), writes an mp4 at the recorded framerate and
a sibling .timestamps.json holding one int64-ns timestamp per frame. The
timestamps are aligned with the robot_data.npz `timestamps` array so a
viewer can scrub proprio + video together.

Usage:
    python decode_svo2.py --svo IN.svo2 --out OUT.mp4
"""
import argparse
import json
import time
from pathlib import Path

import cv2
import numpy as np
import pyzed.sl as sl


def decode(svo_path: Path, out_mp4: Path, quality: int = 85) -> None:
    zed = sl.Camera()
    init = sl.InitParameters()
    init.set_from_svo_file(str(svo_path))
    init.svo_real_time_mode = False
    init.depth_mode = sl.DEPTH_MODE.NONE
    status = zed.open(init)
    if status != sl.ERROR_CODE.SUCCESS:
        raise RuntimeError(f"zed.open failed: {status}")

    n_frames = zed.get_svo_number_of_frames()
    res = zed.get_camera_information().camera_configuration.resolution
    W, H = res.width, res.height
    fps_raw = zed.get_camera_information().camera_configuration.fps
    fps = int(round(fps_raw)) if fps_raw > 0 else 30
    print(f"  {svo_path.name}: {n_frames} frames  {W}x{H}  {fps} Hz  → {out_mp4}")

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(str(out_mp4), fourcc, fps, (W, H))
    if not writer.isOpened():
        raise RuntimeError("cv2.VideoWriter open failed")

    img = sl.Mat()
    rt = sl.RuntimeParameters()
    timestamps_ns = []

    t0 = time.time()
    for i in range(n_frames):
        s = zed.grab(rt)
        if s != sl.ERROR_CODE.SUCCESS:
            timestamps_ns.append(-1)
            continue
        zed.retrieve_image(img, sl.VIEW.LEFT)
        # sl.VIEW.LEFT is 4-channel (BGRA); mp4v wants BGR.
        arr = img.get_data()
        if arr.shape[2] == 4:
            arr = cv2.cvtColor(arr, cv2.COLOR_BGRA2BGR)
        writer.write(arr)
        ts = zed.get_timestamp(sl.TIME_REFERENCE.IMAGE).get_nanoseconds()
        timestamps_ns.append(int(ts))
        if (i + 1) % 500 == 0:
            print(f"    {i + 1}/{n_frames}  ({(i + 1) / (time.time() - t0):.0f} fps)")

    writer.release()
    zed.close()

    ts_path = out_mp4.with_suffix(".timestamps.json")
    json.dump(
        {"timestamps_ns": timestamps_ns, "fps": fps, "width": W, "height": H,
         "n_frames": n_frames, "source": str(svo_path)},
        ts_path.open("w"),
    )
    print(f"  wrote {out_mp4.name} + {ts_path.name}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--svo", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True)
    args = ap.parse_args()
    args.out.parent.mkdir(parents=True, exist_ok=True)
    decode(args.svo, args.out)


if __name__ == "__main__":
    main()