import os
import re
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, MarianMTModel, MarianTokenizer
from openai import OpenAI

# ==========================================
# النماذج المتاحة (مساران):
#   1) Qwen3-1.7B → عبر Transformers، يدعم وضع التفكير (thinking) القابل للتبديل.
#   2) GPT-5      → عبر الـ OpenAI API.
# النموذج المحلي يُحمّل كسولاً عند أول استخدام لتوفير الذاكرة عند الإقلاع.
# ==========================================
QWEN_MODEL = "Qwen/Qwen3-1.7B"
GPT_LABEL  = "GPT-5 (OpenAI API)"

_qwen_models = {}   # سجلّ نماذج Qwen المحمّلة عبر Transformers

# ------------------------------------------
# تحميل Qwen عبر Transformers (مع تكميم int8 على الـ CPU لتوفير الذاكرة)
# ------------------------------------------
def get_qwen_model(model_id):
    if model_id not in _qwen_models:
        print(f"⏳ جاري تحميل نموذج {model_id} ...")
        tok = AutoTokenizer.from_pretrained(model_id)
        if torch.cuda.is_available():
            # على الـ GPU: تحميل عادي بالدقة التلقائية
            mdl = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
        else:
            # على الـ CPU: نحمّل float32 ثم نكمّم الطبقات الخطية إلى int8 (تكميم ديناميكي)
            print("🔻 تطبيق تكميم int8 الديناميكي على CPU لتوفير الذاكرة والتسريع...")
            mdl = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True)
            mdl = mdl.to(torch.float32)   # يحوّل أي أوزان/buffers بقيت bfloat16 إلى float32 (يصلح خطأ التكميم)
            mdl = torch.quantization.quantize_dynamic(mdl, {torch.nn.Linear}, dtype=torch.qint8)
        _qwen_models[model_id] = (tok, mdl)
    return _qwen_models[model_id]

# ==========================================
# وحدة كشف اللغة: تحدد متى نشغّل الجسر العربي تلقائياً
# ==========================================
_ARABIC_CHAR = re.compile(r"[؀-ۿݐ-ݿࢠ-ࣿ]")

def is_arabic(text):
    """يرجّع True إذا كانت أغلب أحرف النص عربية — هكذا يعرف البوت أي مسار يستخدم."""
    letters = [c for c in text if c.isalpha()]
    if not letters:
        return False
    arabic = [c for c in letters if _ARABIC_CHAR.match(c)]
    return (len(arabic) / len(letters)) > 0.4

# ==========================================
# توليد Qwen عبر Transformers مع دعم وضع التفكير القابل للتبديل
# Qwen3 يوصي: وضع التفكير temp=0.6/top_p=0.95 — الوضع العادي temp=0.7/top_p=0.8 — وtop_k=20 دائماً
# ==========================================
def qwen_generate(messages, enable_thinking, tok, mdl, max_new_tokens=1024):
    temperature = 0.6 if enable_thinking else 0.7
    top_p = 0.95 if enable_thinking else 0.8
    device = next(mdl.parameters()).device
    text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking)
    model_inputs = tok([text], return_tensors="pt").to(device)
    generated_ids = mdl.generate(
        **model_inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
        top_k=20,
    )
    return generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

# ==========================================
# جسر الترجمة عبر نماذج opus-mt المحلية (Helsinki-NLP)
# نماذج مخصّصة للترجمة، تعمل بدون إنترنت، وأدق بكثير من ترجمة النموذج لنفسه.
# تُحمّل كسولاً عند أول استخدام حتى لا تُبطئ بدء التشغيل.
# ==========================================
AR_EN_MODEL = "Helsinki-NLP/opus-mt-ar-en"
EN_AR_MODEL = "Helsinki-NLP/opus-mt-en-ar"
_translators = {}

def _get_translator(model_id):
    if model_id not in _translators:
        print(f"⏳ جاري تحميل نموذج الترجمة {model_id} ...")
        tok = MarianTokenizer.from_pretrained(model_id)
        mdl = MarianMTModel.from_pretrained(model_id)
        _translators[model_id] = (tok, mdl)
    return _translators[model_id]

def opus_translate(text, model_id):
    """يترجم النص جملةً جملة حتى لا يتجاوز حد النموذج (512) وتبقى الدقة عالية."""
    tok, mdl = _get_translator(model_id)
    chunks = [c for c in re.split(r"(?<=[.!؟?\n])\s+", text.strip()) if c.strip()]
    if not chunks:
        return ""
    batch = tok(chunks, return_tensors="pt", padding=True, truncation=True, max_length=512)
    generated = mdl.generate(**batch, max_new_tokens=512)
    return " ".join(tok.decode(g, skip_special_tokens=True).strip() for g in generated)

def ar_to_en(text):
    return opus_translate(text, AR_EN_MODEL)

def en_to_ar(text):
    return opus_translate(text, EN_AR_MODEL)

# ==========================================
# تنظيف رسائل المساعد القديمة من الشارات/كواليس التفكير قبل تمريرها للنموذج
# ==========================================
def _strip_assistant_extras(content):
    if not isinstance(content, str):
        return ""
    if "</details>" in content:                                  # إزالة صندوق التفكير (Qwen)
        content = content.split("</details>")[-1].strip()
    if content.startswith(("🌉", "🟢")) and "\n\n" in content:    # إزالة شارة المسار
        content = content.split("\n\n", 1)[-1].strip()
    return content

# ==========================================
# الدالة الأساسية لمعالجة المحادثة بنظام الـ Messages الحديث
# (غلاف يلتقط أي خطأ ويعرضه نصاً في الشات بدل شارة Gradio العامة "Error")
# ==========================================
def chat_engine(user_input, chat_history, model_choice, enable_thinking, api_key, system_prompt, enable_bridge):
    try:
        return _chat_engine_impl(user_input, chat_history, model_choice, enable_thinking, api_key, system_prompt, enable_bridge)
    except Exception as e:
        import traceback
        traceback.print_exc()
        if chat_history is None:
            chat_history = []
        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": f"❌ خطأ عام في المعالجة:\n\n`{type(e).__name__}: {str(e)}`"})
        return chat_history, ""

def _chat_engine_impl(user_input, chat_history, model_choice, enable_thinking, api_key, system_prompt, enable_bridge):
    if not user_input.strip():
        return chat_history, ""

    # 1. تهيئة مصفوفة الرسائل ببرومبت النظام (System Prompt)
    messages = [{"role": "system", "content": system_prompt}]

    # 2. تفكيك التاريخ القديم بصيغة الـ Dicts الجديدة
    for msg in chat_history:
        role = msg["role"]
        content = msg["content"]
        if role == "assistant":
            content = _strip_assistant_extras(content)
        messages.append({"role": role, "content": content})

    # 3. إضافة رسالة المستخدم الجديدة الحالية
    messages.append({"role": "user", "content": user_input})

    # 4. تقليم التاريخ: نُبقي آخر تبادل واحد فقط (سؤال+جواب) حتى لا يتضخّم السياق
    MAX_HISTORY_MSGS = 2
    if len(messages) > 1 + MAX_HISTORY_MSGS + 1:
        messages = [messages[0]] + messages[-(MAX_HISTORY_MSGS + 1):]

    # ------------------------------------------
    # المسار الأول: نموذج Qwen3 المحلي
    # ------------------------------------------
    if model_choice == QWEN_MODEL:
        try:
            # 🌉 الجسر العربي: يقرر البوت تلقائياً متى يحوّل المسار عبر الإنجليزية
            use_bridge = enable_bridge and is_arabic(user_input)

            if use_bridge:
                # 1) ترجمة سؤال المستخدم من العربية إلى الإنجليزية ليفكّر النموذج بلغته الأقوى
                en_input = ar_to_en(user_input)
                gen_messages = [{"role": "system", "content": "You are a smart, concise and helpful assistant. Answer clearly and accurately in English."}]
                for msg in messages[1:-1]:  # نمرر السياق السابق كما هو (نتجاوز الـ system والرسالة الحالية)
                    gen_messages.append(msg)
                gen_messages.append({"role": "user", "content": en_input})
                badge = "🌉 *تم الردّ عبر جسر الترجمة (عربي ← إنجليزي ← عربي)*\n\n"
            else:
                gen_messages = messages
                badge = "🟢 *ردّ مباشر من Qwen*\n\n"

            tok, mdl = get_qwen_model(model_choice)
            output_ids = qwen_generate(gen_messages, enable_thinking, tok, mdl)

            # نفصل التفكير عن الجواب اعتماداً على رمز نهاية التفكير </think> (151668).
            # نقوم بهذا دائماً (حتى في الوضع العادي) حتى لا تتسرّب كواليس التفكير إلى الجواب.
            THINK_END = 151668
            thinking_content = ""
            if THINK_END in output_ids:
                index = len(output_ids) - output_ids[::-1].index(THINK_END)
                thinking_content = tok.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
                assistant_content = tok.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
            elif enable_thinking:
                # وضع تفكير لكن لم يظهر رمز الإغلاق → استُهلك حد التوليد أثناء التفكير
                thinking_content = tok.decode(output_ids, skip_special_tokens=True).strip("\n")
                assistant_content = ""
            else:
                # وضع عادي بلا كتلة تفكير → كامل الإخراج هو الجواب
                assistant_content = tok.decode(output_ids, skip_special_tokens=True).strip("\n")

            # في الوضع العادي لا نعرض صندوق التفكير إطلاقاً
            if not enable_thinking:
                thinking_content = ""

            # إن لم يصل جواب فعلي (انقطع التوليد داخل التفكير) ننبّه المستخدم بدل عرض التفكير كأنه جواب
            if not assistant_content:
                assistant_content = "⚠️ لم يكتمل الجواب ضمن حدّ التوليد. جرّب إيقاف وضع التفكير أو إعادة المحاولة."

            # 2) ترجمة الجواب الإنجليزي رجوعاً إلى العربية
            if use_bridge and assistant_content:
                assistant_content = en_to_ar(assistant_content)

            if thinking_content:
                final_output = f"{badge}<details><summary>🧠 كواليس تفكير النموذج</summary>\n\n{thinking_content}\n\n</details>\n\n{assistant_content}"
            else:
                final_output = f"{badge}{assistant_content}"
        except Exception as e:
            # نُظهر الخطأ الفعلي بدل رسالة Gradio العامة حتى يسهل تشخيصه
            import traceback
            traceback.print_exc()
            final_output = f"❌ خطأ أثناء تشغيل النموذج المحلي ({model_choice}):\n\n`{type(e).__name__}: {str(e)}`"

        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": final_output})
        return chat_history, ""

    # ------------------------------------------
    # المسار الثاني: GPT-5 عبر الـ API
    # ------------------------------------------
    elif model_choice == GPT_LABEL:
        if not api_key.strip():
            chat_history.append({"role": "user", "content": user_input})
            chat_history.append({"role": "assistant", "content": "⚠️ يرجى إدخال الـ API Key أولاً."})
            return chat_history, ""
        try:
            client = OpenAI(api_key=api_key)
            response = client.chat.completions.create(model="gpt-5", messages=messages, max_tokens=1024)
            final_output = response.choices[0].message.content
        except Exception as e:
            final_output = f"❌ خطأ: {str(e)}"

        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": final_output})
        return chat_history, ""

# ==========================================
# بناء واجهة المستخدم (Gradio UI Layout)
# ==========================================
with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate")) as demo:
    gr.Markdown("# 🤖 منصة المحادثة الذكية متعددة النماذج")

    with gr.Row():
        with gr.Column(scale=1, min_width=280):
            gr.Markdown("### ⚙️ إعدادات النظام")
            system_prompt_input = gr.Textbox(
                label="توجيهات النظام (System Prompt)",
                value="أنت مساعد ذكي ومرح، تجيب باختصار ودقة، وتتحدث باللغة العربية المفهومة.",
                lines=4,
                placeholder="اكتب تعليمات الموديل هنا..."
            )
            model_dropdown = gr.Dropdown(choices=[QWEN_MODEL, GPT_LABEL], value=QWEN_MODEL, label="اختر النموذج")
            thinking_checkbox = gr.Checkbox(label="🧠 تفعيل وضع التفكير (Qwen فقط)", value=True)
            bridge_checkbox = gr.Checkbox(label="🌉 الجسر العربي التلقائي (يحسّن أداء النموذج مع العربية)", value=True)
            api_key_input = gr.Textbox(label="OpenAI API Key", type="password")
            clear_button = gr.Button("🧹 مسح الذاكرة")

        with gr.Column(scale=4):
            chatbot = gr.Chatbot(label="شاشة الشات", height=550)
            with gr.Row():
                user_msg = gr.Textbox(placeholder="اكتب سؤالك هنا واضغط Enter...", scale=4)
                submit_button = gr.Button("إرسال 🚀", variant="primary", scale=1)

    submit_event = submit_button.click(
        chat_engine,
        inputs=[user_msg, chatbot, model_dropdown, thinking_checkbox, api_key_input, system_prompt_input, bridge_checkbox],
        outputs=[chatbot, user_msg]
    )
    user_msg.submit(
        chat_engine,
        inputs=[user_msg, chatbot, model_dropdown, thinking_checkbox, api_key_input, system_prompt_input, bridge_checkbox],
        outputs=[chatbot, user_msg]
    )
    clear_button.click(fn=lambda: [], outputs=chatbot, queue=False)

if __name__ == "__main__":
    # على Colab نحتاج share=True لإنشاء رابط عام؛ على Hugging Face Space يجب ألا نفعّله.
    on_hf_space = os.environ.get("SPACE_ID") is not None
    demo.launch(share=not on_hf_space)