import gradio as gr
from pyngrok import ngrok
import os
import torch
from google.colab import drive
import requests

# Ngrok authentication
ngrok.set_auth_token("2oStBdyDH75Ike6t5jLv9AeOVk2_6eX2eNR1hb9sE5APN17Ky")

# Hugging Face repository URLs
HUGGINGFACE_REPO_URL = "https://huggingface.co/spaces/ibrahim313/Lipsing/tree/main"
UPLOAD_FOLDER = '/content/testing/result_output'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Function to download files from Hugging Face
def download_from_huggingface(file_path):
    url = f"{HUGGINGFACE_REPO_URL}/{file_path}"
    local_path = f"/content/lipsync-lab/{file_path}"
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    response = requests.get(url)
    with open(local_path, "wb") as file:
        file.write(response.content)
    return local_path

# Download necessary files from Hugging Face
if not os.path.exists('/content/lipsync-lab'):
    os.makedirs('/content/lipsync-lab')
    download_from_huggingface("checkpoints/mobilenet.pth")
    download_from_huggingface("models/wav2lip.py")
    # Add any additional files needed for the model here

# Mount Google Drive (if in Colab)
if not os.path.exists('/content/drive'):
    drive.mount('/content/drive')

# Define Gradio functions
def setup_environment():
    # Install additional dependencies
    os.system('pip install batch_face --quiet')
    os.system('pip install basicsr==1.4.2 --quiet')
    os.system('pip install gfpgan --quiet')
    os.system('python /content/lipsync-lab/install.py')
    return "Setup complete. You can now upload your files."

def upload_audio(audio_file):
    if audio_file is not None:
        audio_path = os.path.join(UPLOAD_FOLDER, audio_file.name)
        audio_file.save(audio_path)
        return f"Audio file '{audio_file.name}' uploaded successfully!"
    else:
        return "No audio file uploaded."

def upload_video(video_file):
    if video_file is not None:
        video_path = os.path.join(UPLOAD_FOLDER, video_file.name)
        video_file.save(video_path)
        return f"Video file '{video_file.name}' uploaded successfully!"
    else:
        return "No video file uploaded."

# Define the Gradio interface layout
with gr.Blocks() as app:
    gr.Markdown("# AI Lip-Sync Application")
    
    with gr.Tab("Setup"):
        gr.Markdown("Click the button below to set up the environment.")
        setup_button = gr.Button("Setup Environment")
        setup_output = gr.Textbox(label="Setup Status")
        setup_button.click(setup_environment, outputs=setup_output)
    
    with gr.Tab("Upload Files"):
        gr.Markdown("Upload your audio and video files.")
        audio_file = gr.File(label="Upload Audio File")
        video_file = gr.File(label="Upload Video File")
        upload_audio_button = gr.Button("Upload Audio")
        upload_video_button = gr.Button("Upload Video")
        audio_output = gr.Textbox(label="Audio Upload Status")
        video_output = gr.Textbox(label="Video Upload Status")
        
        # Set actions for file uploads
        upload_audio_button.click(upload_audio, inputs=audio_file, outputs=audio_output)
        upload_video_button.click(upload_video, inputs=video_file, outputs=video_output)
    
    with gr.Tab("Settings"):
        gr.Markdown("Settings tab for further configurations (if needed).")
        # Add settings options here if required

# Launch Gradio app and create ngrok tunnel
public_url = ngrok.connect(7860)
print("Public URL:", public_url)
app.launch()