import gradio as gr
import cv2
import numpy as np
from insightface.app import FaceAnalysis
import tempfile
from moviepy.editor import VideoFileClip

# Initialisation du modèle
face_swapper = FaceAnalysis(name="buffalo_l")
face_swapper.prepare(ctx_id=0, det_size=(640, 640))

def process_video(source_img, target_video):
    try:
        # Fichiers temporaires
        temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        
        # Traitement de la vidéo
        source_face = face_swapper.get(np.array(source_img))[0]
        
        cap = cv2.VideoCapture(target_video)
        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        
        out = cv2.VideoWriter(temp_output.name, 
                             cv2.VideoWriter_fourcc(*'mp4v'), 
                             fps, 
                             (frame_width, frame_height))
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            target_faces = face_swapper.get(frame)
            if target_faces:
                frame = swap_face(source_face, target_faces[0], frame)
            
            out.write(frame)
        
        cap.release()
        out.release()
        
        return temp_output.name
        
    except Exception as e:
        print(f"ERREUR: {str(e)}")
        return None

def swap_face(source_face, target_face, frame):
    # Algorithme professionnel d'échange
    src_bbox = source_face.bbox.astype(int)
    tgt_bbox = target_face.bbox.astype(int)
    
    # Extraction et ajustement du visage source
    src_face = cv2.resize(source_face.img[tgt_bbox[1]:tgt_bbox[3], tgt_bbox[0]:tgt_bbox[2]], 
                         (tgt_bbox[2]-tgt_bbox[0], tgt_bbox[3]-tgt_bbox[1]))
    
    # Fusion réaliste
    mask = np.zeros_like(src_face)
    cv2.circle(mask, 
              (mask.shape[1]//2, mask.shape[0]//2), 
              min(mask.shape)//2, 
              (255,255,255), -1)
    
    center = ((tgt_bbox[0]+tgt_bbox[2])//2, (tgt_bbox[1]+tgt_bbox[3])//2)
    output = cv2.seamlessClone(
        src_face, frame, mask, center, cv2.NORMAL_CLONE)
    
    return output

# Interface optimisée pour vidéo
with gr.Blocks() as app:
    gr.Markdown("## 🎬 VideoFaceSwap Pro")
    
    with gr.Row():
        src_img = gr.Image(label="Visage source", type="numpy")
        tgt_video = gr.Video(label="Vidéo cible")
        
    btn = gr.Button("Traiter la vidéo", variant="primary")
    output_video = gr.Video(label="Résultat")
    
    btn.click(
        fn=process_video,
        inputs=[src_img, tgt_video],
        outputs=output_video
    )

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