app.py

import os
import re
import time
import streamlit as st
from dotenv import load_dotenv
from config import *
from utils import normalize_arabic, detect_lang, format_latency, extract_pdf_text
from rag_engine import ENAEngine
from agentic_rag import AgenticRAG

# ══════════════════════════════════════════════════════════
# ⚙️ INITIALIZATION
# ══════════════════════════════════════════════════════════
load_dotenv()
st.set_page_config(page_title="ENA Chatbot PRO", page_icon="🎓", layout="wide")
st.markdown(UI_CSS, unsafe_allow_html=True)

@st.cache_resource
def get_engine():
    token = os.getenv("GROQ_TOKEN") or st.secrets.get("GROQ_TOKEN")
    return ENAEngine(groq_token=token)

@st.cache_resource
def get_agent(_engine):
    return AgenticRAG(_engine)

engine = get_engine()
agent  = get_agent(engine)

# ══════════════════════════════════════════════════════════
# 🛠️ UI HELPERS
# ══════════════════════════════════════════════════════════
def render_sidebar():
    with st.sidebar:
        st.title("🎓 ENA Chatbot PRO")
        st.caption("v3.0 | Optimized Engine")
        
        st.divider()
        lang = st.radio("🌐 Language / اللغة", ["العربية", "Français"], horizontal=True)
        st.session_state.lang = "ar" if lang == "العربية" else "fr"
        
        st.divider()
        show_conf = st.toggle("📊 مؤشر الثقة", value=True)
        show_obs = st.toggle("📡 Observability", value=False)
        
        if show_obs and "logs" in st.session_state and st.session_state.logs:
            log = st.session_state.logs[-1]
            st.metric("Latency", format_latency(log.get("total_ms", 0)))
            with st.expander("Performance Details"):
                st.json(log)
                
        st.divider()
        if st.button("🗑️ مسح المحادثة", use_container_width=True):
            st.session_state.chat_messages = []
            st.session_state.last_sources = []
            st.rerun()
            
        # PDF Section
        st.divider()
        st.markdown("### 📄 مستشار الوثائق (PDF)")
        uploaded_pdf = st.file_uploader("ارفع وثيقة (بلاغ/رائد رسمي)", type=["pdf"])
        if uploaded_pdf:
            with st.spinner("جاري قراءة الملف..."):
                pdf_text = extract_pdf_text(uploaded_pdf.read())
                st.session_state.uploaded_pdf_text = pdf_text
                st.success("✅ تمت قراءة الملف! اسألني ومصادري ستشمل هذا الملف.")
        else:
            st.session_state.uploaded_pdf_text = ""
            
        st.divider()
        
        # DB Status
        try:
            db_count = 0
            try:
                db_count = engine.vectordb._collection.count()
            except:
                # If collection get fails, we treat it as empty 
                db_count = 0

            if db_count > 0:
                st.success(f"✅ المتوفر: {db_count} معلومة")
            else:
                st.warning("⚠️ قاعدة البيانات فارغة!")
                if st.button("🏗️ بناء القاعدة من الملف الأصلي (JSON)"):
                    with st.spinner("✨ جاري المعالجة... يرجى الانتظار دقيقتين"):
                        try:
                            import build_chroma
                            build_chroma.build()
                            st.cache_resource.clear()
                            st.success("✅ تم البناء بنجاح!")
                            st.rerun()
                        except Exception as e:
                            st.error(f"❌ فشل البناء: {e}")
                
                if st.button("📥 فك ضغط ZIP (بديل سريع)"):
                    import zipfile
                    if os.path.exists("chroma_ena_db.zip"):
                        with zipfile.ZipFile("chroma_ena_db.zip", 'r') as z:
                            z.extractall("./temp_db")
                            src = "./temp_db/chroma_ena_db" if os.path.exists("./temp_db/chroma_ena_db") else "./temp_db"
                            if not os.path.exists(CHROMA_PATH): os.makedirs(CHROMA_PATH)
                            import shutil
                            for f in os.listdir(src):
                                shutil.move(os.path.join(src, f), os.path.join(CHROMA_PATH, f))
                        st.cache_resource.clear()
                        st.success("✅ تم فك الضغط!")
                        st.rerun()
        except Exception as e:
            st.error(f"❌ خطأ في الاتصال")
            with st.expander("تفاصيل الخطأ التقني"):
                st.code(str(e))
            if st.button("🧹 مسح وإعادة بناء"):
                import shutil
                if os.path.exists(CHROMA_PATH): shutil.rmtree(CHROMA_PATH)
                st.rerun()
            
        return show_conf

def render_welcome():
    is_ar = st.session_state.get("lang", "ar") == "ar"
    title = "🎓 ENA Chatbot PRO"
    subtitle = "مساعدك الذكي للمناظرات والتكوين بالمدرسة الوطنية للإدارة" if is_ar else "Votre assistant intelligent pour les concours et formations de l'ENA"
    
    st.markdown(f"""
    <div class="welcome-banner">
        <h1>{title}</h1>
        <p>{subtitle}</p>
    </div>
    """, unsafe_allow_html=True)
    
    sug_label = "💡 أسئلة مقترحة:" if is_ar else "💡 Questions suggérées :"
    st.markdown(f'<div style="text-align: {"right" if is_ar else "left"}; margin-top: 20px;">{sug_label}</div>', unsafe_allow_html=True)
    
    cols = st.columns(4)
    ar_questions = ["شروط مناظرة المرحلة العليا", "الوثائق المطلوبة للتسجيل", "برنامج مناظرة أ2", "تاريخ فتح المناظرات القادمة"]
    fr_questions = ["Conditions Cycle Supérieur", "Documents d'inscription", "Programme Concours A2", "Dates des prochains concours"]
    questions = ar_questions if is_ar else fr_questions
    
    for i, q in enumerate(questions):
        if cols[i].button(q, key=f"q_{i}", use_container_width=True):
            st.session_state.pending_q = q
            st.rerun()

# ══════════════════════════════════════════════════════════
# 🚀 MAIN APP
# ══════════════════════════════════════════════════════════
def main():
    show_conf = render_sidebar()
    
    if "chat_messages" not in st.session_state:
        st.session_state.chat_messages = []
    
    if not st.session_state.chat_messages:
        render_welcome()

    # Layout for Chat + Sources
    col_chat, col_sources = st.columns([0.7, 0.3])

    with col_chat:
        # Display history
        for msg in st.session_state.chat_messages:
            with st.chat_message(msg["role"]):
                st.markdown(msg["content"])

        # ══ عرض الاقتراحات بعد آخر إجابة ══
        suggestions = st.session_state.get("last_suggestions", [])
        if suggestions and st.session_state.chat_messages:
            lang = st.session_state.get("lang", "ar")
            sug_label = "💡 قد يهمك أيضاً:" if lang == "ar" else "💡 Vous pourriez aussi demander:"
            st.markdown(f"---\n**{sug_label}**")
            sug_cols = st.columns(len(suggestions))
            for i, sug in enumerate(suggestions):
                # دعم النسختين: dict (جديد) أو string (قديم)
                if isinstance(sug, dict):
                    icon  = sug.get("icon", "💡")
                    title = sug.get("title", "")
                    desc  = sug.get("desc", "")
                    question = sug.get("question", title)
                    label = f"{icon} **{title}**\n_{desc}_"
                else:
                    label    = sug
                    question = sug
                with sug_cols[i]:
                    if isinstance(sug, dict):
                        st.markdown(f"{icon} **{title}**")
                        st.caption(desc)
                        if st.button("اسأل ←" if lang == "ar" else "Demander →", key=f"sug_{len(st.session_state.chat_messages)}_{i}"):
                            st.session_state.last_suggestions = []
                            st.session_state.pending_q = question
                            st.rerun()
                    else:
                        if st.button(sug, key=f"sug_{len(st.session_state.chat_messages)}_{i}"):
                            st.session_state.last_suggestions = []
                            st.session_state.pending_q = question
                            st.rerun()

        # Input
        is_ar = st.session_state.get("lang", "ar") == "ar"
        placeholder = "اسأل عن المناظرات، الشروط، أو إجراءات التكوين..." if is_ar else "Posez des questions sur les concours, conditions ou procédures..."
        prompt = st.chat_input(placeholder)
        pending_q = st.session_state.pop("pending_q", None)
        final_q = prompt or pending_q

        if final_q:
            with st.chat_message("user"):
                st.markdown(final_q)

            with st.chat_message("assistant"):
                t_start = time.time()
                lang    = detect_lang(final_q)
                history = st.session_state.chat_messages[:]
                # نضيف سؤال المستخدم للتاريخ هنا بعد ما نأخذ نسخة من التاريخ القديم
                st.session_state.chat_messages.append({"role": "user", "content": final_q})

                # ══ إذا المستخدم أجاب على أسئلة التأهل ══
                if st.session_state.get("waiting_for_profile"):
                    st.session_state.waiting_for_profile = False
                    with st.spinner("🔍 جاري تقييم ملفك..."):
                        evaluation = agent.evaluate_user_profile(final_q, lang)
                    st.markdown(evaluation)
                    st.session_state.chat_messages.append({"role": "assistant", "content": evaluation})
                    st.session_state.last_tool = "evaluate_profile"
                    st.rerun()

                # ══ Agentic RAG ══
                else:
                    spinner_msg = "🤖 جاري التفكير..." if lang == "ar" else "🤖 Recherche en cours..."
                    with st.spinner(spinner_msg):
                        uploaded_pdf = st.session_state.get("uploaded_pdf_text", "")
                        result = agent.run(final_q, history, uploaded_pdf)

                    full_response = result["answer"]
                    tool_used     = result.get("tool_used", "")

                    # إذا يحتاج معلومات المستخدم
                    if result.get("needs_user_input"):
                        st.session_state.waiting_for_profile = True

                    st.markdown(full_response)

                    # ══ اقتراحات ذكية — نحفظها في session_state قبل rerun ══
                    if tool_used not in ("chat", "ask_user_profile") and not result.get("needs_user_input"):
                        with st.spinner("💡 ..."):
                            suggestions = agent.generate_suggestions(final_q, full_response, lang)
                        st.session_state.last_suggestions = suggestions if suggestions else []
                    else:
                        st.session_state.last_suggestions = []

                    st.session_state.chat_messages.append({"role": "assistant", "content": full_response})
                    st.session_state.last_tool = tool_used
                    st.session_state.last_sources = result.get("sources", [])

                    # Log
                    st.session_state.logs = st.session_state.get("logs", [])
                    st.session_state.logs.append({
                        "total_ms": round((time.time() - t_start) * 1000),
                        "tool": tool_used
                    })
                    st.rerun()

    with col_sources:
        is_ar  = st.session_state.get("lang", "ar") == "ar"
        header = "🤖 الأداة المستخدمة" if is_ar else "🤖 Outil utilisé"
        st.subheader(header)

        tool_used = st.session_state.get("last_tool", "")
        tool_labels = {
            "search_db":         ("🔍 بحث في قاعدة البيانات",     "🔍 Recherche DB"),
            "get_concours_page": ("📄 صفحة المناظرة الكاملة",      "📄 Page complète du concours"),
            "ask_user_profile":  ("👤 تقييم الملف الشخصي",         "👤 Évaluation du profil"),
            "evaluate_profile":  ("✅ تقييم التأهل",                "✅ Évaluation éligibilité"),
            "summarize_url":     ("📝 تلخيص مقال",                  "📝 Résumé article"),
            "chat":              ("💬 محادثة عامة",                 "💬 Conversation"),
            "search_db (fallback)": ("🔍 بحث (احتياطي)",           "🔍 Recherche (fallback)"),
        }

        if tool_used and tool_used in tool_labels:
            label = tool_labels[tool_used][0] if is_ar else tool_labels[tool_used][1]
            st.success(label)

            # شرح ليش اختار هاذي الأداة
            explanations = {
                "search_db":         "سؤال عام — بحثت في قاعدة البيانات",
                "get_concours_page": "سؤال عن شروط/وثائق — جبت الصفحة كاملة مباشرة من الموقع",
                "ask_user_profile":  "سؤال شخصي — طلبت معلوماتك لتقييم تأهلك",
                "evaluate_profile":  "قيّمت ملفك بناءً على الوثائق الرسمية",
                "summarize_url":     "لخّصت المقال المطلوب",
            }
            if tool_used in explanations:
                st.caption(explanations[tool_used])
                
            # عرض المصادر
            sources = st.session_state.get("last_sources", [])
            if sources:
                st.markdown("---")
                src_header = "📚 المصادر ذات الصلة" if is_ar else "📚 Sources liées"
                st.subheader(src_header)
                for i, src in enumerate(sources):
                    url = src.get("url", "")
                    page_name = src.get("page_name", "")
                    content = src.get("content", "")
                    if url:
                        with st.expander(f"المصدر [{i+1}]: {page_name}"):
                            st.markdown(f"[{url}]({url})")
                            if content:
                                st.text_area("النص المستخرج", content, height=120, disabled=True, key=f"src_{i}_{len(st.session_state.chat_messages)}")
        else:
            info_msg = "ستظهر الأداة المستخدمة والمصادر هنا" if is_ar else "L'outil utilisé et les sources apparaîtront ici"
            st.info(info_msg)

if __name__ == "__main__":
    main()