import cv2
import os
import logging
import argparse

# Setup logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

def extract_frames(video_path, output_dir, frame_interval=30):
    try:
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Output directory: {output_dir}")

        video = cv2.VideoCapture(video_path)
        if not video.isOpened():
            raise ValueError(f"Could not open video: {video_path}")

        fps = video.get(cv2.CAP_PROP_FPS)
        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
        logger.info(f"Video: {video_path}, FPS: {fps}, Total frames: {total_frames}")

        frame_count = 0
        saved_count = 0
        video_name = os.path.splitext(os.path.basename(video_path))[0]

        while True:
            ret, frame = video.read()
            if not ret:
                break

            if frame_count % frame_interval == 0:
                frame_path = os.path.join(output_dir, f"{video_name}_frame_{frame_count}.jpg")
                cv2.imwrite(frame_path, frame)
                saved_count += 1
                logger.info(f"Saved frame: {frame_path}")

            frame_count += 1

        video.release()
        logger.info(f"Extracted {saved_count} frames from {video_path}")
    except Exception as e:
        logger.error(f"Error extracting frames from {video_path}: {e}")
        raise

def main():
    parser = argparse.ArgumentParser(description="Extract frames from videos for training dataset.")
    parser.add_argument("--video-dir", default="videos", help="Directory containing input videos")
    parser.add_argument("--output-dir", default="dataset/images/train", help="Directory to save extracted frames")
    parser.add_argument("--frame-interval", type=int, default=30, help="Extract every nth frame")
    args = parser.parse_args()

    video_dir = args.video_dir
    output_dir = args.output_dir
    frame_interval = args.frame_interval

    if not os.path.exists(video_dir):
        logger.error(f"Video directory does not exist: {video_dir}")
        return

    for video_file in os.listdir(video_dir):
        if video_file.endswith((".mp4", ".avi", ".mov")):
            video_path = os.path.join(video_dir, video_file)
            extract_frames(video_path, output_dir, frame_interval)

if __name__ == "__main__":
    logger.info("Starting frame extraction...")
    main()