import os
import cv2
import subprocess
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor

def is_video_file(filename):
    video_exts = ('.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv')
    return filename.lower().endswith(video_exts)

def get_video_resolution(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    cap.release()
    return (width, height)

def resize_video_ffmpeg(input_path, output_path, size=(256, 256)):
    # Overwrite the input file with ffmpeg
    tmp_output = input_path + ".tmp.mp4"
    cmd = [
        "ffmpeg", "-y", "-i", input_path,
        "-vf", f"scale={size[0]}:{size[1]}",
        "-c:v", "libx264", "-preset", "fast", "-crf", "23",
        "-c:a", "copy",
        tmp_output
    ]
    subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    os.replace(tmp_output, output_path)

def collect_video_files(root_folder):
    video_files = []
    for dirpath, _, filenames in os.walk(root_folder):
        for filename in filenames:
            if is_video_file(filename):
                video_files.append(os.path.join(dirpath, filename))
    return video_files

def process_single_video(video_path):
    res = get_video_resolution(video_path)
    if res is None:
        return f"Failed to open {video_path}"
    if res == (256, 256):
        return None
    resize_video_ffmpeg(video_path, video_path)
    return f"Resized {video_path} from {res} to (256, 256)"

def process_videos(root_folder):
    video_files = collect_video_files(root_folder)
    results = []
    with ProcessPoolExecutor() as executor:
        for result in tqdm(executor.map(process_single_video, video_files), total=len(video_files), desc="Processing videos"):
            if result:
                print(result)

if __name__ == "__main__":
    common_videos_dir = os.path.join(os.path.dirname(__file__), "common_videos")
    process_videos(common_videos_dir)