import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ============================================================
# Load Model and Tokenizer
# ============================================================

MODEL_ID = "azeddinShr/LFM2-1.2B-RAG-ARABIC-AdaLoRA"

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    force_download=True
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
print("✅ Model loaded successfully!")

# ============================================================
# Generation Function
# ============================================================

def generate_answer(context, question, max_tokens=150, temperature=0.0):
    """
    Generate answer from context and question
    
    Args:
        context: Background information (Arabic text)
        question: Question to answer (Arabic text)
        max_tokens: Maximum length of generated answer
        temperature: Sampling temperature (0.0 = deterministic)
    
    Returns:
        Generated answer
    """
    # Validate inputs
    if not context or not context.strip():
        return "⚠️ الرجاء إدخال السياق (Please provide context)"
    
    if not question or not question.strip():
        return "⚠️ الرجاء إدخال السؤال (Please provide a question)"
    
    try:
        # Format prompt
        prompt = f"استخدم السياق التالي للإجابة على السؤال:\n\n{context}\n\nالسؤال: {question}"
        
        # Tokenize
        messages = [{"role": "user", "content": prompt}]
        input_ids = tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(model.device)
        
        # Generate
        with torch.no_grad():
            outputs = model.generate(
                input_ids,
                max_new_tokens=int(max_tokens),
                temperature=float(temperature) if temperature > 0 else 0.0,
                do_sample=temperature > 0,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        
        # Decode
        answer = tokenizer.decode(
            outputs[0][input_ids.shape[1]:], 
            skip_special_tokens=True
        ).strip()
        
        return answer
    
    except Exception as e:
        return f"❌ Error: {str(e)}"

# ============================================================
# Example Questions
# ============================================================

long_example = """
المستندات: ['العنوان: ملخص مركز السيطرة على الأمراض والوقاية منها 21 مارس 2020،\nالمقطع: يصيب الفيروس المسبب لمرض كوفيد-19 الناس وينتشر بسهولة من شخص لآخر. وقد تم اكتشاف حالات في معظم بلدان العالم ويتم اكتشاف انتشار مجتمعي في عدد متزايد من البلدان. في 11 مارس، وصفت منظمة الصحة العالمية تفشي مرض كوفيد-19 بأنه جائحة.'، 'العنوان: 2019-nCoV: أداة التعرف والعزل والإبلاغ (3I) المطبقة على فيروس كورونا المستجد الناشئ\nالمقطع: من غير الواضح حاليًا كيف ينتشر 2019-nCoV، ولكن يُشتبه في أنه ينتقل من خلال ملامسة إفرازات الجهاز التنفسي المصابة، مثل فيروسات كورونا الأخرى المعروفة. هناك حالات انتقال مستدام من إنسان إلى إنسان عبر أجيال من الحالات، وخاصة بالقرب من مركز الزلزال في مدينة ووهان. 21 تشير الأدلة الحالية إلى أن الاتصال الوثيق بشخص مصاب هو عامل رئيسي في انتقال المرض. يعرّف مركز السيطرة على الأمراض "الاتصال الوثيق" 33 بأنه التواجد في منطقة بها مريض مؤكد أو على بعد مترين منها أو التعرض المباشر لإفرازات معدية دون معدات الوقاية الشخصية المناسبة. أفادت المرافق الصحية في الصين بانتشار المرض من شخص إلى آخر، "العنوان: أولويات البحث الوبائي للسيطرة على الصحة العامة لتفشي فيروس كورونا المستجد العالمي المستمر (2019-nCoV)\nالمقطع: يحدث انتقال الفيروسات التنفسية عمومًا من خلال قطرات تنفسية كبيرة، ولكن يمكن أن تنتشر بعض الفيروسات التنفسية من خلال الهباء الجوي ذي الجسيمات الدقيقة، ويمكن أن يلعب الانتقال غير المباشر عبر الأدوات دورًا أيضًا. يمكن أن تصيب فيروسات كورونا أيضًا الجهاز الهضمي البشري، وقد يلعب الانتقال البرازي الفموي دورًا أيضًا في هذه الحالة. "إن انتشار فيروس كورونا المستجد في حدائق أموي حيث أصيب أكثر من 300 حالة كان بسبب انتشاره عن طريق البراز والفم، ثم عن طريق الهواء، من خلال فروق الضغط بين أنابيب الصرف الصحي الملوثة، ومصارف أرضية الحمامات والمراحيض. ويبدو أن أول حدث انتشار فائق يمكن التعرف عليه خلال تفشي فيروس كورونا المستجد الحالي قد حدث"، "العنوان: ملخص مركز السيطرة على الأمراض 21 مارس 2020، المقطع: يجب على الأشخاص الذين يعانون من الحمى أو السعال أن يفكروا فيما إذا كانوا مصابين بفيروس كورونا المستجد، اعتمادًا على مكان إقامتهم أو تاريخ سفرهم أو التعرضات الأخرى. يشهد أكثر من نصف الولايات المتحدة مستوى معينًا من انتشار فيروس كورونا المستجد في المجتمع. يمكن الوصول إلى اختبار فيروس كورونا المستجد من خلال مقدمي الخدمات الطبية أو إدارات الصحة العامة، ولكن لا يوجد علاج لهذا الفيروس. يعاني معظم الأشخاص من مرض خفيف ويمكنهم التعافي في المنزل دون رعاية طبية. "]
"""
examples = [
    [
        "أطلقت شركة أوبن إيه آي رسمياً نموذجها الجديد GPT-5 في 7 أغسطس 2025. يمثل GPT-5 قفزة كبيرة في قدرات الذكاء الاصطناعي، حيث يجمع بين التفكير المتقدم والقدرات متعددة الوسائط في نظام واحد موحد. حقق النموذج نتائج قياسية في المعايير الأكاديمية، بما في ذلك 94.6% في اختبار AIME 2025 للرياضيات المتقدمة و74.9% في SWE-bench Verified لمهام البرمجة الواقعية. كما أن GPT-5 أقل عرضة للهلوسة بنسبة 45% مقارنة بـ GPT-4o.",
        "متى تم إطلاق نموذج GPT-5 رسمياً؟"
    ],
    [
        long_example,
        "كيف ينتشر فيروس كورونا المستجد؟"
    ],
    [
        "نيوم هو مشروع ضخم في شمال غرب المملكة العربية السعودية أطلقه ولي العهد الأمير محمد بن سلمان في أكتوبر 2017. تبلغ المساحة الإجمالية المخططة لنيوم 26,500 كيلومتر مربع بتكلفة تقدر بـ 500 مليار دولار. يتضمن المشروع مناطق متعددة منها ذا لاين (The Line)، وهي مدينة خطية مستقبلية، بالإضافة إلى سندالة وأوكساجون وتروجينا. من المخطط أن تكتمل أجزاء رئيسية من المشروع بحلول عام 2030 كجزء من رؤية السعودية 2030.",
        "ما هي المساحة الإجمالية المخططة لمشروع نيوم؟"
    ],
    [
        "في فبراير 2025، وقعت نيوم اتفاقية شراكة مع شركة داتافولت السعودية لتصميم مصنع ضخم للذكاء الاصطناعي سيكون بمثابة مركز بيانات. بموجب الاتفاقية، سيتم تنفيذ المشروع في أوكساجون، مدينة الصناعات النظيفة والمتقدمة التابعة لنيوم، على مراحل. ستشهد المرحلة الأولى استثماراً أولياً قدره 5 مليارات دولار، ومن المتوقع أن يكون جاهزاً للعمل بحلول عام 2028.",
        "كم تبلغ قيمة الاستثمار الأولي لمصنع الذكاء الاصطناعي في نيوم؟"
    ]
]

# ============================================================
# Gradio Interface
# ============================================================

with gr.Blocks(theme=gr.themes.Soft(), title="Arabic QA System") as demo:
    
    gr.Markdown(
        """
        # 🤖 نظام الأسئلة والأجوبة العربي
        # Arabic Question Answering System
        
        Fine-tuned LFM2-1.2B-RAG model for Arabic extractive question answering.
        Provide context in Arabic and ask a question to get an answer extracted from the context.
        
        **Model:** [azeddinShr/LFM2-1.2B-RAG-ARABIC-AdaLoRA](https://huggingface.co/azeddinShr/LFM2-1.2B-RAG-ARABIC-AdaLoRA)
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            context_input = gr.Textbox(
                label="السياق / Context",
                placeholder="أدخل النص الذي يحتوي على المعلومات...\nEnter the text containing the information...",
                lines=8,
                rtl=True
            )
            
            question_input = gr.Textbox(
                label="السؤال / Question",
                placeholder="اطرح سؤالك هنا...\nAsk your question here...",
                lines=2,
                rtl=True
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                max_tokens = gr.Slider(
                    minimum=50,
                    maximum=500,
                    value=150,
                    step=10,
                    label="Max Tokens",
                    info="Maximum length of generated answer"
                )
                
                temperature = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.0,
                    step=0.1,
                    label="Temperature",
                    info="0 = deterministic, higher = more creative"
                )
            
            submit_btn = gr.Button("🔍 Get Answer / احصل على الإجابة", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            answer_output = gr.Textbox(
                label="الإجابة / Answer",
                lines=10,
                rtl=True,
                show_copy_button=True
            )
    
    gr.Examples(
        examples=examples,
        inputs=[context_input, question_input],
        label="📝 Example Questions / أمثلة"
    )
    
    gr.Markdown(
        """
        ---
        ### ℹ️ How to use / كيفية الاستخدام
        
        1. **Paste Context**: Add Arabic text containing information
        2. **Ask Question**: Write your question about the context
        3. **Get Answer**: The model will extract the answer from the context
        
        **Note:** This model is optimized for Modern Standard Arabic and extractive QA tasks.
        
        ### 🔗 Links
         - [Base Model](https://huggingface.co/LiquidAI/LFM2-1.2B-RAG)
        -  [Dataset (ARCD)](https://huggingface.co/datasets/hsseinmz/arcd)
        """
    )
    
    # Connect button to function
    submit_btn.click(
        fn=generate_answer,
        inputs=[context_input, question_input, max_tokens, temperature],
        outputs=answer_output
    )

# ============================================================
# Launch
# ============================================================

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