# app.py - Interface principale Face Swap (Hugging Face Space safe)

import gradio as gr
import cv2
import insightface
from insightface.app import FaceAnalysis
from insightface.model_zoo import get_model
import numpy as np
import os

# ===============================
# 1️⃣ Initialisation FaceAnalysis
# ===============================

app = FaceAnalysis(name="buffalo_l")

# GPU si dispo, sinon CPU
try:
    app.prepare(ctx_id=0, det_size=(640, 640))
    print("✅ FaceAnalysis chargé avec GPU")
except Exception:
    app.prepare(ctx_id=-1, det_size=(640, 640))
    print("⚠️ GPU indisponible → fallback CPU")

# ===============================
# 2️⃣ Téléchargement sécurisé du modèle inswapper
# ===============================

MODEL_PATH = "inswapper_128.onnx"
HF_URL = (
    "https://huggingface.co/ezioruan/inswapper_128.onnx/"
    "resolve/main/inswapper_128.onnx"
)
MIN_BYTES = 200 * 1024 * 1024  # 200 MB minimum (sécurité)

def ensure_inswapper(path: str):
    # Fichier existant mais trop petit = corrompu
    if os.path.exists(path) and os.path.getsize(path) < MIN_BYTES:
        print(f"⚠️ {path} corrompu → suppression")
        os.remove(path)

    # Télécharger si absent
    if not os.path.exists(path):
        print("⬇️ Téléchargement de inswapper_128.onnx depuis Hugging Face...")
        cmd = (
            "wget -L --retry-connrefused --waitretry=1 "
            "--tries=5 --timeout=30 "
            f"-O {path} '{HF_URL}'"
        )
        code = os.system(cmd)
        if code != 0:
            raise RuntimeError("❌ Échec du téléchargement de inswapper_128.onnx")

    # Vérification finale
    size = os.path.getsize(path)
    if size < MIN_BYTES:
        raise RuntimeError(
            f"❌ inswapper_128.onnx invalide (taille {size} bytes)"
        )

    print(f"✅ inswapper_128.onnx prêt ({size // (1024*1024)} MB)")

ensure_inswapper(MODEL_PATH)

# Charger le modèle de swap
swapper = get_model(MODEL_PATH, download=False, download_zip=False)

# ===============================
# 3️⃣ Fonction de face swap
# ===============================

def swap_face(source_img, target_img):
    """
    Swap le visage de source_img vers target_img
    """
    try:
        if source_img is None or target_img is None:
            return None, "❌ Merci de charger deux images"

        # Convertir PIL → OpenCV (BGR)
        source = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
        target = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)

        # Détection des visages
        source_faces = app.get(source)
        target_faces = app.get(target)

        if len(source_faces) == 0:
            return None, "❌ Aucun visage détecté dans l'image source"

        if len(target_faces) == 0:
            return None, "❌ Aucun visage détecté dans l'image cible"

        source_face = source_faces[0]
        result = target.copy()

        # Swap sur chaque visage cible
        for face in target_faces:
            result = swapper.get(
                result,
                face,
                source_face,
                paste_back=True
            )

        # OpenCV → RGB
        result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)

        return result_rgb, "✅ Face swap réussi"

    except Exception as e:
        return None, f"❌ Erreur : {str(e)}"

# ===============================
# 4️⃣ Interface Gradio
# ===============================

with gr.Blocks(title="Face Swapper Pro", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
# 🎭 Face Swapper Pro
### Échangez les visages entre deux images

**Instructions :**
1. Image source = visage à utiliser  
2. Image cible = visage(s) à remplacer  
3. Cliquez sur **Swap Faces**
""")

    with gr.Row():
        source_image = gr.Image(
            label="📸 Image source",
            type="pil",
            height=400
        )
        target_image = gr.Image(
            label="🎯 Image cible",
            type="pil",
            height=400
        )

    swap_btn = gr.Button("🔄 Swap Faces", variant="primary")
    status_text = gr.Textbox(label="Statut", interactive=False)
    result_image = gr.Image(label="✨ Résultat", height=500)

    swap_btn.click(
        fn=swap_face,
        inputs=[source_image, target_image],
        outputs=[result_image, status_text]
    )

    gr.Markdown("""
---
💡 **Conseils**
- Visages bien éclairés
- Angle similaire = meilleur rendu
- Fonctionne avec plusieurs visages sur la cible
""")

# ===============================
# 5️⃣ Lancement
# ===============================

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