app.py

import gradio as gr
import json
import os
from pathlib import Path
from dotenv import load_dotenv
from typing import List, Dict

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.documents import Document

# ================== CONFIG ==================
DATA_PATH = Path("data/allProjects.fixed.jsonl")
CHROMA_DIR = Path("./chroma_db_structured")
MODEL_NAME = "gpt-4o-mini"
EMBEDDING_MODEL = "text-embedding-3-small"

# ================== ENV ==================
load_dotenv()
if not os.getenv("OPENAI_API_KEY"):
    raise EnvironmentError("OPENAI_API_KEY not found.")

# ================== LOAD DATA ==================
def load_projects(path: Path) -> List[Dict]:
    projects = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            data = json.loads(line)
            projects.append({
                "id": data.get("id"),
                "title": data.get("title", ""),
                "content": data.get("content", ""),
                "image": data.get("image", ""),
                "url": data.get("url", "")
            })
    return projects

PROJECTS = load_projects(DATA_PATH)
PROJECT_COUNT = len(PROJECTS)
PROJECTS_BY_ID = {p["id"]: p for p in PROJECTS if p.get("id")}
ALL_TITLES = "\n".join([f"{p['id']}: {p['title']}" for p in PROJECTS])

# ================== VECTOR STORE ==================
embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)

def build_vectorstore():
    """بناء أو تحميل ChromaDB مع التحقق من وجود البيانات"""
    vectorstore = Chroma(
        embedding_function=embeddings,
        persist_directory=str(CHROMA_DIR)
    )
    
    # ✅ التحقق إذا كانت قاعدة البيانات فارغة وإعادة بنائها
    existing_count = vectorstore._collection.count()
    print(f"[ChromaDB] عدد الوثائق الموجودة: {existing_count}")
    
    if existing_count == 0:
        print("[ChromaDB] قاعدة البيانات فارغة، جاري البناء...")
        docs = []
        for p in PROJECTS:
            # دمج العنوان والمحتوى لتحسين البحث
            combined_text = f"عنوان المشروع: {p['title']}\n\nتفاصيل المشروع:\n{p['content']}"
            docs.append(Document(
                page_content=combined_text,
                metadata={
                    "id": p["id"],
                    "title": p["title"],
                    "url": p.get("url", ""),
                    "image": p.get("image", "")
                }
            ))
        
        # إضافة الوثائق على دفعات لتجنب timeout
        BATCH_SIZE = 50
        for i in range(0, len(docs), BATCH_SIZE):
            batch = docs[i:i+BATCH_SIZE]
            vectorstore.add_documents(batch)
            print(f"[ChromaDB] تمت إضافة {min(i+BATCH_SIZE, len(docs))}/{len(docs)}")
        
        print(f"[ChromaDB] ✅ تم البناء بنجاح. إجمالي: {vectorstore._collection.count()}")
    
    return vectorstore

vectorstore = build_vectorstore()
retriever = vectorstore.as_retriever(search_kwargs={"k": 10})

# ================== LLM ==================
llm = ChatOpenAI(model=MODEL_NAME, temperature=0.5)

# ================== HELPERS ==================
def is_project_question(message: str) -> bool:
    keywords = ["مشروع", "مشاريع", "العتبة", "العباسية", "كربلاء"]
    return any(k in message for k in keywords)


def keyword_fallback_search(query: str, top_k: int = 5) -> List[Dict]:
    """
    ✅ بحث احتياطي بالكلمات المفتاحية إذا فشل البحث الدلالي
    يبحث في العناوين والمحتوى مباشرة
    """
    query_words = query.split()
    results = []
    
    for project in PROJECTS:
        score = 0
        title = project.get("title", "")
        content = project.get("content", "")
        
        for word in query_words:
            if len(word) < 3:  # تجاهل الكلمات القصيرة
                continue
            if word in title:
                score += 3  # العنوان أكثر أهمية
            if word in content:
                score += 1
        
        if score > 0:
            results.append((score, project))
    
    results.sort(key=lambda x: x[0], reverse=True)
    return [r[1] for r in results[:top_k]]


def get_best_matching_project(query: str, matched_projects: List[Dict]) -> Dict:
    """
    ✅ يطلب من LLM اختيار المشروع الأنسب من القائمة
    بدلاً من الاعتماد على الترتيب العشوائي من ChromaDB
    """
    if len(matched_projects) == 1:
        return matched_projects[0]
    
    projects_list = "\n".join([
        f"- رقم {p['id']}: {p['title']}"
        for p in matched_projects
    ])
    
    messages = [
        SystemMessage(content="أنت محلل. أجب فقط برقم المشروع الأنسب للسؤال، بدون أي كلام إضافي."),
        HumanMessage(content=f"""
السؤال: {query}

المشاريع المتاحة:
{projects_list}

أجب فقط بـ رقم المشروع الأنسب (مثال: 231)
""")
    ]
    
    response = llm.invoke(messages).content.strip()
    
    # استخراج الرقم من الرد
    import re
    match = re.search(r'\d+', response)
    if match:
        selected_id = int(match.group())
        for p in matched_projects:
            if p.get("id") == selected_id:
                return p
    
    # fallback: أول نتيجة
    return matched_projects[0]

# ================== SMART CHAT ==================
def chat(message, history):
    message = message.strip()
    
    if not is_project_question(message):
        return "⚠️ أنا متخصص فقط في مشاريع العتبة العباسية المقدسة. يرجى طرح سؤال متعلق بالمشاريع."

    # -----------------------------
    # 1️⃣ سؤال عام جداً → إحصائية سريعة
    # -----------------------------
    if len(message.split()) <= 6 and "مشاريع" in message and "مشروع" not in message:
        messages = [
            SystemMessage(content="أنت محلل مشاريع. أجب بإحصائية دقيقة."),
            HumanMessage(content=f"""
سؤال: {message}

عدد المشاريع الكلي: {PROJECT_COUNT}

قائمة العناوين:
{ALL_TITLES}
""")
        ]
        return llm.invoke(messages).content

    # -----------------------------
    # 2️⃣ البحث الدلالي
    # -----------------------------
    docs = retriever.invoke(message)
    
    # ✅ إذا كان البحث الدلالي لم يُعطِ نتائج، استخدم البحث الاحتياطي
    matched_projects = []
    
    if docs:
        # استخراج المشاريع من نتائج ChromaDB
        for doc in docs:
            pid = doc.metadata.get("id")
            if pid and pid in PROJECTS_BY_ID:
                matched_projects.append(PROJECTS_BY_ID[pid])
    
    if not matched_projects:
        print(f"[DEBUG] البحث الدلالي لم يجد نتائج، جاري البحث الاحتياطي...")
        matched_projects = keyword_fallback_search(message)
    
    if not matched_projects:
        return "لا توجد معلومات كافية في قاعدة بيانات المشاريع."

    # بناء السياق
    if docs:
        context = "\n\n".join(doc.page_content for doc in docs)
    else:
        context = "\n\n".join(
            f"عنوان: {p['title']}\nتفاصيل: {p['content']}"
            for p in matched_projects
        )

    # -----------------------------
    # 3️⃣ توليد الإجابة
    # -----------------------------
    messages = [
        SystemMessage(content="أنت مساعد متخصص في مشاريع العتبة العباسية. أجب بدقة وتفصيل."),
        HumanMessage(content=f"""
سؤال المستخدم:
{message}

محتوى المشاريع ذات الصلة:
{context}
""")
    ]

    answer = llm.invoke(messages).content

    # -----------------------------
    # 4️⃣ إضافة بيانات أول مشروع مطابق
    # -----------------------------
    best_project = get_best_matching_project(message, matched_projects)  # ✅ بعد

    # best_project = matched_projects[0] if matched_projects else None

    if best_project:
        answer += "\n\n---\n"
        answer += f"### 📌 رقم المشروع: {best_project.get('id')}\n"
        answer += f"### 🏷️ العنوان: {best_project.get('title')}\n"

        if best_project.get("image"):
            answer += f"\n### 📷 صورة المشروع:\n![project]({best_project['image']})\n"

        if best_project.get("url"):
            answer += f"\n### 🔗 رابط المشروع:\n[{best_project['url']}]({best_project['url']})\n"

    return answer


# ================== UI ==================
INITIAL_MESSAGE = """
مرحبًا 👋  
اسألني عن أي مشروع أو اطلب تحليل عام للمشاريع.
"""

def respond(message, history):
    if history is None:
        history = []
    history.append({"role": "user", "content": message})
    bot_reply = chat(message, history)
    history.append({"role": "assistant", "content": bot_reply})
    return history


with gr.Blocks(title="🕌 مساعد المشاريع") as demo:
    gr.Markdown("# 🕌 مساعد مشاريع العتبة العباسية")
    chatbot = gr.Chatbot(
        value=[{"role": "assistant", "content": INITIAL_MESSAGE}],
        height=600,
        type="messages"
    )
    msg = gr.Textbox(placeholder="اكتب سؤالك هنا...")
    msg.submit(respond, [msg, chatbot], chatbot)

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