import os
import sys
import shutil
import platform
import subprocess
from tqdm import tqdm
import datetime
import cv2
from os import path
from modules.GFPGAN.gradio_inference_gfpgan import main as gfpgan_main

basePath = path.join(os.path.dirname(os.path.abspath(__file__)), 'temp')


def to_frame(video_path):
    work_dir = path.join(basePath, f"{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}")
    unProcessedFramesFolderPath = path.join(work_dir, 'frames')
    if not os.path.exists(unProcessedFramesFolderPath):
        os.makedirs(unProcessedFramesFolderPath)
    vidcap = cv2.VideoCapture(video_path)
    numberOfFrames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = vidcap.get(cv2.CAP_PROP_FPS)
    print("FPS: ", fps, "Frames: ", numberOfFrames)
    for frameNumber in tqdm(range(numberOfFrames)):
        _, image = vidcap.read()
        cv2.imencode('.jpg', image)[1].tofile(
            path.join(unProcessedFramesFolderPath, str(frameNumber).zfill(4) + '.jpg'))  # 中文路径
        # cv2.imwrite(path.join(unProcessedFramesFolderPath, str(frameNumber).zfill(4) + '.jpg'), image)

    wav_file_path = path.join(work_dir, 'audio.wav')
    print('Extracting raw audio...')
    command = 'ffmpeg -loglevel warning -y -i {} -strict -2 {}'.format(video_path, wav_file_path)
    subprocess.call(command, shell=True)
    return unProcessedFramesFolderPath, wav_file_path


# 删除中间临时文件
def clear(dirs):
    # dirs = [path.join(wav2lipPath, 'temp'),
    #         # path.join(gfpganPath, 'temp'),
    #         path.join(basePath, 'outputs/frames'),
    #         ]
    for dir in dirs:
        if os.path.exists(dir):
            try:
                shutil.rmtree(dir)
            except Exception as e:
                print(f'Error: {e}')


def concat_frame(work_path, inputAudioPath):
    restoredFramesPath = work_path + '/restored_imgs/'
    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)
        for filename in tqdm(dir_list[start:end]):
            filename = 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(work_path + '/batch_' + str(batch).zfill(4) + '.avi',
                              cv2.VideoWriter_fourcc(*'DIVX'), 30, size)
        batch = batch + 1

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

    # concat_frame
    concatTextFilePath = work_path + "/concat.txt"
    concatTextFile = open(concatTextFilePath, "w")
    for ips in range(batch):
        concatTextFile.write("file batch_" + str(ips).zfill(4) + ".avi\n")
    concatTextFile.close()

    concatedVideoOutputPath = work_path + "/concated_output.avi"

    subprocess.call(f'ffmpeg -loglevel warning -y -f concat -i {concatTextFilePath} -c copy {concatedVideoOutputPath}',
                    shell=True)
    output_dir = path.join(path.dirname(basePath), f"outputs/restore/")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    finalProcessedOuputVideo = path.join(output_dir, f"{datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.mp4")
    command = 'ffmpeg -loglevel warning -y -i {} -i {} -strict -2 -q:v 1 {}'.format(inputAudioPath,
                                                                                    concatedVideoOutputPath,
                                                                                    finalProcessedOuputVideo)
    subprocess.call(command, shell=platform.system() != 'Windows')
    return finalProcessedOuputVideo


def restore(video_path, only_center_face, upscale):
    unProcessedFramesFolderPath, wav_file_path = to_frame(video_path)
    work_dir = os.path.dirname(unProcessedFramesFolderPath)
    gfpgan_main(unProcessedFramesFolderPath, work_dir, only_center_face=only_center_face, upscale=upscale,
                bg_upsampler=None)  # 修复图像
    finalProcessedOuputVideo = concat_frame(work_dir, wav_file_path)
    return finalProcessedOuputVideo