import cv2
import os
import glob
import shutil
import subprocess

def extract_frames(video_path, output_folder, overwrite=True):
    # Clear output folder if needed
    if overwrite and os.path.exists(output_folder):
        shutil.rmtree(output_folder)
    os.makedirs(output_folder, exist_ok=True)

    # Run ffmpeg to extract frames
    output_pattern = os.path.join(output_folder, "%d.jpg")
    command = [
        "ffmpeg", "-i", video_path,
        "-q:v", "2",  # Quality for JPEG (lower is better)
        output_pattern
    ]
    subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    print(f"✅ Extracted frames to: {output_folder}")

def extract_frames_old(video_path, output_folder):
    # Clear output folder if it exists
    if os.path.exists(output_folder):
        for f in glob.glob(os.path.join(output_folder, "*")):
            os.remove(f)
    else:
        os.makedirs(output_folder)

    cap = cv2.VideoCapture(video_path)
    frame_id = 0

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

        frame_path = os.path.join(output_folder, f"{frame_id}.jpg")
        cv2.imwrite(frame_path, frame)
        frame_id += 1

    cap.release()
    print(f"✅ Extracted {frame_id} frames to: {output_folder}")


def recombine_frames_(video_path, frames_folder, output_video=None):
    # Default output path if not provided
    if output_video is None:
        base, ext = os.path.splitext(video_path)
        output_video = f"{base}_anonymised.mp4"

    # Get frame size and FPS from original video
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()

    # Load all frame paths and sort by frame number
    frame_paths = sorted(glob.glob(os.path.join(frames_folder, "*.jpg")), key=lambda x: int(os.path.basename(x).split('.')[0]))

    if not frame_paths:
        print("❌ No frames found to recombine.")
        return

    # Get frame size from first frame
    frame_sample = cv2.imread(frame_paths[0])
    height, width = frame_sample.shape[:2]

    # Define the video writer
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video, fourcc, fps, (width, height))

    for frame_path in frame_paths:
        frame = cv2.imread(frame_path)
        out.write(frame)

    out.release()
    print(f"✅ Video saved to: {output_video}")

def recombine_frames(video_path, frames_folder, output_video=None, temp_video_path="temp_video_no_audio.mp4"):
    """
    Recombines frames into a video and adds back the original audio using ffmpeg.
    """
    if output_video is None:
        base, ext = os.path.splitext(video_path)
        output_video = f"{base}_anonymised.mp4"

    # Get frame size and FPS from original video
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()

    # Collect and sort frame paths
    frame_paths = sorted(
        glob.glob(os.path.join(frames_folder, "*.jpg")),
        key=lambda x: int(os.path.splitext(os.path.basename(x))[0])
    )

    if not frame_paths:
        print("❌ No frames found to recombine.")
        return

    # Get frame size
    sample_frame = cv2.imread(frame_paths[0])
    height, width = sample_frame.shape[:2]

    # Write video without audio
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))
    for frame_path in frame_paths:
        frame = cv2.imread(frame_path)
        out.write(frame)
    out.release()

    # Use ffmpeg to add original audio
    ffmpeg_command = [
        "ffmpeg",
        "-y",  # Overwrite without asking
        "-i", temp_video_path,
        "-i", video_path,
        "-c:v", "copy",
        "-map", "0:v:0",
        "-map", "1:a:0?",
        "-shortest",
        output_video
    ]
    subprocess.run(ffmpeg_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    # Cleanup
    #os.remove(temp_video_path)
    print(f"✅ Final video with audio saved to: {output_video}")

if __name__ == "__main__":
    input_video = "/Users/sophiemaw/Downloads/CONFIDENTIAL DO NOT SHARE Edna & Paul 29.10.10 Part 2 00.12.46.531.mov"
    output_folder = 'frames'
    extract_frames(input_video,output_folder)