import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from gtts import gTTS
import os
import torch
import spaces # المكتبة الخاصة بـ Hugging Face Spaces لإدارة الـ GPU

print("Loading model and tokenizer...")
model_name = "sshleifer/distilbart-cnn-12-6"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# استخدام دالة الجلب المخصصة لبيئة ZeroGPU وتحديد مسار برمجيات torch
if torch.cuda.is_available():
    model = model.to("cuda")

# إضافة الـ Decorator لكي يتعرف نظام تشغيل الموقع على الدالة ويسمح ببنائها
@spaces.GPU
def process_text(text):
    if not text.strip():
        return "الرجاء إدخال نص صالح للتلخيص.", None
    
    try:
        # نقل المدخلات إلى الـ GPU إذا كان متاحاً
        device = "cuda" if torch.cuda.is_available() else "cpu"
        inputs = tokenizer([text], max_length=1024, return_tensors="pt", truncation=True).to(device)
        
        # توليد التلخيص
        summary_ids = model.generate(
            inputs["input_ids"], 
            num_beams=2, 
            max_length=140, 
            min_length=30, 
            early_stopping=True
        )
        summary_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        
        # توليد الملف الصوتي
        tts = gTTS(text=summary_text, lang='en')
        audio_path = "summary_audio.mp3"
        tts.save(audio_path)
        
        return summary_text, audio_path
    except Exception as e:
        return f"حدث خطأ أثناء معالجة النص: {str(e)}", None

# بناء واجهة Gradio
theme = gr.themes.Soft(primary_hue="blue", secondary_hue="indigo")

with gr.Blocks(theme=theme) as demo:
    gr.Markdown("# 🤖 AI Text Summarizer & Audio Generator")
    gr.Markdown("قم بلصق أي نص إنجليزي طويل للحصول على تلخيص ذكي وملف صوتي للملخص تلقائياً.")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Input Text (النص الأصلي)", lines=10, placeholder="Enter or paste your long text here...")
            submit_btn = gr.Button("Generate Summary & Audio", variant="primary")
        
        with gr.Column():
            output_summary = gr.Textbox(label="AI Summary (التلخيص الذكي)", lines=5)
            output_audio = gr.Audio(label="Audio Playback (الملف الصوتي)", type="filepath")
            
    submit_btn.click(fn=process_text, inputs=input_text, outputs=[output_summary, output_audio])

if __name__ == "__main__":
    demo.launch()