import io
import gradio as gr
from refacer import Refacer
import os
import requests
import tempfile
import subprocess

# Hugging Face URL to download the model
model_url = "https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
model_path = "/home/user/app/inswapper_128.onnx"  # Absolute path for the model in your environment

# Function to download the model if not exists
def download_model():
    if not os.path.exists(model_path):
        print("Downloading the inswapper_128.onnx model...")
        response = requests.get(model_url)
        if response.status_code == 200:
            with open(model_path, 'wb') as f:
                f.write(response.content)
            print("Model downloaded successfully!")
        else:
            print(f"Error: Model download failed. Status code: {response.status_code}")
    else:
        print("Model already exists.")

# Download the model when the script runs
download_model()

# Initialize the Refacer class
refacer = Refacer(force_cpu=False, colab_performance=False)

# Run function for refacing video
def run(video_path, *vars):
    origins = vars[:5]
    destinations = vars[5:10]
    thresholds = vars[10:]

    faces = []
    for k in range(5):
        if origins[k] is not None and destinations[k] is not None:
            faces.append({
                'origin': origins[k],
                'destination': destinations[k],
                'threshold': thresholds[k]
            })

    # Call refacer to process video and get refaced video path
    # Change refaced video output path to /home/user/app/out
    output_path = "/home/user/app/out/output_video.mp4"
    refaced_video_path = refacer.reface(video_path, faces)
    print(f"Refaced video can be found at {refaced_video_path}")
    
    # Use tempfile to create a temporary video file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
        temp_file_path = temp_file.name
        
        # Run ffmpeg to process the video and write the result to the temporary file
        ffmpeg_command = [
            "ffmpeg", "-i", refaced_video_path, "-c:v", "libx264", 
            "-crf", "23", "-preset", "fast", temp_file_path
        ]
        
        try:
            subprocess.run(ffmpeg_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        except subprocess.CalledProcessError as e:
            print(f"ffmpeg error: {e.stderr.decode()}")  # Print stderr output for debugging

        # Read the temporary file and convert it to a BytesIO buffer
        with open(temp_file_path, "rb") as f:
            video_buffer = io.BytesIO(f.read())
        
        return video_buffer  # Gradio will handle the video display

# Prepare Gradio components
with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# Refacer")
    with gr.Row():
        video = gr.Video(label="Original video", format="mp4")
        video2 = gr.Video(label="Refaced video", interactive=False, format="mp4")

    origins, destinations, thresholds = [], [], []
    for i in range(5):
        with gr.Tab(f"Face #{i+1}"):
            with gr.Row():
                origins.append(gr.Image(label="Face to replace"))
                destinations.append(gr.Image(label="Destination face"))
            with gr.Row():
                thresholds.append(gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.2))
    
    with gr.Row():
        button = gr.Button("Reface", variant="primary")

    button.click(fn=run, inputs=[video] + origins + destinations + thresholds, outputs=[video2])

# Launch the Gradio app
demo.queue().launch(show_error=True, server_name="0.0.0.0", server_port=7860)