import torch
import re
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from peft import PeftModel
from threading import Thread
from transformers import TextIteratorStreamer

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"⏳ جاري التحضير باستخدام: {device}")

# 1. دالة التنظيف
def selective_normalization(text):
    if not isinstance(text, str):
        return text
    text = re.sub(r'[\u0617-\u061A\u064B-\u0652]', '', text)
    text = re.sub(r'ـ', '', text)
    text = re.sub(r'[أإآٱ]', 'ا', text)
    text = re.sub(r'ؤ', 'و', text)
    text = re.sub(r'ئ', 'ي', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

# 2. تحميل الموديل الأساسي مرة واحدة
print("⏳ جاري تحميل الموديل الأساسي (AraT5)...")
base_model_name = "UBC-NLP/AraT5-base"
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model_name)

# مسارات المحولات من Hugging Face
path_amman_to_msa = "FarisAlnatsheh/arat5-amman-to-msa-lora" 
path_msa_to_amman = "FarisAlnatsheh/arat5-msa-to-amman-lora"

print("⏳ جاري تركيب المحولات (Adapters) للنموذجين...")
try:
    # تحميل أوزان الاتجاه الأول 
    model = PeftModel.from_pretrained(base_model, path_amman_to_msa, adapter_name="amman2msa")
    
    # تحميل أوزان الاتجاه الثاني وتركيبها
    model.load_adapter(path_msa_to_amman, adapter_name="msa2amman")
    
    model = model.to(device)
    model.eval()
    print("✅ تم تحميل النموذجين بنجاح ودمجهما!")
except Exception as e:
    print(f"❌ حدث خطأ في تحميل الأوزان، تأكد من مسارات الملفات: {e}")

# 3. دالة الترجمة التفاعلية
def translate_text(text, direction):
    clean_text = selective_normalization(text)
    
    if direction == "من لهجة عمّان ⬅️ إلى فصحى":
        input_text = "Amman to MSA: " + clean_text
        model.set_adapter("amman2msa")
    else:
        input_text = "MSA to Amman: " + clean_text
        model.set_adapter("msa2amman")
        
    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        truncation=True,
        max_length=64
    ).to(device)
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    generation_kwargs = dict(
        **inputs,
        max_new_tokens=64,
        num_beams=1,            
        do_sample=True,          
        temperature=0.2,         
        top_k=30,                
        top_p=0.70,              
        repetition_penalty=1.15,  
        streamer=streamer        
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    translated_text = ""
    for new_text in streamer:
        translated_text += new_text
        yield translated_text

# 4. واجهة المستخدم (Gradio)
custom_css = """
#rtl-radio {
    direction: rtl;
    text-align: right;
}
.rtl-text {
    direction: rtl !important;
}
.rtl-textbox {
    direction: rtl !important;
}
.rtl-textbox > label > span {
    right: 45px !important;
    left: auto !important;
}
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css=custom_css) as demo:
    gr.Markdown("<h1 style='text-align: center;'>محول الجمل بين اللهجات العربية</h1>")
    
    with gr.Row():
        direction = gr.Radio(
            choices=["من فصحى ⬅️ إلى لهجة عمّان", "من لهجة عمّان ⬅️ إلى فصحى"],
            value="من فصحى ⬅️ إلى لهجة عمّان",
            label="اختر اتجاه الترجمة",
            elem_id="rtl-radio" 
        )
    
    with gr.Row():
        output_text = gr.Textbox(
            label="الترجمة", 
            lines=5, 
            interactive=False, 
            rtl=True,            
            text_align="right",  
            elem_classes="rtl-textbox"
        )

        input_text = gr.Textbox(
            label="النص الأصلي", 
            placeholder="اكتب هنا أو اختر مثالاً من الأسفل...", 
            lines=5, 
            rtl=True,            
            text_align="right",  
            elem_classes="rtl-textbox"
        )
    
    translate_btn = gr.Button("ترجم الآن 🚀", variant="primary")
    
    gr.Markdown("<h3 dir='rtl' style='text-align: right;'>💡 جمل مقترحة للتجربة :</h3>")
    
    with gr.Column(visible=True) as msa_examples_col:
        with gr.Row():
            btn1 = gr.Button("كيف حالك يا صديقي ؟ يا رجل كان الامتحان صعبا جدا!", size="sm", elem_classes="rtl-text")
            btn2 = gr.Button("الجو في عمان اليوم رائع جداً ومناسب للتنزه.", size="sm", elem_classes="rtl-text")
            btn3 = gr.Button("أريد حجز طاولة في المطعم الكبير في وسط البلد.", size="sm", elem_classes="rtl-text")
            btn7 = gr.Button("متى ستبدأ المحاضرة الأولى اليوم؟", size="sm", elem_classes="rtl-text")

        with gr.Row():
            btn8 = gr.Button("لم أفهم أي شيء من شرح الأستاذ اليوم.", size="sm", elem_classes="rtl-text")
            btn9 = gr.Button("لماذا الإنترنت بطيء جداً اليوم؟", size="sm", elem_classes="rtl-text")
            btn10 = gr.Button("قم بالركوب في هذه السيارة.", size="sm", elem_classes="rtl-text")
            btn11 = gr.Button("أطلب منك أن تسرع في الصعود إلى السيارة الآن وإلا سوف نتأخر عن موعد العشاء.", size="sm", elem_classes="rtl-text")
            
        btn1.click(lambda: "كيف حالك يا صديقي ؟ يا رجل كان الامتحان صعبا جدا!", None, input_text)
        btn2.click(lambda: "الجو في عمان اليوم رائع جداً ومناسب للتنزه.", None, input_text)
        btn3.click(lambda: "أريد حجز طاولة في المطعم الكبير في وسط البلد.", None, input_text)
        btn7.click(lambda: "متى ستبدأ المحاضرة الأولى اليوم؟", None, input_text)
        btn8.click(lambda: "لم أفهم أي شيء من شرح الأستاذ اليوم.", None, input_text)
        btn9.click(lambda: "لماذا الإنترنت بطيء جداً اليوم؟", None, input_text)
        btn10.click(lambda: "قم بالركوب في هذه السيارة.", None, input_text)
        btn11.click(lambda: "أطلب منك أن تسرع في الصعود إلى السيارة الآن وإلا سوف نتأخر عن موعد العشاء.", None, input_text)

    with gr.Column(visible=False) as amm_examples_col:
        with gr.Row():
            btn4 = gr.Button("لو سمحت, وين بقدر ألاقي محمد؟ بدي اياه باشي مهم.", size="sm", elem_classes="rtl-text")
            btn5 = gr.Button("قديش حق هالقطعة؟ لازم اشتريها علشان شكلها حلو.", size="sm", elem_classes="rtl-text")
            btn6 = gr.Button("يزم، هاد الامتحان كان كثير صعب!", size="sm", elem_classes="rtl-text")
            btn12 = gr.Button("اركب بهالسيارة بسرعة.", size="sm", elem_classes="rtl-text")

        with gr.Row():
            btn13 = gr.Button("ما فهمتش ولا اشي من شرح الأستاذ اليوم.", size="sm", elem_classes="rtl-text")
            btn14 = gr.Button("الجو اليوم حلو شو رايك بعد ما نخلص اكل نروح نشتري اواعي العيد من وسط البلد.", size="sm", elem_classes="rtl-text")
            btn15 = gr.Button("كيفك يا زلمة؟ شو اخبارك؟ من زمان ما شفناك.", size="sm", elem_classes="rtl-text")
            
        btn4.click(lambda: "لو سمحت, وين بقدر ألاقي محمد؟ بدي اياه باشي مهم.", None, input_text)
        btn5.click(lambda: "قديش حق هالقطعة؟ لازم اشتريها علشان شكلها حلو.", None, input_text)
        btn6.click(lambda: "يزم، هاد الامتحان كان كثير صعب!", None, input_text)
        btn12.click(lambda: "اركب بهالسيارة بسرعة.", None, input_text)
        btn13.click(lambda: "ما فهمتش ولا اشي من شرح الأستاذ اليوم.", None, input_text)
        btn14.click(lambda: "الجو اليوم حلو شو رايك بعد ما نخلص اكل نروح نشتري اواعي العيد من وسط البلد.", None, input_text)
        btn15.click(lambda: "كيفك يا زلمة؟ شو اخبارك؟ من زمان ما شفناك.", None, input_text)

    def toggle_ui(dir_val):
        if dir_val == "من فصحى ⬅️ إلى لهجة عمّان":
            return gr.update(visible=True), gr.update(visible=False), "", ""
        else:
            return gr.update(visible=False), gr.update(visible=True), "", ""
    
    direction.change(
        fn=toggle_ui, 
        inputs=direction, 
        outputs=[msa_examples_col, amm_examples_col, input_text, output_text]
    )
        
    translate_btn.click(fn=translate_text, inputs=[input_text, direction], outputs=output_text)

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