# -*- coding: utf-8 -*-
"""
Educational AI Assistant for Arabic Lessons - Hugging Face Spaces Version
Spécialement optimisée pour la reconnaissance vocale arabe
Version ultra-robuste avec protection contre les erreurs d'index
"""

import gradio as gr
import torch
import os
import tempfile
import warnings
import gc
import time

warnings.filterwarnings("ignore")

# Fonction pour libérer la mémoire GPU
def clear_gpu_memory():
    """Libère la mémoire GPU."""
    try:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            gc.collect()
            print("Mémoire GPU libérée")
    except Exception as e:
        print(f"Erreur lors de la libération de la mémoire GPU: {e}")

# Fonction pour extraire l'audio d'une vidéo
def extract_audio(video_path):
    """Extrait l'audio d'un fichier vidéo."""
    try:
        import ffmpeg
        audio_path = tempfile.mktemp(suffix=".wav")
        print(f"Extraction de l'audio depuis {video_path} vers {audio_path}")
        
        # Commande ffmpeg simplifiée
        (ffmpeg
         .input(video_path)
         .output(audio_path, 
                acodec='pcm_s16le',  # Format PCM non compressé
                ac=1,               # Mono
                ar='16k')           # 16kHz (optimal pour Whisper)
         .run(capture_stdout=True, capture_stderr=True, overwrite_output=True))
        
        print("Extraction audio réussie")
        return audio_path
    except Exception as e:
        print(f"Erreur lors de l'extraction audio: {e}")
        return None

# Fonction pour transcrire l'audio
def transcribe_audio(audio_path, progress=None):
    """Transcrit l'audio arabe."""
    if not audio_path:
        return "Veuillez télécharger un fichier audio ou vidéo."
    
    if progress:
        try:
            progress(0.3, "Chargement du modèle ASR spécialisé pour l'arabe...")
        except Exception as e:
            print(f"Erreur lors de la mise à jour de la progression: {e}")
    
    try:
        from transformers import pipeline
        
        # Utiliser un modèle spécialisé pour l'arabe
        model_name = "Salama1429/KalemaTech-Arabic-STT-ASR-based-on-Whisper-Small"
        
        # Déterminer l'appareil à utiliser
        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        print(f"Utilisation de l'appareil: {device}")
        
        # Paramètres ultra-légers pour éviter les problèmes de mémoire
        asr = pipeline(
            "automatic-speech-recognition",
            model=model_name,
            chunk_length_s=10,  # Chunks courts
            batch_size=4,       # Petit batch
            device=device,
        )
        
        if progress:
            try:
                progress(0.5, "Transcription en cours...")
            except Exception as e:
                print(f"Erreur lors de la mise à jour de la progression: {e}")
        
        print(f"Transcription du fichier audio: {audio_path}")
        
        # Ajout d'un bloc try/except spécifique pour la transcription
        try:
            result = asr(audio_path, generate_kwargs={"language": "arabic"})
            
            # Vérification robuste du résultat avec protection contre les erreurs d'index
            if result and isinstance(result, dict) and "text" in result:
                transcription = result["text"]
                print("Transcription réussie")
            else:
                print("Format de résultat inattendu")
                transcription = "La transcription n'a pas produit de texte reconnaissable."
        except IndexError as idx_err:
            print(f"Erreur d'index pendant la transcription: {idx_err}")
            transcription = "La transcription a échoué en raison d'un problème avec le fichier audio. Essayez un fichier plus court ou de meilleure qualité."
        except Exception as inner_err:
            print(f"Erreur interne pendant la transcription: {inner_err}")
            transcription = f"La transcription a rencontré un problème: {str(inner_err)}"
        
        # Nettoyage des fichiers temporaires
        if "temp" in audio_path and os.path.exists(audio_path):
            try:
                os.remove(audio_path)
                print(f"Fichier audio temporaire supprimé: {audio_path}")
            except Exception as e:
                print(f"Erreur lors de la suppression du fichier temporaire: {e}")
        
        return transcription if transcription else "La transcription a échoué ou l'audio était silencieux."
    
    except Exception as e:
        print(f"Erreur pendant la transcription: {e}")
        # Nettoyage des fichiers temporaires en cas d'erreur
        if "temp" in audio_path and os.path.exists(audio_path):
            try:
                os.remove(audio_path)
            except:
                pass
        
        return f"La transcription a échoué: {str(e)}"

# Fonction pour résumer le texte
def summarize_text(text, progress=None):
    """Résume le texte arabe en entrée."""
    if not text or not isinstance(text, str) or "transcription a échoué" in text.lower():
        return "Veuillez fournir un texte à résumer (transcription nécessaire d'abord)."
    
    if progress:
        try:
            progress(0.6, "Chargement du modèle de résumé...")
        except Exception as e:
            print(f"Erreur lors de la mise à jour de la progression: {e}")
    
    try:
        from transformers import pipeline
        
        # Libérer la mémoire avant de charger un nouveau modèle
        clear_gpu_memory()
        
        # Utiliser un modèle de résumé pour l'arabe
        model_name = "malmarjeh/t5-arabic-text-summarization"
        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        
        summarizer = pipeline("summarization", model=model_name, device=device)
        
        if progress:
            try:
                progress(0.7, "Génération du résumé...")
            except Exception as e:
                print(f"Erreur lors de la mise à jour de la progression: {e}")
        
        print("Résumé du texte en cours...")
        
        # Vérifier que le texte n'est pas vide
        if not text.strip():
            return "Le texte à résumer est vide."
        
        # Limiter la longueur du texte pour éviter les problèmes de mémoire
        max_input_length = 1024
        if len(text) > max_input_length:
            text = text[:max_input_length]
            print(f"Texte tronqué à {max_input_length} caractères pour éviter les problèmes de mémoire")
        
        # Vérification robuste du résultat
        try:
            summary_result = summarizer(text, max_length=150, min_length=30, do_sample=False)
            
            if summary_result and isinstance(summary_result, list) and len(summary_result) > 0 and 'summary_text' in summary_result[0]:
                summary = summary_result[0]['summary_text']
                print("Résumé réussi")
                return summary
            else:
                print("Format de résultat de résumé inattendu")
                return "Le résumé a échoué en raison d'un format de résultat inattendu."
        except IndexError as idx_err:
            print(f"Erreur d'index pendant le résumé: {idx_err}")
            return "Le résumé a échoué en raison d'un problème avec le texte. Le texte est peut-être trop court."
        except Exception as inner_err:
            print(f"Erreur interne pendant le résumé: {inner_err}")
            return f"Le résumé a rencontré un problème: {str(inner_err)}"
    except Exception as e:
        print(f"Erreur pendant le résumé: {e}")
        return f"Le résumé a échoué: {str(e)}"

# Fonction pour générer des questions
def generate_questions(text, progress=None):
    """Génère des questions basées sur le texte arabe."""
    if not text or not isinstance(text, str) or "transcription a échoué" in text.lower():
        return "Veuillez fournir un texte pour générer des questions."
    
    if progress:
        try:
            progress(0.8, "Génération des questions...")
        except Exception as e:
            print(f"Erreur lors de la mise à jour de la progression: {e}")
    
    try:
        # Version simplifiée avec des questions génériques
        questions = """1. ما هي الفكرة الرئيسية للدرس؟
2. ما هي أهم المفاهيم التي تم شرحها في هذا الدرس؟
3. كيف يمكن تطبيق المعلومات الواردة في هذا الدرس في الحياة العملية؟
4. ما هي العلاقة بين المفاهيم المختلفة التي تم تقديمها؟
5. لو طلب منك تلخيص هذا الدرس في ثلاث نقاط، ماذا ستكتب؟"""
        
        return questions
    except Exception as e:
        print(f"Erreur dans generate_questions: {e}")
        return f"Erreur lors de la génération des questions: {str(e)}"

# Fonction pour générer une carte mentale
def generate_mind_map(text, progress=None):
    """Génère une représentation de carte mentale simple."""
    if not text or not isinstance(text, str) or "transcription a échoué" in text.lower() or "résumé a échoué" in text.lower():
        return "Veuillez fournir un texte pour générer une carte mentale."
    
    if progress:
        try:
            progress(0.9, "Création de la carte mentale...")
        except Exception as e:
            print(f"Erreur lors de la mise à jour de la progression: {e}")
    
    try:
        # Version ultra-simplifiée avec vérification des indices
        lines = text.split('.')
        
        # Création d'une structure Markdown basique
        markdown_map = "# خريطة ذهنية للدرس\n\n"
        
        # Titre principal - avec vérification que lines n'est pas vide
        if lines and len(lines) > 0 and lines[0].strip():
            main_idea = lines[0].strip()[:100] + ("..." if len(lines[0]) > 100 else "")
            markdown_map += f"## الموضوع الرئيسي\n- {main_idea}\n\n"
        else:
            markdown_map += f"## الموضوع الرئيسي\n- (لم يتم العثور على موضوع رئيسي)\n\n"
        
        # Points principaux (limités à 3 pour simplicité)
        markdown_map += "## النقاط الرئيسية\n\n"
        
        # Vérifier s'il y a des lignes supplémentaires
        if len(lines) > 1:
            count = 0
            for i in range(1, min(len(lines), 10)):  # Limiter à 10 lignes maximum pour éviter les problèmes
                if i < len(lines) and lines[i] and lines[i].strip() and len(lines[i].strip()) > 15:
                    point_text = lines[i].strip()[:80] + ("..." if len(lines[i]) > 80 else "")
                    markdown_map += f"### نقطة {count+1}\n- {point_text}\n\n"
                    count += 1
                    if count >= 3:
                        break
        
        # Si aucun point n'a été trouvé, ajouter un message
        if "### نقطة" not in markdown_map:
            markdown_map += "### ملاحظة\n- لم يتم العثور على نقاط رئيسية كافية في النص\n\n"
        
        return markdown_map
    except Exception as e:
        print(f"Erreur dans generate_mind_map: {e}")
        return f"Erreur lors de la création de la carte mentale: {str(e)}"

# Fonction principale de traitement
def process_lesson(file_obj, progress=gr.Progress()):
    """Traite le fichier audio/vidéo avec une approche ultra-légère."""
    if file_obj is None:
        return "Veuillez télécharger un fichier.", "", "", ""
    
    try:
        try:
            progress(0.1, "Préparation du traitement...")
        except Exception as e:
            print(f"Erreur lors de la mise à jour de la progression: {e}")
        
        file_path = file_obj.name
        file_name, file_extension = os.path.splitext(file_path)
        file_extension = file_extension.lower()
        
        audio_path = None
        try:
            progress(0.2, "Préparation du fichier...")
        except Exception as e:
            print(f"Erreur lors de la mise à jour de la progression: {e}")
        
        # Traitement du fichier selon son type
        try:
            if file_extension in ['.mp4', '.mov', '.avi', '.mkv']:
                audio_path = extract_audio(file_path)
                if not audio_path:
                    return "Échec de l'extraction audio.", "", "", ""
            elif file_extension in ['.mp3', '.wav', '.ogg', '.flac', '.m4a']:
                audio_path = file_path
            else:
                return f"Type de fichier non pris en charge: {file_extension}. Veuillez télécharger un fichier audio ou vidéo.", "", "", ""
        except Exception as e:
            print(f"Erreur lors de la préparation du fichier: {e}")
            return f"Erreur lors de la préparation du fichier: {str(e)}", "", "", ""
        
        # 1. Transcription
        try:
            transcription = transcribe_audio(audio_path, progress)
        except Exception as e:
            print(f"Erreur lors de la transcription: {e}")
            transcription = f"La transcription a échoué: {str(e)}"
        
        # Libérer la mémoire après transcription
        clear_gpu_memory()
        
        # 2. Résumé (seulement si la transcription a réussi)
        try:
            if isinstance(transcription, str) and "échoué" not in transcription.lower():
                summary = summarize_text(transcription, progress)
            else:
                summary = "Impossible de générer un résumé car la transcription a échoué."
        except Exception as e:
            print(f"Erreur lors du résumé: {e}")
            summary = f"Le résumé a échoué: {str(e)}"
        
        # Libérer la mémoire après résumé
        clear_gpu_memory()
        
        # 3. Génération de questions (simplifiée)
        try:
            questions = generate_questions(transcription, progress)
        except Exception as e:
            print(f"Erreur lors de la génération des questions: {e}")
            questions = f"La génération des questions a échoué: {str(e)}"
        
        # 4. Génération de la carte mentale (simplifiée)
        try:
            mind_map = generate_mind_map(summary, progress)
        except Exception as e:
            print(f"Erreur lors de la génération de la carte mentale: {e}")
            mind_map = f"La génération de la carte mentale a échoué: {str(e)}"
        
        try:
            progress(1.0, "Terminé!")
        except Exception as e:
            print(f"Erreur lors de la mise à jour de la progression finale: {e}")
            
        return transcription, summary, questions, mind_map
    
    except Exception as e:
        print(f"Erreur générale: {e}")
        return f"Une erreur s'est produite: {str(e)}", "", "", ""

# Interface Gradio simplifiée
def create_interface():
    with gr.Blocks(theme=gr.themes.Soft(), title="مساعد الدروس بالذكاء الاصطناعي") as demo:
        gr.Markdown("## مساعد الدروس بالذكاء الاصطناعي")
        gr.Markdown("قم برفع ملف صوتي أو فيديو لدرس باللغة العربية، وسيقوم التطبيق بتحويله إلى نص، تلخيصه، توليد أسئلة، وإنشاء خريطة ذهنية.")
        
        with gr.Row():
            input_file = gr.File(label="رفع ملف الدرس (صوت أو فيديو)", file_types=["audio", "video"])
        
        process_button = gr.Button("معالجة الدرس")
        
        with gr.Tabs():
            with gr.TabItem("النص الكامل"):
                output_transcription = gr.Textbox(label="النص المكتوب للدرس", lines=15, interactive=False)
            with gr.TabItem("الملخص"):
                output_summary = gr.Textbox(label="ملخص الدرس", lines=10, interactive=False)
            with gr.TabItem("أسئلة الفهم"):
                output_questions = gr.Textbox(label="أسئلة مقترحة", lines=10, interactive=False)
            with gr.TabItem("الخريطة الذهنية"):
                output_mindmap = gr.Markdown(label="خريطة ذهنية (بصيغة ماركداون)")
        
        process_button.click(
            fn=process_lesson,
            inputs=input_file,
            outputs=[output_transcription, output_summary, output_questions, output_mindmap]
        )
        
        gr.Markdown("--- Developed with AI ---")
    
    return demo

# Point d'entrée principal pour Hugging Face Spaces
demo = create_interface()

# Cette ligne est nécessaire pour Hugging Face Spaces
if __name__ == "__main__":
    demo.launch()