#!/usr/bin/env python3
"""
extract_frames.py

Extract frames from a video file using OpenCV.

Usage examples:
python extract_frames.py -i sample.mp4 -o frames_dir --every-frame 10
python extract_frames.py -i sample.mp4 -o frames_dir --every-sec 1 --display
"""

import cv2
import os
import argparse
from pathlib import Path
from datetime import timedelta

def parse_args():
    p = argparse.ArgumentParser(description="Extract frames from a video file using OpenCV.")
    p.add_argument("-i", "--input", required=True, help="Path to input video file.")
    p.add_argument("-o", "--output", default="frames", help="Output directory to save extracted frames.")
    group = p.add_mutually_exclusive_group()
    group.add_argument("--every-frame", type=int, default=1,
                       help="Save every Nth frame (default=1 -> save every frame).")
    group.add_argument("--every-sec", type=float,
                       help="Save one frame every N seconds (overrides --every-frame).")
    p.add_argument("--display", action="store_true", help="Show frames in a window while extracting.")
    p.add_argument("--max-frames", type=int, default=0, help="Stop after saving this many frames (0 = no limit).")
    p.add_argument("--prefix", default="frame", help="Filename prefix for saved frames.")
    return p.parse_args()

def ensure_output_dir(path: Path):
    path.mkdir(parents=True, exist_ok=True)

def seconds_to_hhmmss(sec: float) -> str:
    td = timedelta(seconds=round(sec))
    return str(td)

def main():
    args = parse_args()
    input_path = Path(args.input)
    output_dir = Path(args.output)
    ensure_output_dir(output_dir)

    # Open the video capture
    cap = cv2.VideoCapture(str(input_path))
    if not cap.isOpened():
        print(f"Error: Could not open video file: {input_path}")
        return

    # Get video properties
    fps = cap.get(cv2.CAP_PROP_FPS) or 0
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
    duration_sec = total_frames / fps if fps > 0 else 0

    print(f"Input: {input_path}")
    print(f"Resolution: {int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))}x{int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))}")
    print(f"FPS: {fps:.2f}")
    print(f"Total frames (approx): {total_frames}")
    print(f"Duration (sec): {duration_sec:.2f}")
    print(f"Saving to directory: {output_dir.resolve()}")

    # Determine the saving interval
    if args.every_sec is not None:
        # Convert seconds to frames interval (may be fractional; round to nearest frame)
        if fps <= 0:
            print("Warning: FPS unknown, falling back to every-frame behavior.")
            save_every_n_frames = max(1, args.every_frame)
        else:
            save_every_n_frames = max(1, int(round(args.every_sec * fps)))
        print(f"Saving one frame every {args.every_sec} sec -> every {save_every_n_frames} frames (approx).")
    else:
        save_every_n_frames = max(1, args.every_frame)
        print(f"Saving every {save_every_n_frames} frames.")

    saved_count = 0
    read_count = 0
    frame_idx = 0

    # Loop over frames
    while True:
        ret, frame = cap.read()
        if not ret:
            # End of video or read error
            break
        read_count += 1
        frame_idx += 1

        # Decide whether to save this frame
        if (frame_idx - 1) % save_every_n_frames == 0:
            # Compute timestamp (in seconds) of current frame
            timestamp_sec = frame_idx / fps if fps > 0 else 0.0
            timestamp_text = seconds_to_hhmmss(timestamp_sec)

            # Construct filename with zero-padded frame index and timestamp
            filename = f"{args.prefix}_{frame_idx:06d}_{timestamp_text.replace(':','-')}.png"
            output_file = output_dir / filename

            # Save the image (PNG keeps good quality)
            cv2.imwrite(str(output_file), frame)
            saved_count += 1

            # Print progress
            print(f"[{saved_count}] Saved frame {frame_idx} @ {timestamp_sec:.2f}s -> {output_file.name}")

            # Stop if reached max frames limit
            if args.max_frames and saved_count >= args.max_frames:
                print("Reached max_frames limit. Stopping extraction.")
                break

        # Optionally display the frame (scaled to fit window)
        if args.display:
            disp = frame.copy()
            h, w = disp.shape[:2]
            max_w, max_h = 960, 540
            scale = min(1.0, max_w / w, max_h / h)
            if scale < 1.0:
                disp = cv2.resize(disp, (int(w*scale), int(h*scale)), interpolation=cv2.INTER_LINEAR)
            cv2.imshow("Frame Extractor - Press 'q' to quit display", disp)
            # waitKey small delay to allow window to refresh; 1ms is okay for fast extraction
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                print("User requested display quit. Continuing extraction without display.")
                args.display = False
                cv2.destroyAllWindows()

    # Release resources
    cap.release()
    cv2.destroyAllWindows()

    print("Extraction complete.")
    print(f"Total frames read: {read_count}")
    print(f"Total frames saved: {saved_count}")
    print(f"Saved images are in: {output_dir.resolve()}")

if __name__ == "__main__":
    main()