# -*- coding:UTF-8 -*-
#!/usr/bin/env python
import spaces
import numpy as np
import gradio as gr
import gradio.exceptions
import roop.globals
from roop.core import (
    start,
    decode_execution_providers,
)
from roop.processors.frame.core import get_frame_processors_modules
from roop.utilities import normalize_output_path
import os
import random
from PIL import Image
import onnxruntime as ort
import cv2
from roop.face_analyser import get_one_face

# Set up the application theme and styling
custom_css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.disclaimer {
    background-color: #f8d7da;
    border: 1px solid #f5c6cb;
    color: #721c24;
    padding: 10px;
    margin-bottom: 20px;
    border-radius: 5px;
}
.instructions {
    background-color: #e2f0d9;
    border: 1px solid #c5e0b4;
    padding: 10px;
    margin-bottom: 20px;
    border-radius: 5px;
}
"""

@spaces.GPU
def swap_face(source_file, target_file, doFaceEnhancer):
    """
    Swap faces between source and target images.
    
    Args:
        source_file: Image containing the face to be used
        target_file: Image where the face will be swapped onto
        doFaceEnhancer: Whether to enhance the resulting face
        
    Returns:
        Path to the output image with the swapped face
    """
    # Create a temporary directory for processing
    session_dir = "temp"
    os.makedirs(session_dir, exist_ok=True)
    
    # Generate random filenames to avoid conflicts
    source_filename = f"source_{random.randint(1000, 9999)}.jpg"
    target_filename = f"target_{random.randint(1000, 9999)}.jpg"
    output_filename = f"output_{random.randint(1000, 9999)}.jpg"
    
    # Save uploaded images to the temporary directory
    source_path = os.path.join(session_dir, source_filename)
    target_path = os.path.join(session_dir, target_filename)
    
    source_image = Image.fromarray(source_file)
    source_image.save(source_path)
    
    target_image = Image.fromarray(target_file)
    target_image.save(target_path)
    
    print("Processing source image:", source_path)
    print("Processing target image:", target_path)
    
    # Check if a face is detected in the source image
    source_face = get_one_face(cv2.imread(source_path))
    if source_face is None:
        raise gradio.exceptions.Error("No face detected in the source image. Please upload a clear photo with a face.")
    
    # Check if a face is detected in the target image
    target_face = get_one_face(cv2.imread(target_path))
    if target_face is None:
        raise gradio.exceptions.Error("No face detected in the target image. Please upload a clear photo with a face.")
    
    # Set up output path
    output_path = os.path.join(session_dir, output_filename)
    normalized_output_path = normalize_output_path(source_path, target_path, output_path)
    
    # Determine which processors to use
    frame_processors = ["face_swapper", "face_enhancer"] if doFaceEnhancer else ["face_swapper"]
    
    # Validate processors
    for frame_processor in get_frame_processors_modules(frame_processors):
        if not frame_processor.pre_check():
            print(f"Pre-check failed for {frame_processor}")
            raise gradio.exceptions.Error(f"The {frame_processor} module failed to initialize. Please try again.")
    
    # Configure roop settings
    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    roop.globals.output_path = normalized_output_path
    roop.globals.frame_processors = frame_processors
    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.execution_providers = decode_execution_providers(['cpu'])
    roop.globals.reference_face_position = 0
    roop.globals.similar_face_distance = 0.6
    roop.globals.max_memory = 60
    roop.globals.execution_threads = 8
    
    # Start the face swapping process
    start()
    
    return normalized_output_path

# Create the Gradio interface with improved descriptions and layout
with gr.Blocks(css=custom_css) as app:
    gr.Markdown("# AI Face Swap Tool")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            <div class="disclaimer">
                <h3>⚠️ Ethical Usage Guidelines</h3>
                <p>This tool is provided for educational and creative purposes only. Please use responsibly:</p>
                <ul>
                    <li>Do not use to create misleading or deceptive content</li>
                    <li>Do not use for harassment or to misrepresent individuals</li>
                    <li>Always get consent when using someone's likeness</li>
                    <li>Be aware that creating deepfakes may violate laws in some jurisdictions</li>
                </ul>
            </div>
            """)
            
            gr.Markdown("""
            <div class="instructions">
                <h3>📝 How to Use This Tool</h3>
                <ol>
                    <li><b>Source Image:</b> Upload a clear photo of the face you want to use</li>
                    <li><b>Target Image:</b> Upload the photo where you want to place the face</li>
                    <li><b>Face Enhancer:</b> Toggle this option to improve the quality of the swapped face</li>
                    <li>Click the "Swap Face" button and wait for processing to complete</li>
                </ol>
                <p><b>Tips for best results:</b></p>
                <ul>
                    <li>Use high-resolution images with clear, front-facing faces</li>
                    <li>Ensure good lighting in both photos</li>
                    <li>Similar face angles work best</li>
                    <li>Processing may take 10-30 seconds depending on image size</li>
                </ul>
            </div>
            """)
    
    with gr.Row():
        with gr.Column():
            source_image = gr.Image(label="Source Face", info="Upload a photo with the face you want to use")
        with gr.Column():
            target_image = gr.Image(label="Target Image", info="Upload the photo where you want to place the face")
    
    with gr.Row():
        face_enhancer = gr.Checkbox(
            label="Apply Face Enhancement", 
            info="Enable this to improve the quality and realism of the result",
            value=True
        )
    
    with gr.Row():
        swap_button = gr.Button("Swap Face", variant="primary")
    
    output_image = gr.Image(label="Result")
    
    swap_button.click(
        fn=swap_face,
        inputs=[source_image, target_image, face_enhancer],
        outputs=output_image
    )
    
    gr.Markdown("""
    ### How It Works
    This application uses AI to detect faces in both images, then carefully transplants the face from 
    the source image onto the target image while preserving natural appearance. The optional face 
    enhancement feature uses additional AI processing to improve skin texture and details.
    """)

# Launch the application
app.launch()