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

from themes import IndonesiaTheme  # Impor tema custom dari themes.py


def swap_face(source_file, target_file, doFaceEnhancer):

    source_path = "input.jpg"
    target_path = "target.jpg"

    # Simpan gambar sumber dan target
    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)

    # Pengaturan globals untuk proses
    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
    )
    
    # Pilih processor yang digunakan
    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(
        "[-] Proses dimulai..",
        roop.globals.source_path,
        roop.globals.target_path,
        roop.globals.output_path,
    )

    # Pre-check frame processor
    for frame_processor in get_frame_processors_modules(
        roop.globals.frame_processors
    ):
        if not frame_processor.pre_check():
            return None, "Proses gagal, pre-check tidak berhasil."

    # Mulai proses
    start()

     # Setelah selesai, hapus file sementara
    try:
        os.remove(source_path)
        os.remove(target_path)
        print("[-] File input, target, dan output telah dihapus.")
    except Exception as e:
        print(f"[!] Gagal menghapus file: {e}")
        return None, "Proses berhasil, namun ada masalah saat menghapus file."

    return output_path, "Proses berhasil, 🥹 wajah telah ditukar. 🥰 Mantap banget ya!!."

# CSS untuk styling antarmuka
css = """
#col-left, #col-mid, #col-right {
    margin: 0 auto;
    max-width: 400px;
    padding: 10px;
    border-radius: 15px;
    background-color: #f9f9f9;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
#banner {
    width: 100%;
    text-align: center;
    margin-bottom: 20px;
}
#run-button {
    background-color: #ff4b5c;
    color: white;
    font-weight: bold;
    padding: 10px;
    border-radius: 10px;
    cursor: pointer;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
}
#footer {
    text-align: center;
    margin-top: 20px;
    color: silver;
}
"""

# Membuat antarmuka Gradio dengan tema IndonesiaTheme
with gr.Blocks(css=css, theme=IndonesiaTheme()) as TukarWajah:
        # Tambahkan banner
    gr.HTML("""
        <div style='text-align: center;'>
        <div><h1>Selamat Datang di Aplikasi Tukar Wajah</h1></div>
        <div><p>Unggah gambar sumber dan target untuk melakukan pertukaran wajah. Anda juga dapat menggunakan fitur Face Enhancer jika diinginkan.</p></div>
            <img src='https://i.ibb.co.com/Snwc7Lb/banner-tw.jpg' alt='Banner' style='width: 100%; height: auto;'/>
        </div>
    """)

    
    # Layout dua kolom
    with gr.Row():
        with gr.Column(elem_id="col-left"):
            gr.Markdown("### 🖼️ Unggah Foto Wajah Sumber")
            source_image = gr.Image(label="Gambar Sumber", type="numpy")
        with gr.Column(elem_id="col-right"):
            gr.Markdown("### 🖼️ Unggah Foto Target")
            target_image = gr.Image(label="Gambar Target", type="numpy")

    # Output di bawah kolom
    with gr.Row():
        result_image = gr.Image(label="🖼️ Hasil Pertukaran Wajah")
        status_output = gr.Textbox(label="Status Proses")
    
    # Checkbox dan tombol untuk memulai proses
    gr.Markdown("### ✅ Gunakan fitur ini untuk hasil maksimal.")
    doFaceEnhancer = gr.Checkbox(label="Gunakan Face Enhancer?", value=False)
    run_button = gr.Button(value="⭐ Mulai Proses Pertukaran Wajah ⭐", elem_id="run-button")

    run_button.click(fn=swap_face, inputs=[source_image, target_image, doFaceEnhancer], outputs=[result_image, status_output])

    # Tambahkan footer di bagian bawah
    gr.HTML("""
    <footer id="footer">
        Transfer Energi Semesta Digital © 2024 __drat. | 🇮🇩 Untuk Indonesia Jaya!
    </footer>
    """)


# Menjalankan aplikasi
if __name__ == "__main__":
    TukarWajah.queue(api_open=False).launch(show_api=False)