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




# 1. تحميل النموذج
model_id = "yasserrmd/kallamni-1.2b-v1"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 2. System Prompt + Few-shot
system_prompt = {
    "role": "system",
    "content": (
        "أنت مساعد إماراتي باللهجة الإماراتية المحكية، تجاوب دايمًا "
        "بأسلوب عفوي وقصير مثل كلام الربع، بدون فصحى."
    )
}
few_shot = [
    {"role": "user", "content": "شحالَك اليوم؟"},
    {"role": "assistant", "content": "الحمدلله زين، وانت كيفك؟"},
    {"role": "user", "content": "وين ناوي تسير عقب الدوام؟"},
    {"role": "assistant", "content": "يمكن أمر على المول وأتعشى ويا الربع."},
]

# 3. دالة التوليد
@spaces.GPU
def chat_fn(message, history):
    try:
        # Gradio ChatInterface history = list of dicts [{"role":..., "content":...}]
        messages = [system_prompt] + few_shot + history + [{"role": "user", "content": message}]

        # تجهيز الإدخال
        input_ids = tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            tokenize=True,
        ).to(model.device)

        # التوليد
        output = model.generate(
            input_ids,
            do_sample=True,
            temperature=0.1,
            min_p=0.15,
            repetition_penalty=1.05,
            max_new_tokens=60,
        )

        decoded = tokenizer.decode(output[0], skip_special_tokens=False)

        # استخراج رد المساعد الأخير
        try:
            a_start = decoded.rindex("<|im_start|>assistant") + len("<|im_start|>assistant")
            a_end = decoded.index("<|im_end|>", a_start)
            answer = decoded[a_start:a_end].strip()
        except ValueError:
            answer = decoded.strip()

        return answer

    except Exception as e:
        return f"[خطأ داخلي]: {str(e)}"

# 4. CSS للـ RTL
css = """
#chat-container { direction: rtl; text-align: right; }
"""

# 5. واجهة Gradio
with gr.Blocks(css=css, fill_height=True) as demo:
    gr.HTML(
        """<div style="text-align: center;">
             <img src="https://huggingface.co/spaces/yasserrmd/Kallamni-chat/resolve/main/assets/logo.png" 
                  alt="Logo" width="120">
           </div>"""
    )
    gr.ChatInterface(
        fn=chat_fn,
        type="messages",
        examples = [
            "وين ناوي تسير عقب ما تخلص الدوام اليوم؟",
            "شرايك في الجو هالأيام، والله تحسه حر وايد؟",
            "كيف تقضي الويكند عادةً ويا العيال والربع؟",
            "شو أحلى أكلة دوم تحبها من طبخ الوالدة؟",
            "وين أحلى مكان دوم تاخذ منه قهوة الصبح؟",
        ],
        title="💬 شات باللهجة الإماراتية",
        cache_examples=True,
        theme="soft",
        fill_height=True
    )

# 6. تشغيل Debug
if __name__ == "__main__":
    demo.launch(debug=True)