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

# Initialisation sécurisée
try:
    face_app = FaceAnalysis(name="buffalo_l")
    face_app.prepare(ctx_id=0, det_size=(640, 640))
except Exception as e:
    print(f"ERREUR initialisation : {str(e)}")
    face_app = None

def safe_face_swap(source_img, target_img):
    try:
        if face_app is None:
            raise ValueError("Le modèle de détection n'est pas chargé")
        
        # Conversion robuste
        source = np.array(source_img, dtype=np.uint8)
        target = np.array(target_img, dtype=np.uint8)
        
        # Vérification des dimensions
        if source.ndim != 3 or target.ndim != 3:
            raise ValueError("Format d'image invalide")
            
        # Détection sécurisée
        source_faces = face_app.get(cv2.cvtColor(source, cv2.COLOR_RGB2BGR))
        target_faces = face_app.get(cv2.cvtColor(target, cv2.COLOR_RGB2BGR))
        
        if not source_faces:
            raise ValueError("Aucun visage détecté dans l'image source")
        if not target_faces:
            raise ValueError("Aucun visage détecté dans l'image cible")
            
        # Swap simplifié mais stable
        src_face = source_faces[0].bbox.astype(int)
        tgt_face = target_faces[0].bbox.astype(int)
        
        result = target.copy()
        result[tgt_face[1]:tgt_face[3], tgt_face[0]:tgt_face[2]] = \
            cv2.resize(source[src_face[1]:src_face[3], src_face[0]:src_face[2]], 
                     (tgt_face[2]-tgt_face[0], tgt_face[3]-tgt_face[1]))
        
        return result
    except Exception as e:
        print(f"ERREUR traitement : {str(e)}")
        # Retourne l'image cible avec message d'erreur
        if target_img is not None:
            return target_img
        return None

# Interface renforcée
with gr.Blocks(title="FaceSwap Pro") as app:
    gr.Markdown("""
    ## 🔄 FaceSwap Professionnel
    *Échange de visages stable et sécurisé*
    """)
    
    with gr.Row():
        with gr.Column():
            src = gr.Image(label="Source", type="numpy")
            tgt = gr.Image(label="Cible", type="numpy")
            btn = gr.Button("Exécuter", variant="primary")
        result = gr.Image(label="Résultat", interactive=False)
    
    # Gestion des erreurs dans l'UI
    error_box = gr.Textbox(visible=False, label="Erreur")
    
    def wrapped_face_swap(src_img, tgt_img):
        try:
            return safe_face_swap(src_img, tgt_img), ""
        except Exception as e:
            return None, f"Erreur : {str(e)}"
    
    btn.click(
        fn=wrapped_face_swap,
        inputs=[src, tgt],
        outputs=[result, error_box]
    )

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        debug=True
    )