import os
import subprocess
import time
from datetime import datetime
from functools import partial
from concurrent.futures import ThreadPoolExecutor
from collections import deque
import cv2
from flask import Flask, request, jsonify, send_file
from flask_ngrok2 import run_with_ngrok
from subprocess import call
from tqdm import tqdm
import numpy as np
from ffmpy import FFmpeg

from main import call_wav2lip, call_gfpgan, merge

# !pip install flask flask-ngrok2 pyngrok

app = Flask(__name__)
auth_token = '2XQTtjJqVg4B4ryQVgauTPIGeiK_3JUWVJcMhXwpPxz2sc9KB'
root_dir = '/content/wav2lip-gfpgan'
jobs_dir = os.path.join('/content', 'jobs')


@app.route('/', methods=['GET'])
def index():
    return jsonify({'hello': 'world!'})


# New route to serve files from /content/input
@app.route('/job/<job_id>/<filename>', methods=['GET'])
def download_file(job_id, filename):
    try:
        file_path = os.path.join(jobs_dir, job_id, filename)
        print(file_path)
        return send_file(file_path, as_attachment=True)
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/wav2lip', methods=['POST'])
def wav2lip():
    try:
        # Get uploaded files
        video_file = request.files['video']
        audio_file = request.files['audio']
        job_id = str(int(time.time()))
        job_path = os.path.join(jobs_dir, job_id)
        os.makedirs(job_path, exist_ok=True)

        # Save the files to a temporary directory
        video_path = os.path.join(job_path, video_file.filename)
        audio_path = os.path.join(job_path, audio_file.filename)
        video_file.save(video_path)
        audio_file.save(audio_path)

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

        output_filename = 'output.mp4'
        output_mp4 = os.path.join(job_path, output_filename)
        call_gfpgan(wav2lip_mp4, audio_path, output_mp4)

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

        return jsonify({'url': f'/job/{job_id}/{output_filename}'})

    except Exception as e:
        return jsonify({'error': str(e)}), 500


if __name__ == '__main__':
    run_with_ngrok(app, auth_token=auth_token)
    app.run()


def test():
    # request
    import requests
    ngrok_url = f"http://74c0-34-87-172-60.ngrok-free.app"
    url = f"{ngrok_url}/wav2lip"
    print(url)
    video_path = '/Users/taoluo/Downloads/oIy5B4-vHVw.4.6588496370531551262.0.jpg'
    audio_path = '/Users/taoluo/Downloads/test_audio.mp3'
    files = {'video': ('video.jpg', open(video_path, 'rb')), 'audio': ('audio.mp3', open(audio_path, 'rb'))}
    headers = {'ngrok-skip-browser-warning': 'true'}
    response = requests.post(url, files=files, headers=headers)
    # Print the response
    print(response.json())
    data = response.json()
    print(ngrok_url + data['url'])