import os
import subprocess
import re
import shutil

def convert_video(video_path, output_path):
    print(f"Converting video, {video_path}")

    # Codec that works with openframeworks!
    
    command = [
        "ffmpeg", 
        "-i", video_path,                  # Input file
        "-c:v", "libx264",                 # Video codec
        "-profile:v", "high",              # Video profile
        "-c:a", "copy",                    # Copy audio without re-encoding
        output_path                       # Output file
    ]

    results = subprocess.run(command, stderr=subprocess.PIPE, text=True)
    output = results.stderr

    print(f"Video converted {output}")

def get_frame_rate(video_path):
    """
    Get the frame rate of a video file using ffmpeg.
    """
    command = [
        'ffmpeg',
        '-i', video_path,
        '-vcodec', 'png',
        '-f', 'null',
        '-'
    ]
    
    result = subprocess.run(command, stderr=subprocess.PIPE, text=True)
    output = result.stderr

    match = re.search(r'(\d+(\.\d+)?) fps', output)
    if match:
        return float(match.group(1))
    else:
        raise ValueError("Frame rate not found in video information.")

def extract_frames_from_video(video_path, output_folder):
    """
    Extract frames from a video file at 1 fps.
    """
    os.makedirs(output_folder, exist_ok=True)

    command = [
        'ffmpeg',
        '-i', video_path,                 # Input file
        '-vf', 'fps=1/10',                   # Extract 1 frame per second
        '-vsync', 'vfr',                  # Variable frame rate
        os.path.join(output_folder, 'frame_%09d.png')  # Output pattern
    ]

    subprocess.run(command, check=True)

def rename_frames(output_folder, frame_rate):
    """
    Rename the extracted frames to match their original frame numbers.
    """
    frame_files = sorted(f for f in os.listdir(output_folder) if f.startswith('frame_') and f.endswith('.png'))
    for i, filename in enumerate(frame_files):
        # Calculate the original frame number
        original_frame_number = int(i * frame_rate)
        
        # Create the new filename with the original frame number
        new_filename = f"{original_frame_number:09d}.png"
        os.rename(
            os.path.join(output_folder, filename),
            os.path.join(output_folder, new_filename)
        )

def process_videos_in_folder(folder_path):
    """
    Process all video files in the specified folder and extract frames from each video.
    """
    for filename in os.listdir(folder_path):
        if filename.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.flv')):

            video_path = os.path.join(folder_path, filename)

            output_folder = os.path.join(folder_path, os.path.splitext(filename)[0])  # Create a folder for each video

            # New location for .mp4
            video_copy_location = os.path.join(output_folder, filename)

            # Create frames folder
            frames_folder = os.path.join(output_folder, "frames")
            
            if not os.path.exists(output_folder):
                os.mkdir(output_folder)

            if not os.path.exists(frames_folder):
                os.mkdir(frames_folder)

            print(f'Processing video: {video_path}')
            
            try:
                #convert_video(video_path, video_copy_location)
                frame_rate = get_frame_rate(video_path)
                print(f'Original frame rate: {frame_rate} fps')
                extract_frames_from_video(video_path, frames_folder)
                rename_frames(frames_folder, frame_rate)
                print(f'Frames saved and renamed to: {frames_folder}')
            except ValueError as e:
                print(f'Error processing video {video_path}: {e}')

if __name__ == "__main__":
    # Path to the folder containing videos
    folder_path = 'main_dataset'
    
    process_videos_in_folder(folder_path)