import gradio as gr
import os
import re
import faiss
import numpy as np
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer

# ========== إعداد الـ Client ==========
client = InferenceClient(
    model="Qwen/Qwen3-8B",
    token=os.environ["HF_TOKEN"]
)

embedding_model = SentenceTransformer("BAAI/bge-m3")
DISEASES_FOLDER = "diseases"

# ========== تقسيم النص ==========
def split_text(text, chunk_size=300, chunk_overlap=50):
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start += chunk_size - chunk_overlap
    return chunks

# ========== الدالة الرئيسية ==========
def ask_in_disease(question, disease_name, k=5):
    if disease_name and disease_name.strip():
        filename = os.path.join(DISEASES_FOLDER, f"{disease_name}.txt")
        if not os.path.exists(filename):
            return f"❌ الملف {filename} غير موجود!"

        with open(filename, "r", encoding="utf-8") as f:
            text = f.read()

        texts = split_text(text)
        embeddings = embedding_model.encode(texts, convert_to_numpy=True)

        index = faiss.IndexFlatL2(embeddings.shape[1])
        index.add(embeddings)

        q_emb = embedding_model.encode([question], convert_to_numpy=True)
        _, indices = index.search(q_emb, k)

        context = "\n\n".join([texts[i] for i in indices[0]])

        messages = [
            {
                "role": "system",
                "content": (
                    "أنت مساعد زراعي ذكي متخصص في أمراض النباتات.\n"
                    "أجب بلغة السؤال لو السؤال عربي جاوب بالعربي ولو السؤال انجليزي جاوب بالانجليزي اوعي تتكلم اي لغه غير اللغتين دول.\n"
                    "استخدم المعلومات الموجودة في النصوص المسترجعة فقط ولو المعلومه مش موجوده بشكل مباشر حاول تستنجها من المعلومات الموجودة.\n"
                    "إذا لم تجد إجابة واضحة قل: لا توجد معلومات كافية.\n"
                    "لا تخترع معلومات."
                )
            },
            {
                "role": "user",
                "content": f"المعلومات:\n{context}\n\nالسؤال:\n{question}"
            }
        ]
    else:
        messages = [
            {
                "role": "system",
                "content": (
                    "أنت مساعد زراعي ذكي متخصص في أمراض النباتات.\n"
                    "أجب بلغة السؤال لو السؤال عربي جاوب بالعربي ولو السؤال انجليزي جاوب بالانجليزي اوعي تتكلم اي لغه غير اللغتين دول.\n"
                    "إذا لم تجد إجابة واضحة قل: لا توجد معلومات كافية.\n"
                    "لا تخترع معلومات."
                )
            },
            {
                "role": "user",
                "content": f"السؤال:\n{question}"
            }
        ]

    response = client.chat_completion(
        messages=messages,
        max_tokens=400,
        temperature=0.2,
        extra_body={"chat_template_kwargs": {"enable_thinking": False}}
    )

    output = response.choices[0].message.content.strip()
    output = re.sub(r"<think>.*?</think>", "", output, flags=re.DOTALL).strip()
    return output

# ========== Gradio Interface ==========
demo = gr.Interface(
    fn=ask_in_disease,
    inputs=[
        gr.Textbox(label="اكتب سؤالك", placeholder="مثال: ما هي أعراض المرض؟"),
        gr.Textbox(label="اسم المرض", placeholder="مثال: Grape___Black_rot")
    ],
    outputs=gr.Textbox(label="الإجابة"),
    title="🌿 المساعد الزراعي الذكي",
    description="اسأل عن أي مرض نباتي وهنجاوبك!"
)

demo.launch()