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
import uuid

# Global variable to store history
history = []

def swap_face(source_file, target_files, doFaceEnhancer):
    global history
    output_paths = []

    source_path = f"input_{uuid.uuid4()}.jpg"
    source_image = Image.fromarray(source_file)
    source_image.save(source_path)

    for target_file in target_files:
        target_path = target_file.name
        
        print(f"Processing image: {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 = f"output_{uuid.uuid4()}.jpg"
        roop.globals.output_path = normalize_output_path(
            roop.globals.source_path, roop.globals.target_path, output_path
        )
        if doFaceEnhancer:
            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()
        output_paths.append(output_path)

    # Add to history
    history.extend(output_paths)
    
    return output_paths, gr.update(value=history)

def get_history():
    return history


with gr.Blocks() as app:

    with gr.Tab("Face Swap"):
        with gr.Row():
            source_input = gr.Image(label="Source Image")
            target_input = gr.File(label="Target Image(s)", file_count="multiple")
        
        enhance_checkbox = gr.Checkbox(label="Use Face Enhancer", info="Apply face enhancement")
        swap_button = gr.Button("Swap Faces")
        output_gallery = gr.Gallery(label="Output Images")
    
    with gr.Tab("History"):
        history_gallery = gr.Gallery(label="Previous Outputs")
        refresh_button = gr.Button("Refresh History")
    
    swap_button.click(
        fn=swap_face,
        inputs=[source_input, target_input, enhance_checkbox],
        outputs=[output_gallery, history_gallery]
    )
    
    refresh_button.click(
        fn=get_history,
        inputs=[],
        outputs=[history_gallery]
    )

app.launch()