import sys
import os
import argparse
from pathlib import Path

sys.path.append(os.path.abspath(sys.path[0] + '/../../'))

parser = argparse.ArgumentParser()
parser.add_argument('-s', '--start', type=int)
parser.add_argument('-t', '--duration', type=int)
parser.add_argument('--fps', type=str)
parser.add_argument('--subset', type=str)
parser.add_argument('datadir', type=str)
args = parser.parse_args()

os.chdir(args.datadir)

if args.subset is not None:
    video_dir = Path(f"videos/{args.subset}")
else:
    video_dir = Path(f"video")
for video_path in video_dir.glob("*.*"):
    # Extract frames from video
    image_dir = "images" if args.subset is None else f"images/{args.subset}"
    os.makedirs(f"{image_dir}/{video_path.stem}", exist_ok=True)
    extra_args = []
    if args.start is not None:
        extra_args.append(f"-ss {args.start}")
    if args.duration is not None:
        extra_args.append(f"-t {args.duration}")
    if args.fps is not None:
        extra_args.append(f"-vf fps={args.fps}")
    extra_args = ' '.join(extra_args)
    os.system(f"ffmpeg -i {video_path} {extra_args} -f image2 -q:v 2 "
              f"{image_dir}/{video_path.stem}/image%03d.png")