#!/usr/bin/env python3
import argparse
import sys
from pathlib import Path

import cv2 as cv


parser = argparse.ArgumentParser(description="Extract selected frames from a video with OpenCV.")
parser.add_argument("video", help="video file to read")
parser.add_argument(
    "--output-dir",
    default="output/frames",
    help="directory for extracted images",
)
parser.add_argument(
    "--frames",
    nargs="+",
    type=int,
    required=True,
    help="0-based frame numbers to save",
)
parser.add_argument("--prefix", default="frame", help="output filename prefix")
args = parser.parse_args()

video_path = Path(args.video)
output_dir = Path(args.output_dir)

capture = cv.VideoCapture(str(video_path))
if not capture.isOpened():
    sys.exit(f"Could not open video: {video_path}")

try:
    total_frames = int(capture.get(cv.CAP_PROP_FRAME_COUNT))
    fps = capture.get(cv.CAP_PROP_FPS)
    width = int(capture.get(cv.CAP_PROP_FRAME_WIDTH))
    height = int(capture.get(cv.CAP_PROP_FRAME_HEIGHT))
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"video: {video_path}")
    print(f"frames: {total_frames}")
    print(f"fps: {fps:.2f}")
    print(f"size: {width}x{height}")

    for frame_number in args.frames:
        if frame_number < 0:
            sys.exit(f"Frame index must be zero or greater: {frame_number}")
        if total_frames and frame_number >= total_frames:
            sys.exit(
                f"Frame {frame_number} is outside the video length of {total_frames} frames"
            )

        capture.set(cv.CAP_PROP_POS_FRAMES, frame_number)
        ok, frame = capture.read()
        if not ok or frame is None:
            sys.exit(f"Could not read frame {frame_number}")

        output_path = output_dir / f"{args.prefix}-{frame_number:04d}.jpg"
        if not cv.imwrite(str(output_path), frame):
            sys.exit(f"Could not write frame image: {output_path}")

        timestamp = frame_number / fps if fps > 0 else None
        time_label = f"{timestamp:.2f}s" if timestamp is not None else "unknown"
        print(f"saved: frame={frame_number} time={time_label} path={output_path}")
finally:
    capture.release()