import os
import sys
import requests

from main import call_wav2lip, call_gfpgan, merge

root_dir = '/content/jobs'
os.makedirs(root_dir, exist_ok=True)


def download_file(url, local_filename):
    response = requests.get(url, stream=True)
    with open(local_filename, 'wb') as file:
        for chunk in response.iter_content(chunk_size=8192):
            if chunk:
                file.write(chunk)


def main(job_id, video_url, audio_url):
    job_path = os.path.join(root_dir, job_id)
    os.makedirs(job_path, exist_ok=True)

    if video_url.startswith('http'):
        video_file = os.path.basename(video_url)
        video_path = os.path.join(job_path, video_file)
        download_file(video_url, video_path)
    else:
        video_path = video_url

    if audio_url.startswith('http'):
        audio_file = os.path.basename(audio_url)
        audio_path = os.path.join(job_path, audio_file)
        download_file(audio_url, audio_path)
    else:
        audio_path = audio_url

    assert os.path.isfile(video_path), f'Video {video_path} not exist.'
    assert os.path.isfile(audio_path), f'Audio {audio_path} not exist.'

    wav2lip_mp4 = os.path.join(job_path, 'wav2lip.mp4')
    call_wav2lip(video_path, audio_path, wav2lip_mp4)
    call_gfpgan(wav2lip_mp4)

    output_filename = 'output.mp4'
    output_mp4 = os.path.join(job_path, output_filename)
    merge(job_path, audio_path, output_mp4)
    return output_mp4


if __name__ == '__main__':
    job_id = sys.argv[1]
    video_url = sys.argv[2]
    audio_url = sys.argv[3]

    output_mp4 = main(job_id, video_url, audio_url)

    from google.colab import files

    files.download(output_mp4)