# app.py

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

# ====================== 1) Load model ======================
model_name = "ALLaM-AI/ALLaM-7B-Instruct-preview"  # ✅ الصحيح
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16,
    trust_remote_code=True,
    revision="main"
)

# ====================== 2) System prompts ======================
allowed_keywords = [
    # 🌴 Palm-related
    "palm", "palms", "نخلة", "نخل", "نخلي", "فسيلة", "فسائل", "جذع", "سعف", "خوص", "جريد", "شماريخ", "عذوق", "مزرعة نخل", "offshoot", "date palm",
    # 🐛 Diseases
    "مرض النخيل", "ذبول", "لفحة", "اللفحة السوداء", "فطريات", "قمل أبيض", "ذبول الفيوزاريم", "تبقعات الأوراق", "نقص المغنيسيوم", "نقص البوتاسيوم", "اصفرار الخوص", "احتراق الجريد", "الثمر ما يكتمل",
    "Black scorch", "Fusarium wilt", "Rachis blight", "Leaf spot", "Mites", "Insects", "White bugs", "Parlatoria", "Magnesium deficiency", "Potassium deficiency", "Manganese deficiency",
    # 🍇 Dates
    "تمر", "تمور", "تمر سكري", "سكري", "خلاص", "عجوة", "مجدول", "روثانة", "رطب", "برني", "عنبري", "صفاوي", "صقعي", "خضري", "فاخر", "premium dates", "best dates", "Ajwa", "Medjool", "Sukkary", "Khalas", "Safawi", "Sagai", "Khudri", "Ruthana", "Barni", "Anbara",
    # 🌱 Care & Maintenance
    "ري", "سقي", "تسميد", "مبيد", "متى أسقي النخلة", "نصائح العناية", "رش", "تقليم", "عناية", "كيف أعتني",
    "palm care", "how to water", "fertilizer", "pest control", "sunlight", "organic spray", "how to prune",
    # 🌦️ Stats & Weather
    "كم نخلة", "كم تمر", "عدد النخل", "الجو", "الطقس", "الحرارة", "الرطوبة", "مناسب للتلقيح", "is it good weather", "weather", "temperature", "total palms", "healthy palms", "sick palms",
    # 🛒 Consumer
    "أين أشتري تمر", "أفضل تمر", "تمور مغشوشة", "تمور القصيم", "تمور المدينة", "جودة التمر", "التغليف", "شراء تمر", "buy dates", "where to find", "identify good dates", "how to store dates",
    # 📸 App features
    "افتح الكاميرا", "حلل الصورة", "قيم النخلة", "قيم التمر", "camera", "analyze", "scan", "image detection",
    # ℹ️ General
    "تطبيق لينة", "نظام لينة", "عن لينة", "Lina app", "explain Lina", "help with Lina", "what is Lina"
]

greetings = ["سلام", "السلام عليكم", "أهلاً", "هاي", "hi", "hello", "hey"]
intro_questions = ["من أنت", "مين انت", "what is this", "who are you", "explain the app", "about lina"]

# ====================== 3) Chat function ======================
def chat(user_message):
    # Check if input is allowed
    if not any(keyword in user_message.lower() for keyword in allowed_keywords + greetings + intro_questions):
        return "❗ عذرًا، هذا النظام متخصص فقط بالنخيل والتمور. حاول أن تسأل عن شيء متعلق بذلك 🌴."

    # Build prompt
    prompt = f"<|user|>\n{user_message}\n<|assistant|>"

    # Generate answer
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = model.generate(
        **inputs,
        max_new_tokens=500,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id
    )
    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)

    # Extract the assistant's reply
    reply = decoded_output.split("<|assistant|>")[-1].strip()
    return reply

# ====================== 4) Gradio Interface ======================
iface = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=2, placeholder="اكتب سؤالك عن النخيل أو التمر هنا..."),
    outputs=gr.Textbox(),
    title="مساعد لينة - Lina Assistant 🌴",
    description="اسألني عن النخيل، التمور، العناية، الأمراض، الطقس أو أي شيء متعلق! 🍃"
)

iface.launch()