import argparse
import os
import sys

from tqdm import tqdm
from ffmpy import FFmpeg
import cv2


def cli(args):
    folder_path = args.job_folder

    restoredFramesPath = os.path.join(folder_path, 'restored_imgs')
    os.makedirs(restoredFramesPath, exist_ok=True)
    dir_list = os.listdir(restoredFramesPath)
    dir_list.sort()
    batch = 0
    batchSize = 300
    for i in tqdm(range(0, len(dir_list), batchSize)):
        img_array = []
        start, end = i, i + batchSize
        print("processing ", start, end)

        batch_video = os.path.join(folder_path, 'batch_' + str(batch).zfill(4) + '.avi')
        if os.path.isfile(batch_video):
            print(f'{batch_video} existed, skipped.')
            batch = batch + 1
            continue

        for filename in tqdm(dir_list[start:end]):
            filename = os.path.join(restoredFramesPath, filename)
            img = cv2.imread(filename)
            if img is None:
                continue
            height, width, layers = img.shape
            size = (width, height)
            img_array.append(img)

        out = cv2.VideoWriter(batch_video, cv2.VideoWriter_fourcc(*'DIVX'), 30, size)
        batch = batch + 1
        print(f'batch video dumped: {batch_video}')

        for j in range(len(img_array)):
            out.write(img_array[j])
        out.release()

    concatTextFilePath = os.path.join(folder_path, "concat.txt")

    concatTextFile = open(concatTextFilePath, "w")
    for ips in range(batch):
        filename = os.path.join(folder_path, "batch_" + str(ips).zfill(4) + ".avi")
        concatTextFile.write(f"file '{filename}'\n")
    concatTextFile.close()

    ff = FFmpeg(
        inputs={
            concatTextFilePath: '-f concat -safe 0 ',
            args.audio: None
        },
        outputs={args.output: '-y -c:v libx264 -c:a aac -shortest'}
    )
    # !ffmpeg -y -f concat -safe 0 -i {concatTextFilePath} -i {audio} -c:v libx264 -c:a aac -shortest {output}
    ff.run()


if __name__ == '__main__':
    # merge images in restored_imgs folder as video, merge audio
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-j',
        '--job_folder',
        help='job_folder, ./restored_imgs must under it')
    parser.add_argument('-a', '--audio', type=str, help='audio file path')
    parser.add_argument('-o', '--output', type=str, help='output file path')
    args = parser.parse_args()

    cli(args)