"""Record synchronized camera frames + calibrated servo states into a dataset.

For each tick it grabs a frame from the camera and reads the calibrated servo
positions (radians about the joint zero, using calib/<config>.json), and saves the
pair under datasets/<task>/:

  datasets/<task>/
    frames/000000.jpg     # camera frame
    states/000000.npy     # float32 array of calibrated joint radians (see meta.json)
    timestamps.npy        # (N,) wall-clock time per frame
    meta.json             # config, servo ids, joint names, signs, columns, counts

Servo torque is disabled so you can back-drive / teleoperate while recording.

Usage:
  python record.py <task_name> --config gripper [--camera 0] [--fps 30]
                   [--png] [--no-preview] [--overwrite] [--port /dev/tty...]
"""
import argparse
import json
import os
import sys
import time

import numpy as np
import cv2

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from feetech_bus import Bus, ticks_to_rad          # noqa: E402
from configs import calib_path                       # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
DATASETS = os.path.join(HERE, "datasets")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("task", help="dataset / task name (folder under datasets/)")
    ap.add_argument("--config", default="gripper")
    ap.add_argument("--camera", type=int, default=0, help="OpenCV camera index (default 0)")
    ap.add_argument("--fps", type=float, default=30.0)
    ap.add_argument("--png", action="store_true", help="save lossless PNG instead of JPG")
    ap.add_argument("--no-preview", action="store_true", help="don't open a preview window")
    ap.add_argument("--overwrite", action="store_true", help="clear an existing dataset of this name")
    ap.add_argument("--port", default=None)
    a = ap.parse_args()

    # --- calibration ---
    cpath = calib_path(a.config)
    try:
        calib = json.load(open(cpath))
    except FileNotFoundError:
        sys.exit(f"No calibration at {cpath}. Run calibrate_servo.py first.")
    if not calib:
        sys.exit(f"{cpath} is empty. Run calibrate_servo.py first.")
    servo_ids = sorted(int(s) for s in calib)
    joints = [calib[str(s)]["joint"] for s in servo_ids]
    signs = [int(calib[str(s)]["sign"]) for s in servo_ids]

    # --- dataset dirs ---
    ddir = os.path.join(DATASETS, a.task)
    fdir, sdir = os.path.join(ddir, "frames"), os.path.join(ddir, "states")
    if os.path.isdir(fdir) and os.listdir(fdir):
        if not a.overwrite:
            sys.exit(f"Dataset '{a.task}' already has frames. Use --overwrite or a new name.")
        for d in (fdir, sdir):
            for f in os.listdir(d):
                os.remove(os.path.join(d, f))
    os.makedirs(fdir, exist_ok=True)
    os.makedirs(sdir, exist_ok=True)
    ext = "png" if a.png else "jpg"

    # --- camera ---
    cap = cv2.VideoCapture(a.camera)
    if not cap.isOpened():
        sys.exit(f"Could not open camera {a.camera}.")

    # --- bus ---
    bus = Bus(a.port) if a.port else Bus()
    for sid in servo_ids:
        if not bus.ping(sid):
            cap.release()
            bus.close()
            sys.exit(f"Servo {sid} not responding.")
        bus.torque(sid, False)  # back-drivable while recording

    meta = dict(task=a.task, config=a.config, camera=a.camera, fps_target=a.fps, image_ext=ext,
                servo_ids=servo_ids, joints=joints, signs=signs,
                state_columns=[f"{j}_rad" for j in joints],
                started=time.strftime("%Y-%m-%d %H:%M:%S"))
    json.dump(meta, open(os.path.join(ddir, "meta.json"), "w"), indent=2)

    print(f"Recording '{a.task}' -> {ddir}")
    print(f"  camera {a.camera}, servos {servo_ids} -> {joints}")
    print("  press 'q' in the preview window, or Ctrl-C, to stop.\n")

    period = 1.0 / a.fps
    i = 0
    ts = []
    jpg_params = [] if a.png else [cv2.IMWRITE_JPEG_QUALITY, 95]
    try:
        while True:
            t0 = time.perf_counter()
            ok, frame = cap.read()
            if not ok:
                print("  camera read failed — stopping.")
                break
            state = np.full(len(servo_ids), np.nan, dtype=np.float32)
            for k, (sid, sgn) in enumerate(zip(servo_ids, signs)):
                pos = bus.read_pos(sid)
                if pos is not None:
                    state[k] = ticks_to_rad(pos, sgn)

            name = f"{i:06d}"
            cv2.imwrite(os.path.join(fdir, f"{name}.{ext}"), frame, jpg_params)
            np.save(os.path.join(sdir, f"{name}.npy"), state)
            ts.append(time.time())
            i += 1

            if not a.no_preview:
                disp = frame.copy()
                cv2.putText(disp, f"{a.task}  #{i}  {np.array2string(state, precision=3)}",
                            (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                cv2.imshow("recording (press q to stop)", disp)
                if cv2.waitKey(1) & 0xFF == ord("q"):
                    break

            dt = time.perf_counter() - t0
            if dt < period:
                time.sleep(period - dt)
    except KeyboardInterrupt:
        print("\n  stopped (Ctrl-C).")
    finally:
        np.save(os.path.join(ddir, "timestamps.npy"), np.asarray(ts))
        meta.update(num_frames=i, ended=time.strftime("%Y-%m-%d %H:%M:%S"),
                    duration_s=round(ts[-1] - ts[0], 3) if len(ts) > 1 else 0.0)
        json.dump(meta, open(os.path.join(ddir, "meta.json"), "w"), indent=2)
        cap.release()
        cv2.destroyAllWindows()
        bus.close()
        rate = i / (ts[-1] - ts[0]) if len(ts) > 1 else 0.0
        print(f"\nSaved {i} frames ({rate:.1f} fps) to {ddir}")


if __name__ == "__main__":
    main()