import os
import cv2
import copy
import spaces
import gradio as gr
import insightface
import onnxruntime
import numpy as np
from PIL import Image
from typing import List, Union

# ─── CONFIGURACIÓN ─────────────────────────────────────────
MODEL_PATH = "./inswapper_128.onnx"
DET_SIZE = (320, 320)

# ─── CARGA DE MODELOS A NIVEL DE MÓDULO ────────────────────
# ZeroGPU: los modelos se cargan en 'cuda' aquí.
# Fuera de @spaces.GPU usa emulación CUDA; dentro, GPU real.

# Face swapper (ONNX - GPU)
face_swapper = insightface.model_zoo.get_model(MODEL_PATH)

# Face analyser (InsightFace - CPU para evitar conflictos ONNX CUDA en ZeroGPU)
face_analyser = insightface.app.FaceAnalysis(
    name="buffalo_l",
    root="./checkpoints",
    providers=["CPUExecutionProvider"]
)
face_analyser.prepare(ctx_id=0, det_size=DET_SIZE)

# ─── FUNCIONES AUXILIARES ──────────────────────────────────
def get_many_faces(frame: np.ndarray):
    """Obtiene caras ordenadas de izquierda a derecha"""
    try:
        face = face_analyser.get(frame)
        return sorted(face, key=lambda x: x.bbox[0])
    except (IndexError, TypeError):
        return None

def swap_face(source_faces, target_faces, source_index, target_index, temp_frame):
    """Pega la cara fuente en la imagen objetivo"""
    source_face = source_faces[source_index]
    target_face = target_faces[target_index]
    return face_swapper.get(temp_frame, target_face, source_face, paste_back=True)

# ─── FUNCIÓN GPU (decorada para ZeroGPU) ───────────────────
@spaces.GPU(duration=90)
def process_image(source_img: Image.Image, target_img: Image.Image):
    """
    Pipeline principal de face swapping.
    Corre en GPU Zero con duración de 90s (detección CPU + swap GPU).
    """
    if source_img is None or target_img is None:
        return None, "Faltan imágenes. Sube ambas."

    # Convertir a BGR para OpenCV
    target_cv = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
    source_cv = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)

    # Detectar caras (CPU - más estable en ZeroGPU)
    target_faces = get_many_faces(target_cv)
    source_faces = get_many_faces(source_cv)

    if target_faces is None or len(target_faces) == 0:
        return None, "No se detectaron caras en la imagen objetivo"

    if source_faces is None or len(source_faces) == 0:
        return None, "No se detectaron caras en la imagen fuente"

    num_target = len(target_faces)
    num_source = len(source_faces)

    temp_frame = copy.deepcopy(target_cv)

    # Lógica de reemplazo
    if num_source == 1:
        # Una cara fuente -> reemplazar todas las caras objetivo
        for i in range(num_target):
            temp_frame = swap_face(source_faces, target_faces, 0, i, temp_frame)
    else:
        # Múltiples caras fuente -> mapeo 1 a 1
        iterations = min(num_source, num_target)
        for i in range(iterations):
            temp_frame = swap_face(source_faces, target_faces, i, i, temp_frame)

    # Convertir de vuelta a RGB
    result_img = Image.fromarray(cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB))
    return result_img, f"✅ Swap completado: {num_target} cara(s) reemplazada(s)"

# ─── INTERFAZ GRADIO ───────────────────────────────────────
with gr.Blocks(title="Swapperface - Face Swap", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎭 Swapperface
    ### Face Swapper con InsightFace + ZeroGPU
    Sube una imagen fuente (cara a copiar) y una imagen objetivo (donde pegar la cara).
    """)

    with gr.Row():
        with gr.Column():
            source_input = gr.Image(
                label="Imagen Fuente (Cara a copiar)",
                type="pil",
                image_mode="RGB",
                height=400
            )
        with gr.Column():
            target_input = gr.Image(
                label="Imagen Objetivo (Donde pegar)",
                type="pil",
                image_mode="RGB",
                height=400
            )

    swap_btn = gr.Button("🔄 Realizar Face Swap", variant="primary")

    with gr.Row():
        output_image = gr.Image(label="Resultado", type="pil", height=400)
        output_text = gr.Textbox(label="Estado", interactive=False)

    # Ejemplos (cache desactivado para ZeroGPU - no hay GPU en startup)
    gr.Examples(
        examples=[
            ["./examples/source1.jpg", "./examples/target1.jpg"],
            ["./examples/source2.jpg", "./examples/target2.jpg"],
        ],
        inputs=[source_input, target_input],
        outputs=[output_image, output_text],
        fn=process_image,
        cache_examples=False,
    )

    swap_btn.click(
        fn=process_image,
        inputs=[source_input, target_input],
        outputs=[output_image, output_text]
    )

    gr.Markdown("""
    ---
    ⚡ **ZeroGPU**: La inferencia corre en GPU dinámica. Puede haber cola si hay muchos usuarios.
    """)

if __name__ == "__main__":
    demo.queue().launch()