import cv2
import os
import shutil
import time

def extract_frames(video_path, frames_dir, quality=95):
    """Extract frames as high-quality JPGs (better quality + smaller than PNG)."""
    os.makedirs(frames_dir, exist_ok=True)

    existing = [f for f in os.listdir(frames_dir) if f.startswith("frame_") and f.endswith(".jpg")]
    last_idx = max([int(f.split("_")[1].split(".")[0]) for f in existing]) if existing else -1

    cap = cv2.VideoCapture(video_path)
    frame_paths = []
    idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frame_path = os.path.join(frames_dir, f"frame_{idx:05d}.jpg")
        if idx > last_idx:
            # High quality JPG
            cv2.imwrite(frame_path, frame, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
        frame_paths.append(frame_path)
        idx += 1
    cap.release()
    return frame_paths

def frames_to_video(frames_dir, output_video_path, fps, use_ffmpeg=True, crf=17):
    """
    Prefer ffmpeg for much better quality and optional hardware encoding.
    Falls back to OpenCV if ffmpeg fails.
    """
    frames = sorted([
        os.path.join(frames_dir, f)
        for f in os.listdir(frames_dir)
        if f.endswith('.jpg') and f.startswith("swapped_")
    ])
    if not frames:
        print("No swapped frames found.")
        return

    if use_ffmpeg:
        # Create a temporary file list for ffmpeg
        list_file = os.path.join(frames_dir, "frames.txt")
        with open(list_file, "w") as f:
            for fp in frames:
                f.write(f"file '{os.path.abspath(fp)}'\n")

        # High quality + try NVENC if available on A100
        cmd = [
            "ffmpeg", "-y",
            "-f", "concat", "-safe", "0",
            "-r", str(fps),
            "-i", list_file,
            "-c:v", "libx264",          # change to h264_nvenc if your Space supports it
            "-preset", "slow",
            "-crf", str(crf),
            "-pix_fmt", "yuv420p",
            output_video_path
        ]
        try:
            import subprocess
            subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            os.remove(list_file)
            return
        except Exception as e:
            print(f"ffmpeg failed ({e}), falling back to OpenCV")

    # Fallback
    first = cv2.imread(frames[0])
    h, w = first.shape[:2]
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (w, h))
    for fp in frames:
        out.write(cv2.imread(fp))
    out.release()