import os
import sys
import subprocess
import argparse
from pathlib import Path

def extract_first_frame(video_path, output_path):
    """
    Extract the first frame from a video using ffmpeg
    """
    try:
        # Use ffmpeg to extract the first frame
        cmd = [
            'ffmpeg',
            '-i', str(video_path),
            '-vframes', '1',
            '-f', 'image2',
            '-y',  # Overwrite output file if it exists
            str(output_path)
        ]
        
        # Run ffmpeg with suppressed output
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print(f"✓ Extracted frame: {output_path}")
            return True
        else:
            print(f"✗ Failed to extract frame from {video_path}")
            print(f"Error: {result.stderr}")
            return False
            
    except FileNotFoundError:
        print("Error: ffmpeg not found. Please install ffmpeg and make sure it's in your PATH")
        return False
    except Exception as e:
        print(f"Error processing {video_path}: {e}")
        return False

def is_video_file(filename):
    """
    Check if file is a video based on common extensions
    """
    video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv', '.webm', '.m4v', '.3gp'}
    return Path(filename).suffix.lower() in video_extensions

def batch_extract_frames(input_folder, output_folder):
    """
    Batch process all videos in input folder and extract first frames
    """
    input_path = Path(input_folder)
    output_path = Path(output_folder)
    
    # Check if input folder exists
    if not input_path.exists():
        print(f"Error: Input folder '{input_folder}' does not exist")
        return
    
    # Create output folder if it doesn't exist
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Find all video files
    video_files = [f for f in input_path.iterdir() if f.is_file() and is_video_file(f.name)]
    
    if not video_files:
        print(f"No video files found in '{input_folder}'")
        return
    
    print(f"Found {len(video_files)} video file(s)")
    print(f"Output folder: {output_folder}")
    print("-" * 50)
    
    successful = 0
    failed = 0
    
    for video_file in video_files:
        # Create output filename (video name + .png)
        output_filename = video_file.stem + '.png'
        output_file_path = output_path / output_filename
        
        print(f"Processing: {video_file.name}")
        
        if extract_first_frame(video_file, output_file_path):
            successful += 1
        else:
            failed += 1
    
    print("-" * 50)
    print(f"Processing complete!")
    print(f"Successful: {successful}")
    print(f"Failed: {failed}")

def main():
    parser = argparse.ArgumentParser(description='Extract first frames from videos in a folder')
    parser.add_argument('input_folder', help='Input folder containing video files')
    parser.add_argument('output_folder', help='Output folder for extracted frames')
    
    args = parser.parse_args()
    
    batch_extract_frames(args.input_folder, args.output_folder)

if __name__ == "__main__":
    # You can also run it directly by modifying these paths
    if len(sys.argv) == 1:
        print("Usage examples:")
        print("python frame_extractor.py /path/to/videos /path/to/output")
        print("python frame_extractor.py C:\\Videos C:\\Frames")
        print("\nOr modify the script to set default paths:")
        
        # Uncomment and modify these lines to set default paths
        # input_folder = "path/to/your/video/folder"
        # output_folder = "path/to/your/output/folder"
        # batch_extract_frames(input_folder, output_folder)
    else:
        main()