import gradio as gr
import torch
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
from moviepy.video.io.VideoFileClip import VideoFileClip
import librosa
import soundfile as sf
import numpy as np
import re
from gtts import gTTS
from langdetect import detect
import os

# ======================
# 1. الإعدادات الأساسية
# ======================
device = "cuda" if torch.cuda.is_available() else "cpu"

# ======================
# 2. تحميل النماذج
# ======================
pipe = pipeline("automatic-speech-recognition", model="openai/whisper-medium", device=0 if device == "cuda" else -1)
bart_model = AutoModelForSeq2SeqLM.from_pretrained("ahmedabdo/arabic-summarizer-bart")
bart_tokenizer = AutoTokenizer.from_pretrained("ahmedabdo/arabic-summarizer-bart")

# تحميل نموذج الإجابة بالعربية
qa_model_name = "ZeyadAhmed/AraElectra-Arabic-SQuADv2-QA"
qa_pipeline = pipeline("question-answering", model=qa_model_name, tokenizer=qa_model_name, device=0 if device == "cuda" else -1)

# ======================
# 3. التخصيصات
# ======================
PRIMARY_COLOR = "#2A4D6E"  # اللون الكحلي
EXAMPLE_AUDIO_PATH = "_⁨ما_هو_الذكاء_الإصطناعي_؟_فِهموجرافيك_١⁩.m4a"  

# ======================
# 4. الدوال المساعدة
# ======================
def clean_text_advanced(text):
    """تنظيف النص مع الحفاظ على الهيكل الأساسي"""
    text = re.sub(r'\b\d+[.)]\s*', '', text)  # إزالة الترقيم الرقمي
    text = re.sub(r'[ـ\-_*#@&]', ' ', text)    # إزالة الرموز الخاصة
    text = re.sub(r'\s+', ' ', text).strip()   # تقليل الفراغات
    return text

def prepare_context(text, max_length=1024):
    """تحضير السياق مع الحفاظ على الجمل الكاملة"""
    cleaned = clean_text_advanced(text)
    sentences = [s.strip() for s in cleaned.split('.') if s.strip()]
    selected = []
    total_len = 0

    for sentence in reversed(sentences):
        if total_len + len(sentence.split()) > max_length:
            break
        selected.append(sentence)
        total_len += len(sentence.split())

    return '. '.join(reversed(selected))

def convert_audio_to_text(uploaded_file):
    try:
        if not uploaded_file:
            return "⛔ الرجاء رفع ملف أولاً"

        input_path = uploaded_file if isinstance(uploaded_file, str) else uploaded_file.name
        output_path = "/tmp/processed.wav"

        # معالجة ملفات الفيديو
        if input_path.split('.')[-1].lower() in ['mp4', 'avi', 'mov', 'mkv']:
            video = VideoFileClip(input_path)
            if video.audio:
                video.audio.write_audiofile(output_path, codec='pcm_s16le')
            else:
                return "⛔ لا يوجد صوت في الفيديو!"
        else:
            output_path = input_path

        audio, rate = librosa.load(output_path, sr=16000)
        transcripts = []
        # تقسيم الصوت إلى مقاطع للتعامل مع الملفات الكبيرة
        for start in np.arange(0, len(audio)/rate, 30):
            end = min(start + 30, len(audio)/rate)
            segment = audio[int(start*rate):int(end*rate)]
            sf.write(f"/tmp/segment_{int(start)}.wav", segment, rate)
            transcripts.append(pipe(f"/tmp/segment_{int(start)}.wav")["text"])
        return " ".join(transcripts)
    except Exception as e:
        return f"⛔ خطأ: {str(e)}"
        
def process_example_audio():
    try:
        if not os.path.exists(EXAMPLE_AUDIO_PATH):
            return "⛔ الملف التجريبي غير موجود"
        return convert_audio_to_text(EXAMPLE_AUDIO_PATH)
    except Exception as e:
        return f"⛔ خطأ في معالجة المثال: {str(e)}"

def summarize_text(text):
    cleaned_text = clean_text_advanced(text)
    inputs = bart_tokenizer(cleaned_text, return_tensors="pt", max_length=1024, truncation=True).to(device)
    summary_ids = bart_model.generate(
        inputs.input_ids,
        max_length=150,
        num_beams=4,
        early_stopping=True
    )
    return bart_tokenizer.decode(summary_ids[0], skip_special_tokens=True)

def answer_question(text, question):
    if not question.strip() or not text.strip():
        return "⛔️ الرجاء إدخال النص والسؤال بشكل صحيح"

    # تقسيم النص إلى شرائح صغيرة بحيث لا تزيد كل شريحة عن 256 كلمة
    words = text.split()
    chunk_size = 256
    segments = [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]

    best_answer = None
    best_score = -1

    # تطبيق نموذج الإجابة على كل شريحة واختيار الإجابة ذات أعلى score
    for seg in segments:
        result = qa_pipeline({'question': question, 'context': seg})
        if result['score'] > best_score:
            best_score = result['score']
            best_answer = result['answer']

    return best_answer


def text_to_speech(text):
    if not text.strip():
        return None
    try:
        tts = gTTS(text=text, lang='ar' if detect(text) == 'ar' else 'en', slow=False)
        output = "/tmp/tts.wav"
        tts.save(output)
        return output
    except Exception as e:
        return f"⛔ خطأ في توليد الصوت: {str(e)}"

# ======================
# 5. واجهة المستخدم (التصميم المحسّن)
# ======================
with gr.Blocks(css=f"""
.gradio-container {{
    background: #f8f9fa;
    font-family: 'Noto Sans Arabic', sans-serif;
}}

.header-section {{
    background: {PRIMARY_COLOR};
    padding: 2rem;
    border-radius: 10px;
    color: white;
    text-align: center;
    margin-bottom: 2rem;
}}

.tab-buttons {{
    border-bottom: 3px solid {PRIMARY_COLOR} !important;
}}

.tab-buttons button {{
    color: {PRIMARY_COLOR} !important;
    background: #f8f9fa !important;
    transition: all 0.3s !important;
}}

.tab-buttons button.selected {{
    border-top: 3px solid {PRIMARY_COLOR} !important;
    background: rgba(42, 77, 110, 0.1) !important;
}}

.custom-button {{
    background: {PRIMARY_COLOR} !important;
    color: white !important;
    border-radius: 8px;
    padding: 12px 25px;
    margin: 5px;
}}

.custom-button:hover {{
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(42, 77, 110, 0.3);
}}

audio {{
    width: 100%;
    margin: 1rem 0;
}}
""") as demo:

    gr.Markdown("""
    <div class="header-section">
        <h1>فهـــيم 🧠</h1>
        <p>منصة الذكاء الاصطناعي المتكاملة لفهم المحتوى المرئي والسمعي</p>
    </div>
    """)

    with gr.Tabs():
        # تبويب استخراج النص
        with gr.TabItem("🎤 استخراج النص"):
            gr.Markdown("### <span style='color:#2A4D6E'>🔊 مثال تجريبي</span>")
            gr.Audio(EXAMPLE_AUDIO_PATH)
            example_btn = gr.Button("تجربة المثال ⚡", elem_classes="custom-button")
            gr.Markdown("<div style='color:#2A4D6E; margin-top:10px'>اضغط على الزر أعلاه أو ارفع ملفك الخاص</div>")

            file_input = gr.File(file_types=[".wav", ".mp3", ".mp4"])
            gr.Markdown("<span style='color:#2A4D6E'>🖱️ يدعم الصوت والفيديو (MP3, WAV, MP4)</span>")

            with gr.Row():
                extract_btn = gr.Button("بدء التحليل 🚀", elem_classes="custom-button")
                clear_btn = gr.Button("مسح الكل 🗑️", elem_classes="custom-button")

            extracted_text = gr.Textbox(label="📝 النص المستخرج", lines=8)

        # تبويب التلخيص
        with gr.TabItem("📃 التلخيص"):
            with gr.Row():
                summarize_btn = gr.Button("توليد الملخص ✨", elem_classes="custom-button")
                tts_btn = gr.Button("تحويل للصوت 🔊", elem_classes="custom-button")
            summary_output = gr.Textbox(label="📌 الملخص المكثف", lines=6)
            audio_output = gr.Audio()

        # تبويب الأسئلة
        with gr.TabItem("❓ الأسئلة"):
            gr.Markdown("### <span style='color:#2A4D6E'>💡 الأسئلة الذكية</span>")
            with gr.Row():
                q1_btn = gr.Button("الموضوع الرئيسي", elem_classes="custom-button")
                q2_btn = gr.Button("الأفكار المفتاحية", elem_classes="custom-button")
                q3_btn = gr.Button("الاستنتاجات", elem_classes="custom-button")

            question_input = gr.Textbox(placeholder="...اكتب سؤالك هنا")

            with gr.Row():
                answer_btn = gr.Button("الحصول على إجابة 🎯", elem_classes="custom-button")
                clear_answer_btn = gr.Button("مسح الإجابة 🗑️", elem_classes="custom-button")

            answer_output = gr.Textbox(label="💡 الإجابة الدقيقة", lines=4)

    # ======================
    # توصيل الأحداث
    # ======================
    extract_btn.click(convert_audio_to_text, inputs=file_input, outputs=extracted_text)
    example_btn.click(process_example_audio, outputs=extracted_text)
    clear_btn.click(lambda: [None, "", ""], outputs=[file_input, extracted_text, summary_output])
    summarize_btn.click(summarize_text, inputs=extracted_text, outputs=summary_output)
    tts_btn.click(text_to_speech, inputs=summary_output, outputs=audio_output)
    q1_btn.click(lambda: "ما هو الموضوع الرئيسي للنص؟", outputs=question_input)
    q2_btn.click(lambda: "ما هي أهم النقاط في النص؟", outputs=question_input)
    q3_btn.click(lambda: "ما هي الاستنتاجات الرئيسية؟", outputs=question_input)
    answer_btn.click(answer_question, inputs=[extracted_text, question_input], outputs=answer_output)
    clear_answer_btn.click(lambda: "", outputs=answer_output)

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