import os
import tempfile
from pathlib import Path
from typing import Tuple, Optional, List

import numpy as np
import gradio as gr
from PIL import Image

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


def setup_roop_config(
    source_path: str,
    target_path: str,
    output_path: str,
    use_face_enhancer: bool = False,
    execution_providers: List[str] = ["cuda"],
) -> None:
    """
    Configure roop global settings for face swapping.
    
    Args:
        source_path: Path to the source image
        target_path: Path to the target image
        output_path: Path for the output image
        use_face_enhancer: Whether to use face enhancer
        execution_providers: List of execution providers
    """
    # Set paths
    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    roop.globals.output_path = normalize_output_path(
        source_path, target_path, output_path
    )
    
    # Set processors
    roop.globals.frame_processors = ["face_swapper", "face_enhancer"] if use_face_enhancer else ["face_swapper"]
    
    # Set other configurations
    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(execution_providers)
    roop.globals.execution_threads = suggest_execution_threads()


def swap_face(
    source_img: np.ndarray, 
    target_img: np.ndarray, 
    use_face_enhancer: bool = False,
    execution_provider: str = "cuda"
) -> Optional[np.ndarray]:
    """
    Swap faces between source and target images.
    
    Args:
        source_img: Source image as numpy array
        target_img: Target image as numpy array
        use_face_enhancer: Whether to enhance the face after swapping
        execution_provider: Hardware acceleration provider (cuda, cpu, etc.)
        
    Returns:
        Resulting image as numpy array or None if processing failed
    """
    try:
        # Create temporary directory for processing
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_dir_path = Path(temp_dir)
            
            # Save input images to temporary files
            source_path = str(temp_dir_path / "source.jpg")
            target_path = str(temp_dir_path / "target.jpg")
            output_path = str(temp_dir_path / "output.jpg")
            
            Image.fromarray(source_img).save(source_path)
            Image.fromarray(target_img).save(target_path)
            
            # Configure roop
            setup_roop_config(
                source_path=source_path,
                target_path=target_path,
                output_path=output_path,
                use_face_enhancer=use_face_enhancer,
                execution_providers=[execution_provider]
            )
            
            # Check if processors are available
            for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
                if not frame_processor.pre_check():
                    raise RuntimeError(f"Frame processor {frame_processor.__name__} failed pre-check")
            
            # Process the face swap
            start()
            
            # Return the result if file exists
            if os.path.isfile(output_path):
                return np.array(Image.open(output_path))
            else:
                raise FileNotFoundError("Output file was not created")
                
    except Exception as e:
        gr.Warning(f"Face swap failed: {str(e)}")
        return None


# UI Components
TITLE = "Face Swap"
DESCRIPTION = """
Upload your source and target images to swap faces. 
Optionally, use the face enhancer feature for HD Results.
"""

FOOTER = """
<div style="text-align: center; margin-top: 20px;">
Poop poop!
</div>
"""

# Create Gradio app with improved UI
with gr.Blocks(title=TITLE, theme=gr.themes.Soft()) as app:
    gr.Markdown(f"# {TITLE}")
    gr.Markdown(DESCRIPTION)
    
    with gr.Row():
        with gr.Column():
            source_image = gr.Image(label="Source Face", type="numpy")
        with gr.Column():
            target_image = gr.Image(label="Target Image", type="numpy")
    
    with gr.Row():
        with gr.Column():
            use_enhancer = gr.Checkbox(
                label="Use Face Enhancer", 
                value=False,
                info="Apply face enhancement for better quality (slower)"
            )
        with gr.Column():
            provider = gr.Dropdown(
                label="Execution Provider",
                choices=["cuda", "cpu", "coreml", "directml", "openvino"],
                value="cuda",
                info="Hardware acceleration (CUDA recommended for NVIDIA GPUs)"
            )
    
    with gr.Row():
        swap_btn = gr.Button("Swap Face", variant="primary")
    
    output_image = gr.Image(label="Result")
    
    swap_btn.click(
        fn=swap_face,
        inputs=[source_image, target_image, use_enhancer, provider],
        outputs=output_image
    )
    
    gr.HTML(FOOTER)

if __name__ == "__main__":
    app.launch(share=True)