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

# === Descarga y configuración de Roop ===
if not os.path.exists("roop"):
    subprocess.run(["git", "clone", "https://github.com/s0md3v/roop.git"])
    os.chdir("roop")

    subprocess.run(["pip", "uninstall", "-y", "jax", "jaxlib", "numpy", "tensorflow"])
    subprocess.run(["pip", "install", "numpy==1.23.5", "jax==0.4.13", "jaxlib==0.4.13", "tensorflow==2.12.0"])
    subprocess.run(["pip", "install", "-r", "requirements.txt"])
    
    os.chdir("..")

def procesar_lote(img_rostro, imgs_objetivo):
    input_dir = "roop/batch_input"
    output_dir = "roop/batch_output"
    
    shutil.rmtree(input_dir, ignore_errors=True)
    shutil.rmtree(output_dir, ignore_errors=True)
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    rostro_path = os.path.join("roop", "RB.jpg")
    img_rostro.save(rostro_path)

    for img_file in imgs_objetivo:
        image = Image.open(img_file)
        image.save(os.path.join(input_dir, os.path.basename(img_file)))

    comando = [
        "python", "roop/run.py",
        "--source_face", rostro_path,
        "--target_folder", input_dir,
        "--output_folder", output_dir,
        "--frame-processor", "face_swapper"
    ]
    subprocess.run(comando)

    zip_path = "roop/resultado.zip"
    shutil.make_archive(zip_path.replace(".zip", ""), 'zip', output_dir)
    
    return zip_path

gr.Interface(
    fn=procesar_lote,
    inputs=[
        gr.Image(label="Rostro base (RB.jpg)", type="pil"),
        gr.File(file_types=["image"], label="Imágenes a procesar", file_count="multiple")
    ],
    outputs=gr.File(label="Descargar resultados ZIP"),
    title="Intercambio de rostros por lote con Roop",
    description="Sube una imagen de rostro base y varias imágenes objetivo. El sistema procesará todas y devolverá un ZIP con los resultados.",
).launch()