import os
import shutil
import zipfile
from PIL import Image
import gradio as gr
import uuid
import subprocess

# Directorios
INPUT_DIR = "batch_images"
OUTPUT_DIR = "batch_results"
RESULTS_DIR = "comparaciones"
os.makedirs(INPUT_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)

def face_swap_batch(rb_img, lote_imgs):
    # Limpiar directorios
    shutil.rmtree(INPUT_DIR, ignore_errors=True)
    shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
    shutil.rmtree(RESULTS_DIR, ignore_errors=True)
    os.makedirs(INPUT_DIR, exist_ok=True)
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    os.makedirs(RESULTS_DIR, exist_ok=True)

    # Guardar imagen de rostro base
    rb_path = os.path.join(INPUT_DIR, "RB.jpg")
    rb_img.save(rb_path)

    # Guardar imágenes por lote
    img_paths = []
    for i, img in enumerate(lote_imgs):
        name = f"img_{i}.jpg"
        path = os.path.join(INPUT_DIR, name)
        img.save(path)
        img_paths.append(path)

    # Procesar imágenes con Roop
    for path in img_paths:
        filename = os.path.basename(path)
        output_img = os.path.join(OUTPUT_DIR, f"swapped_{filename}")
        subprocess.run([
            "python", "run.py",
            "-s", rb_path,
            "-t", path,
            "-o", output_img,
            "--execution-provider", "cpu",
            "--frame-processor", "face_swapper", "face_enhancer"
        ])

    # Crear comparaciones lado a lado
    resultados = []
    for path in img_paths:
        filename = os.path.basename(path)
        original = Image.open(path).resize((256, 256))
        generado_path = os.path.join(OUTPUT_DIR, f"swapped_{filename}")
        if not os.path.exists(generado_path):
            continue
        generado = Image.open(generado_path).resize((256, 256))
        combinado = Image.new("RGB", (512, 256))
        combinado.paste(original, (0, 0))
        combinado.paste(generado, (256, 0))
        out_path = os.path.join(RESULTS_DIR, f"comparado_{filename}")
        combinado.save(out_path)
        resultados.append(out_path)

    # Descarga
    if len(resultados) == 1:
        return resultados[0]
    else:
        zip_path = f"resultados_{uuid.uuid4().hex[:6]}.zip"
        with zipfile.ZipFile(zip_path, 'w') as zipf:
            for file in resultados:
                zipf.write(file, os.path.basename(file))
        return zip_path

# Interfaz Gradio
gr.Interface(
    fn=face_swap_batch,
    inputs=[
        gr.Image(label="🧠 Rostro base", type="pil"),
        gr.File(label="🖼️ Imágenes por lote", file_types=["image"], file_count="multiple")
    ],
    outputs=gr.File(label="📦 Descargar ZIP o imagen"),
    title="Deep Fake Lote - Yepo Hz",
    description="Sube una imagen de rostro base y varias imágenes objetivo. El sistema aplicará el cambio de rostro en lote. Resultado visual lado a lado."
).launch()