import argparse
import av
import cv2
import time

DEVICE = "/dev/video4"
TARGET_FPS = 30
WIDTH = 1920
HEIGHT = 1080


def main():
    parser = argparse.ArgumentParser(description="Capture from V4L2 via PyAV, optionally save to MP4.")
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default="",
        help="Output MP4 path (e.g. capture.mp4). If empty, only preview, no file.",
    )
    args = parser.parse_args()

    # Open v4l2 device with same style options as ffmpeg CLI
    container = av.open(
        DEVICE,
        format="v4l2",
        options={
            "input_format": "mjpeg",
            "video_size": f"{WIDTH}x{HEIGHT}",
            "framerate": str(TARGET_FPS),
        },
    )

    stream = container.streams.video[0]
    frame_count = 0
    t_start = time.perf_counter()

    writer = None
    if args.output:
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        writer = cv2.VideoWriter(
            args.output, fourcc, float(TARGET_FPS), (WIDTH, HEIGHT)
        )
        if not writer.isOpened():
            container.close()
            raise RuntimeError(f"Could not open VideoWriter for {args.output}")
        print(f"Recording to {args.output} @ {TARGET_FPS} fps, {WIDTH}x{HEIGHT}")

    try:
        for frame in container.decode(stream):
            t0 = time.perf_counter()
            img = frame.to_ndarray(format="bgr24")  # OpenCV-compatible
            t1 = time.perf_counter()

            if writer is not None:
                writer.write(img)

            cv2.imshow("pyav_cam", img)
            key = cv2.waitKey(1) & 0xFF
            t2 = time.perf_counter()

            frame_count += 1
            if frame_count % 30 == 0:
                elapsed = t2 - t_start
                print(
                    f"Frame {frame_count}: convert={1000*(t1-t0):.1f}ms, "
                    f"imshow+waitKey={1000*(t2-t1):.1f}ms, avg_fps={frame_count/elapsed:.1f}"
                )

            if key == ord("q"):
                break
    finally:
        container.close()
        if writer is not None:
            writer.release()
            print(f"Saved {frame_count} frames to {args.output}")
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()