import gradio as gr
import cv2
import numpy as np
from PIL import Image
import tempfile
import os

class VideoInpaintingApp:
    def __init__(self):
        self.current_video = None
        self.current_frame = None
        self.fps = 30
        self.total_frames = 0
        
    def load_video(self, video_path):
        """Charge la vidéo et retourne des informations"""
        if video_path is None:
            return None, "Aucune vidéo chargée", 0, 0, 0
        
        cap = cv2.VideoCapture(video_path)
        self.fps = int(cap.get(cv2.CAP_PROP_FPS))
        self.total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = self.total_frames / self.fps if self.fps > 0 else 0
        
        cap.release()
        self.current_video = video_path
        
        info = f"✅ Vidéo chargée: {self.total_frames} frames, {self.fps} FPS, {duration:.2f}s"
        max_frame = max(0, self.total_frames - 1)
        return video_path, info, 0, max_frame, max_frame
    
    def get_frame(self, video_path, frame_number):
        """Récupère une frame spécifique de la vidéo"""
        if video_path is None:
            return None
        
        try:
            cap = cv2.VideoCapture(video_path)
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(frame_number))
            ret, frame = cap.read()
            cap.release()
            
            if ret:
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                self.current_frame = frame_rgb
                return Image.fromarray(frame_rgb)
        except Exception as e:
            print(f"Erreur lors de la lecture de la frame: {e}")
        
        return None
    
    def apply_inpainting(self, video_path, start_frame, end_frame, mask_image, progress=gr.Progress()):
        """Applique l'inpainting sur la plage de frames sélectionnée"""
        if video_path is None:
            return None, "❌ Veuillez charger une vidéo"
        
        if mask_image is None:
            return None, "❌ Veuillez dessiner un masque sur la zone à supprimer"
        
        try:
            progress(0, desc="Préparation...")
            
            # Ouvrir la vidéo
            cap = cv2.VideoCapture(video_path)
            fps = int(cap.get(cv2.CAP_PROP_FPS))
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            # Créer un fichier temporaire pour la sortie
            output_path = tempfile.mktemp(suffix='.mp4')
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
            
            # Extraire le masque
            if isinstance(mask_image, dict) and 'mask' in mask_image:
                mask = np.array(mask_image['mask'])
            else:
                mask = np.array(mask_image)
            
            # Redimensionner le masque si nécessaire
            if mask.shape[:2] != (height, width):
                mask = cv2.resize(mask, (width, height))
            
            # Convertir en masque binaire
            if len(mask.shape) == 3:
                mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
            _, mask_binary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
            
            # Traiter chaque frame
            for frame_idx in range(total_frames):
                ret, frame = cap.read()
                if not ret:
                    break
                
                # Appliquer l'inpainting seulement dans la plage sélectionnée
                if start_frame <= frame_idx <= end_frame:
                    inpainted = cv2.inpaint(frame, mask_binary, 3, cv2.INPAINT_TELEA)
                    out.write(inpainted)
                else:
                    out.write(frame)
                
                # Mise à jour de la progression
                progress((frame_idx + 1) / total_frames, desc=f"Frame {frame_idx + 1}/{total_frames}")
            
            cap.release()
            out.release()
            
            return output_path, "✅ Inpainting terminé avec succès!"
            
        except Exception as e:
            return None, f"❌ Erreur: {str(e)}"

# Créer l'instance de l'application
app = VideoInpaintingApp()

# Interface Gradio
with gr.Blocks(title="Inpainting Vidéo IA", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎬 Inpainting Vidéo avec IA
    ### Supprimez des objets indésirables de vos vidéos
    
    **Instructions:**
    1. Chargez votre vidéo
    2. Sélectionnez la plage de frames avec les curseurs
    3. Dessinez sur la zone à supprimer (pinceau blanc)
    4. Cliquez sur "Supprimer l'objet sélectionné"
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📁 1. Charger la vidéo")
            video_input = gr.Video(label="Votre vidéo")
            video_info = gr.Textbox(label="Informations", interactive=False)
            
            gr.Markdown("### ⏱️ 2. Sélectionner la plage")
            start_frame = gr.Slider(
                minimum=0, 
                maximum=100, 
                value=0, 
                step=1, 
                label="Frame de début"
            )
            end_frame = gr.Slider(
                minimum=0, 
                maximum=100, 
                value=100, 
                step=1, 
                label="Frame de fin"
            )
            current_frame_slider = gr.Slider(
                minimum=0,
                maximum=100,
                value=0,
                step=1,
                label="Prévisualiser la frame"
            )
        
        with gr.Column(scale=2):
            gr.Markdown("### 🎨 3. Dessiner la zone à supprimer")
            frame_preview = gr.Image(
                label="Dessinez sur la zone à supprimer",
                type="pil",
                tool="sketch",
                brush=gr.Brush(colors=["#FFFFFF"], default_size=20)
            )
    
    gr.Markdown("### ⚙️ 4. Lancer l'inpainting")
    process_btn = gr.Button("🚀 Supprimer l'objet sélectionné", variant="primary", size="lg")
    
    status_text = gr.Textbox(label="Statut", interactive=False)
    output_video = gr.Video(label="Vidéo traitée")
    
    # Événements
    video_input.change(
        fn=app.load_video,
        inputs=[video_input],
        outputs=[video_input, video_info, start_frame, end_frame, current_frame_slider]
    )
    
    current_frame_slider.change(
        fn=app.get_frame,
        inputs=[video_input, current_frame_slider],
        outputs=[frame_preview]
    )
    
    process_btn.click(
        fn=app.apply_inpainting,
        inputs=[video_input, start_frame, end_frame, frame_preview],
        outputs=[output_video, status_text]
    )

demo.launch()