# -* coding:UTF-8 -*
# !/usr/bin/env python
import numpy as np
import gradio as gr
import roop.globals
from roop.core import (
    start,
    decode_execution_providers,
    suggest_max_memory,
    suggest_execution_threads,
)
from roop.processors.frame.core import get_frame_processors_modules
from roop.utilities import normalize_output_path
import os
from PIL import Image

# Loader HTML and CSS
loader_html = """
<div id="loader" style="
  position: fixed;
  top: 0;
  left: 0;
  width: 100%;
  height: 100%;
  background: rgba(255, 255, 255, 0.9);
  display: none; /* Hidden by default */
  justify-content: center;
  align-items: center;
  z-index: 1000;
">
  <div class="spinner" style="
    width: 50px;
    height: 50px;
    border: 4px solid #f3f3f3;
    border-top: 4px solid #3498db;
    border-radius: 50%;
    animation: spin 1s linear infinite;
  "></div>
</div>

<style>
@keyframes spin {
  0% { transform: rotate(0deg); }
  100% { transform: rotate(360deg); }
}
</style>
"""

# ============================================================
# ADDED: Custom CSS for the button
# ============================================================
button_css = """
<style>
/* Style for the button */
.custom-button {
    background-color: black; /* Black background */
    color: white; /* White text */
    border: none; /* No border */
    padding: 15px 30px; /* Padding */
    font-size: 16px; /* Font size */
    border-radius: 25px; /* Rounded corners (fully rounded) */
    cursor: pointer; /* Pointer cursor on hover */
    transition: background-color 0.3s ease; /* Smooth transition */
}

/* Hover effect */
.custom-button:hover {
    background-color: #333333; /* Slightly lighter black on hover */
}

/* Active (click) effect */
.custom-button:active {
    background-color: #666666; /* Even lighter black on click */
}
</style>
"""

article_text = """
<div style="text-align: center;">
    <p>Leer zelf Generative AI-tools te bouwen en te begrijpen!</p>
    <a href="https://www.glossybranding.com/masterclass">Schrijf je in voor de AI Masterclass</a>
</div>
"""


def swap_face(source_file, target_file, doFaceEnhancer):
    # Show the loader
    show_loader = """
    <script>
    document.getElementById('loader').style.display = 'flex';
    </script>
    """
    
    source_path = "input.jpg"
    target_path = "target.jpg"

    source_image = Image.fromarray(source_file)
    source_image.save(source_path)
    target_image = Image.fromarray(target_file)
    target_image.save(target_path)

    print("source_path: ", source_path)
    print("target_path: ", target_path)

    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    output_path = "output.jpg"
    roop.globals.output_path = normalize_output_path(
        roop.globals.source_path, roop.globals.target_path, output_path
    )
    if doFaceEnhancer == True:
        roop.globals.frame_processors = ["face_swapper", "face_enhancer"]
    else:
        roop.globals.frame_processors = ["face_swapper"]
    roop.globals.headless = True
    roop.globals.keep_fps = True
    roop.globals.keep_audio = True
    roop.globals.keep_frames = False
    roop.globals.many_faces = False
    roop.globals.video_encoder = "libx264"
    roop.globals.video_quality = 18
    roop.globals.max_memory = suggest_max_memory()
    roop.globals.execution_providers = decode_execution_providers(["cuda"])
    roop.globals.execution_threads = suggest_execution_threads()

    print(
        "start process",
        roop.globals.source_path,
        roop.globals.target_path,
        roop.globals.output_path,
    )

    for frame_processor in get_frame_processors_modules(
        roop.globals.frame_processors
    ):
        if not frame_processor.pre_check():
            return

    start()
    
    # Hide the loader
    hide_loader = """
    <script>
    document.getElementById('loader').style.display = 'none';
    </script>
    """
    
    return output_path, hide_loader


# ============================================================
# CHANGES START HERE
# ============================================================

# Use Gradio Blocks for more flexibility
with gr.Blocks() as demo:
    
    # Add the loader HTML
    gr.HTML(loader_html)
    
    # Add the face-swap interface
    with gr.Column():  # Stack everything vertically
        source_image = gr.Image(label="Source Image")
        target_image = gr.Image(label="Target Image")
        do_face_enhancer = gr.Checkbox(label="Face Enhancer?", info="Do face enhancer?")
        submit_button = gr.Button("Swap Faces")
        output_image = gr.Image(label="Output Image")
        gr.HTML(article_text)
    
    # Define the function to call on button click
    submit_button.click(
        fn=swap_face,
        inputs=[source_image, target_image, do_face_enhancer],
        outputs=[output_image, gr.HTML()]  # Output image and HTML (to hide loader)
    )

# ============================================================
# CHANGES END HERE
# ============================================================

# Launch the app
demo.launch(show_api=False)