import gradio as gr
import torch
from transformers import (
    BlipProcessor, 
    BlipForQuestionAnswering,
    pipeline,
    AutoTokenizer 
)
from modelscope.pipelines import pipeline as ms_pipeline
from modelscope.outputs import OutputKeys
from PIL import Image
import os
import logging
import tempfile
import moviepy.editor as mp

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class MultimodalProcessor:
    def __init__(self):
        self.load_models()
        self.temp_dir = tempfile.mkdtemp()
    
    def load_models(self):
        """Charge les modèles avec gestion d'erreurs"""
        try:
            logger.info("Chargement des modèles...")
            self.blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
            self.blip_model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
            self.audio_transcriber = pipeline("automatic-speech-recognition", 
                                           model="openai/whisper-base")
            
            # Initialisation du pipeline text-to-video ModelScope
            self.video_pipeline = ms_pipeline(
                'text-to-video-synthesis',
                model='damo/text-to-video-synthesis',
                model_revision='v1.0.0'
            )
            
            logger.info("Modèles chargés avec succès")
        except Exception as e:
            logger.error(f"Erreur lors du chargement des modèles: {str(e)}")
            raise

    def transcribe_video(self, video_path):
        """Transcrit une vidéo en utilisant Whisper"""
        try:
            if video_path is None:
                return ""
            
            # Extraction de l'audio de la vidéo
            video = mp.VideoFileClip(video_path)
            audio_path = os.path.join(self.temp_dir, "temp_audio.wav")
            video.audio.write_audiofile(audio_path, logger=None)
            
            # Transcription avec Whisper
            transcription = self.audio_transcriber(audio_path)["text"]
            
            # Nettoyage
            os.remove(audio_path)
            video.close()
            
            return transcription
        except Exception as e:
            logger.error(f"Erreur lors de la transcription vidéo: {str(e)}")
            return "Erreur lors de la transcription vidéo."

    def analyze_image(self, image):
        """Analyse une image et retourne une description"""
        try:
            if image is None:
                return ""
            
            questions = [
                "What is in the picture?",
                "What are the main colors?",
                "What is the setting or background?",
                "What is happening in the image?",
            ]
            
            responses = {}
            for question in questions:
                inputs = self.blip_processor(images=image, text=question, return_tensors="pt")
                outputs = self.blip_model.generate(**inputs)
                answer = self.blip_processor.decode(outputs[0], skip_special_tokens=True)
                responses[question] = answer
            
            description = (
                f"This image shows {responses['What is in the picture?']}. "
                f"The main colors are {responses['What are the main colors?']}. "
                f"The setting is {responses['What is the setting or background?']}. "
                f"In the scene, {responses['What is happening in the image?']}"
            )
            
            return description
        except Exception as e:
            logger.error(f"Erreur lors de l'analyse de l'image: {str(e)}")
            return "Erreur lors de l'analyse de l'image."

    def transcribe_audio(self, audio_path):
        """Transcrit un fichier audio"""
        try:
            if audio_path is None:
                return ""
            return self.audio_transcriber(audio_path)["text"]
        except Exception as e:
            logger.error(f"Erreur lors de la transcription audio: {str(e)}")
            return "Erreur lors de la transcription audio."

    def generate_video(self, prompt):
        """Génère une vidéo à partir d'un prompt"""
        try:
            if not prompt:
                return None
                
            output_path = os.path.join(self.temp_dir, "generated_video.mp4")
            result = self.video_pipeline({
                'text': prompt,
                'output_path': output_path
            })
            
            if os.path.exists(output_path):
                return output_path
            else:
                raise Exception("La vidéo n'a pas été générée correctement")
                
        except Exception as e:
            logger.error(f"Erreur lors de la génération de vidéo: {str(e)}")
            return None

    def process_inputs(self, image, audio, video, text):
        """Traite les entrées multimodales"""
        try:
            combined_input = []
            
            # Analyse de l'image
            if image is not None:
                image_desc = self.analyze_image(image)
                combined_input.append(f"Visual description: {image_desc}")
            
            # Transcription audio
            if audio is not None:
                audio_text = self.transcribe_audio(audio)
                combined_input.append(f"Audio content: {audio_text}")
            
            # Transcription vidéo
            if video is not None:
                video_text = self.transcribe_video(video)
                combined_input.append(f"Video content: {video_text}")
            
            # Ajout du texte additionnel
            if text:
                combined_input.append(f"Additional context: {text}")
            
            # Création du prompt final
            final_prompt = " ".join(combined_input) if combined_input else "Empty scene"
            
            # Génération de la vidéo
            output_video = self.generate_video(final_prompt)
            
            return output_video, final_prompt
            
        except Exception as e:
            logger.error(f"Erreur lors du traitement des entrées: {str(e)}")
            return None, "Une erreur est survenue lors du traitement des entrées."

def create_interface():
    """Crée l'interface Gradio"""
    processor = MultimodalProcessor()
    
    interface = gr.Interface(
        fn=processor.process_inputs,
        inputs=[
            gr.Image(type="pil", label="Télécharger une image"),
            gr.Audio(type="filepath", label="Télécharger un fichier audio"),
            gr.Video(label="Télécharger une vidéo"),
            gr.Textbox(label="Entrez du texte additionnel")
        ],
        outputs=[
            gr.Video(label="Vidéo générée"),
            gr.Textbox(label="Prompt généré")
        ],
        title="Générateur de Vidéo Multimodal",
        description="""
        Cette application combine vos entrées multimodales pour générer une vidéo :
        - Images : analyse et description
        - Audio : transcription
        - Vidéo : transcription
        - Texte : contexte additionnel
        
        Le résultat est une nouvelle vidéo générée à partir de la description combinée.
        """
    )
    
    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()