import gradio as gr
import whisper
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import MT5Tokenizer, AutoModelForSeq2SeqLM
from transformers import AutoTokenizer, AutoModelWithLMHead
import graphviz
import torch

# تحويل الصوت إلى نص
model_whisper = whisper.load_model("base")

# تلخيص النص
tokenizer = MT5Tokenizer.from_pretrained("csebuetnlp/mT5_multilingual_XLSum")
model_summarizer = AutoModelForSeq2SeqLM.from_pretrained("csebuetnlp/mT5_multilingual_XLSum")

# تحميل نموذج AraGPT2 العربي الخفيف


aragpt_tokenizer = AutoTokenizer.from_pretrained("aubmindlab/aragpt2-base")
aragpt_model = AutoModelForCausalLM.from_pretrained("aubmindlab/aragpt2-base")
aragpt_model.eval()


def generate_arabic_explanation(prompt):
    inputs = aragpt_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
    outputs = aragpt_model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.8)
    return aragpt_tokenizer.decode(outputs[0], skip_special_tokens=True)
    
def transcribe_summarize_answer(audio_path):
    result = model_whisper.transcribe(audio_path)
    text = result["text"]

    inputs = tokenizer("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)
    summary_ids = model_summarizer.generate(inputs["input_ids"], max_length=150, min_length=40, length_penalty=2.0, num_beams=4)
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

    dot = graphviz.Digraph(comment="Mind Map", format='png')
    dot.node('central', '🧠 فكرة الدرس')
    for i, sentence in enumerate(summary.split(".")):
        if sentence.strip():
            node_id = f"n{i}"
            dot.node(node_id, sentence.strip())
            dot.edge('central', node_id)
    mindmap_path = "/tmp/mindmap"
    dot.render(mindmap_path, cleanup=True)

    prompt = summary + "\n\nاشرح هذه الأفكار بشكل مبسط:"
    answer = generate_arabic_explanation(prompt)

    return text, summary, mindmap_path + ".png", answer

gr.Interface(
    fn=transcribe_summarize_answer,
    inputs=gr.Audio(source="upload", type="filepath", label="🎙️ ارفع ملف صوتي للدرس"),
    outputs=[
        gr.Textbox(label="📜 النص الكامل"),
        gr.Textbox(label="✂️ الملخص"),
        gr.Image(label="🧠 الخريطة الذهنية المرئية"),
        gr.Textbox(label="🤖 شرح AraGPT2 باللغة العربية")
    ],
    title="SmartLessonMap + AraGPT2 🤖🧠",
    description="تطبيق ذكي يحول الدروس الصوتية إلى نص، ملخص، خريطة ذهنية، وشرح باللغة العربية باستخدام نموذج خفيف."
).launch()