import gradio as gr
import roop.globals
import roop.core
import roop.processors.frame.core
import os
import shutil

# Ensure roop thinks it's headless so it doesn't try to launch its own GUI
roop.globals.headless = True
roop.globals.execution_providers = ['cpu'] # Default to CPU, can be changed
roop.globals.execution_threads = 1
roop.globals.frame_processors = ['face_swapper']
roop.globals.keep_fps = True
roop.globals.keep_frames = False
roop.globals.skip_audio = False
roop.globals.many_faces = False
roop.globals.reference_face_position = 0
roop.globals.reference_frame_number = 0
roop.globals.similar_face_distance = 0.85
roop.globals.temp_frame_format = 'png'
roop.globals.temp_frame_quality = 100
roop.globals.output_video_encoder = 'libx264'
roop.globals.output_video_quality = 35
roop.globals.max_memory = None


def swap_face(source_file, target_file, enable_enhancer):
    if source_file is None or target_file is None:
        return None
    
    source_path = source_file if isinstance(source_file, str) else source_file.name
    target_path = target_file if isinstance(target_file, str) else target_file.name
    
    # Force strict check: Allow images ONLY.
    target_name, target_ext = os.path.splitext(os.path.basename(target_path))
    if target_ext.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.webp']:
        raise gr.Error("Video processing is disabled on this Space due to high CPU usage. Please use Images only.")
        return None

    output_ext = '.png'
    output_path = os.path.join(os.getcwd(), f"output_{target_name}{output_ext}")

    print(f"Source: {source_path}")
    print(f"Target: {target_path}")
    print(f"Output: {output_path}")
    print(f"Enhancer Enabled: {enable_enhancer}")

    # Set globals
    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    roop.globals.output_path = output_path
    roop.globals.execution_providers = ['cpu']
    
    # Configure processors
    processors = ['face_swapper']
    if enable_enhancer:
        processors.append('face_enhancer')
    roop.globals.frame_processors = processors
    
    # CRITICAL FIX: Clear the cache of frame processors so it reloads based on the new list
    roop.processors.frame.core.FRAME_PROCESSORS_MODULES = []

    # Use correct provider name to avoid ONNX Runtime warnings
    roop.globals.execution_providers = ['CPUExecutionProvider']

    try:
        roop.core.start()
        if os.path.exists(output_path):
            return output_path
        else:
            print("Error: Output file not created.")
            return None
    except Exception as e:
        print(f"Error during processing: {e}")
        return None

# Define the interface
with gr.Blocks(title="Roop Web UI") as demo:
    
    # Custom Banner for AI Guruji
    gr.HTML("""
    <div style="background-color: #0b0f19; padding: 20px; text-align: center; border-radius: 10px; margin-bottom: 20px;">
        <img src="https://i.ibb.co/k2cBZj6q/logo-1.png" alt="AI Guruji Logo" style="width: 120px; height: 120px; border-radius: 50%; display: block; margin: 0 auto;">
        <h1 style="color: white; font-family: sans-serif; margin-top: 10px; font-size: 28px; font-weight: bold;">Al Guruji</h1>
        <div style="margin-top: 15px; display: flex; justify-content: center; gap: 10px;">
            <a href="https://www.youtube.com/@AI_Guruji_97" target="_blank" style="background-color: #ff0000; color: white; padding: 8px 20px; border-radius: 20px; text-decoration: none; font-weight: bold; font-family: sans-serif;">YouTube</a>
            <a href="https://buymeacoffee.com/bdas7906" target="_blank" style="background-color: #FFDD00; color: black; padding: 8px 20px; border-radius: 20px; text-decoration: none; font-weight: bold; font-family: sans-serif;">Buy me a coffee</a>
            <a href="https://www.instagram.com/aiguruji97/" target="_blank" style="background-color: #E1306C; color: white; padding: 8px 20px; border-radius: 20px; text-decoration: none; font-weight: bold; font-family: sans-serif;">Instagram</a>
        </div>
        <div style="margin-top: 20px; color: white; font-family: sans-serif; font-size: 14px;">
            <p>If multiple users are using this space (3-4 users), generation might be slow.</p>
            <p><strong>Duplicate this Space</strong> for faster generation without waiting!</p>
            <p style="margin-top: 10px; font-size: 16px;">❤️ <strong>Please Like this Space if you find it useful!</strong> ❤️</p>
            <div style="display: flex; justify-content: center; align-items: center; margin-top: 10px;">
                <a href="https://huggingface.co/spaces/huggingface/facefusion?duplicate=true" target="_blank" style="display: inline-flex; align-items: center; background-color: white; color: black; padding: 5px 15px; border-radius: 6px; text-decoration: none; font-weight: bold; border: 1px solid #ccc;">
                    <img src="https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-md-dark.svg" alt="Duplicate this Space" style="height: 20px; margin-right: 5px;">
                    Duplicate Space
                </a>
            </div>
        </div>
    </div>
    """)
    
    # gr.Markdown("# Roop Face Swap Web Interface")
    gr.Markdown("### Upload a source face and a target image to swap faces.")
    gr.Markdown("⚠️ **NOTE: Video processing is DISABLED on this Space due to high CPU usage.** ⚠️")
    
    with gr.Row():
        with gr.Column():
            source_input = gr.Image(type="filepath", label="Source Face (Image)")
            # Changed to Strict Image Input
            target_input = gr.Image(type="filepath", label="Target (Image ONLY)")
            enable_enhancer = gr.Checkbox(label="Enable Face Enhancer (GFPGAN)", value=False)
            submit_btn = gr.Button("Swap Face", variant="primary")
        
        with gr.Column():
            result_image = gr.Image(label="Result Image", interactive=False)
            # Video output removed

    submit_btn.click(
        fn=swap_face,
        inputs=[source_input, target_input, enable_enhancer],
        outputs=[result_image]
    )

if __name__ == "__main__":
    demo.launch(inbrowser=True)