import gradio as gr
from refacer import Refacer
import ngrok
import time
import threading

# Initialize Refacer
def initialize_refacer():
    global refacer
    print("Initializing Refacer...")
    start_time = time.time()
    refacer = Refacer(force_cpu=False, colab_performance=False)
    print(f"Refacer initialized in {time.time() - start_time:.2f} seconds")

# Ngrok connection
def connect(token, port, options):
    account = None
    if token:
        if ':' in token:
            # token = authtoken:username:password
            token, username, password = token.split(':', 2)
            account = f"{username}:{password}"

    if not options.get('authtoken_from_env'):
        options['authtoken'] = token
    if account:
        options['basic_auth'] = account

    try:
        public_url = ngrok.connect(f"127.0.0.1:{port}", **options).url()
        print(f'ngrok connected to localhost:{port}! URL: {public_url}')
    except Exception as e:
        print(f'ngrok connection aborted: {e}')

# Run reface
def run(*vars):
    video_path = vars[0]
    origins = vars[1:(num_faces + 1)]
    destinations = vars[(num_faces + 1):(num_faces * 2) + 1]
    thresholds = vars[(num_faces * 2) + 1:]

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

    return refacer.reface(video_path, faces)

# UI setup
origin = []
destination = []
thresholds = []
num_faces = 5

from gradio.themes.utils.theme_dropdown import create_theme_dropdown

dropdown, js = create_theme_dropdown()

with gr.Blocks(theme='victorisgeek/gray') as demo:
    with gr.Row(equal_height=True):
        with gr.Column(scale=10):
            gr.Markdown("# Refacer")
    with gr.Row():
        video = gr.Video(label="Original video", format="mp4")
        video2 = gr.Video(label="Refaced video", interactive=False, format="mp4")

    for i in range(num_faces):
        with gr.Tab(f"Face #{i + 1}"):
            with gr.Row():
                origin.append(gr.Image(label="Face to replace"))
                destination.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] + origin + destination + thresholds, outputs=[video2])

# Start initialization in a separate thread to prevent blocking
init_thread = threading.Thread(target=initialize_refacer)
init_thread.start()

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