import os
import shutil
import sys
import requests
from gradio_client import Client

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')
    print('='*40)
    call_wav2lip(video_path, audio_path, wav2lip_mp4)
    print('='*40)
    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


def generate_audio(text):
    client = Client("https://digitalxingtong-xingtong-3da-bert-vits2.hf.space/")
    result = client.predict(
        text,  # str  in 'Text' Textbox component
        "ign",  # str (Option from: ['ign']) in 'Speaker' Dropdown component
        0.2,  # int | float (numeric value between 0 and 1) in '语调变化' Slider component
        0.6,  # int | float (numeric value between 0.1 and 1.5) in '感情变化' Slider component
        0.8,  # int | float (numeric value between 0.1 and 1.4) in '音节发音长度变化' Slider component
        1,  # int | float (numeric value between 0.1 and 2) in '语速' Slider component
        fn_index=0
    )
    print(result)
    success, audio_file, ogg_file = result
    print(success, audio_file, ogg_file)
    shutil.move(audio_file)

if __name__ == '__main__':
    """
    ! python run.py 2 'http://1.com.1.jpg' /content/wav2lip-gfpgan/inputs/kimk_audio.mp3
    """
    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)