import os, glob
from utils.commons.multiprocess_utils import multiprocess_run_tqdm

from data_gen.utils.path_converter import PathConverter, pc

# mp4_names = glob.glob("/home/tiger/datasets/raw/CelebV-HQ/video/*.mp4")

def extract_img_job(video_name, raw_img_dir=None):
    if raw_img_dir is not None:
        out_path = raw_img_dir
    else:
        out_path = pc.to(video_name.replace(".mp4", ""), "vid", "gt")
    os.makedirs(out_path, exist_ok=True)
    ffmpeg_path = "/usr/bin/ffmpeg"
    cmd = f'{ffmpeg_path} -i {video_name} -vf fps={25},scale=w=512:h=512 -qmin 1 -q:v 1 -start_number 0 -v quiet {os.path.join(out_path, "%8d.jpg")}'
    os.system(cmd)

if __name__ == '__main__':
    import argparse, glob, tqdm, random
    parser = argparse.ArgumentParser()
    parser.add_argument("--vid_dir", default='/home/tiger/datasets/raw/CelebV-HQ/video')
    parser.add_argument("--ds_name", default='CelebV-HQ')
    parser.add_argument("--num_workers", default=64, type=int)
    parser.add_argument("--process_id", default=0, type=int)
    parser.add_argument("--total_process", default=1, type=int)
    args = parser.parse_args()
    vid_dir = args.vid_dir
    ds_name = args.ds_name
    if ds_name in ['lrs3_trainval']:
        mp4_name_pattern = os.path.join(vid_dir, "*/*.mp4")
    elif ds_name in ['TH1KH_512', 'CelebV-HQ']:
        vid_names = glob.glob(os.path.join(vid_dir, "*.mp4"))
    elif ds_name in ['lrs2', 'lrs3', 'voxceleb2']:
        vid_name_pattern = os.path.join(vid_dir, "*/*/*.mp4")
        vid_names = glob.glob(vid_name_pattern)
    elif ds_name in ["RAVDESS", 'VFHQ']:
        vid_name_pattern = os.path.join(vid_dir, "*/*/*/*.mp4")
        vid_names = glob.glob(vid_name_pattern)
    vid_names = sorted(vid_names)
    
    process_id = args.process_id
    total_process = args.total_process
    if total_process > 1:
        assert process_id <= total_process -1
        num_samples_per_process = len(vid_names) // total_process
        if process_id == total_process:
            vid_names = vid_names[process_id * num_samples_per_process : ]
        else:
            vid_names = vid_names[process_id * num_samples_per_process : (process_id+1) * num_samples_per_process]
    
    for i, res in multiprocess_run_tqdm(extract_img_job, vid_names, num_workers=args.num_workers, desc="extracting images"):
        pass