main.py

"""
Saudi Legal AI — HuggingFace Space
v7.0 | Gemini-first · Groq fallback · FastAPI + HTML frontend
"""
import os, re, time, json, gc, hashlib
from collections import deque, defaultdict, Counter
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from pydantic import BaseModel

from groq import Groq
from huggingface_hub import login
import google.generativeai as genai
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import SentenceTransformerEmbeddings
from datasets import load_dataset
from rank_bm25 import BM25Okapi
from rapidfuzz import fuzz, process as fuzz_proc
import chromadb

# ══════════════════════════════════════════════════════════
# Config
# ══════════════════════════════════════════════════════════
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
HF_TOKEN     = os.getenv("HF_TOKEN", "")
GEMINI_KEY   = os.getenv("GEMINI_API_KEY", "")
HF_REPO_ID   = os.getenv("HF_REPO_ID", "WafaaFraih/saudi-legal-moj")

# ══════════════════════════════════════════════════════════
# Global State
# ══════════════════════════════════════════════════════════
vectorstore:      Optional[Chroma]    = None
bm25_index:       Optional[BM25Okapi] = None
bm25_texts:       list = []
bm25_metas:       list = []
embeddings                            = None
ACTIVE_MODELS:    list = []
GEMINI_AVAILABLE: bool = False
gemini_clients:   dict = {}
groq_client:      Optional[Groq] = None

request_log  = deque(maxlen=500)
api_stats    = defaultdict(int)
user_ips     = defaultdict(list)
_trans_cache  = {}
_legal_cache  = {}
_answer_cache = {}
_rewrite_cache= {}

GEMINI_PRIORITY = [
    "models/gemini-2.5-flash",
    "models/gemini-2.0-flash",
    "models/gemma-3-27b-it",
]

GEMINI_SAFETY = [
    {"category": "HARM_CATEGORY_HARASSMENT",        "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH",       "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
]

# ══════════════════════════════════════════════════════════
# Manual supplements (not in HF dataset)
# ══════════════════════════════════════════════════════════
MANUAL = [
    # نظام العمل
    {"text":"المادة الثالثة والثمانون: عند انتهاء عقد العمل يستحق العامل مكافأة عن مدة خدمته تحسب على أساس أجر نصف شهر عن كل سنة من السنوات الخمس الأولى، وأجر شهر عن كل سنة بعد ذلك. وتحسب المكافأة على أساس آخر أجر تقاضاه العامل.",
     "article_number":"المادة الثالثة والثمانون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    {"text":"المادة الثامنة والثمانون: إذا أنهى صاحب العمل عقد العمل دون سبب مشروع وجب عليه دفع تعويض يعادل أجر خمسة عشر يوماً عن كل سنة خدمة ولا يقل عن أجر شهرين.",
     "article_number":"المادة الثامنة والثمانون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    {"text":"المادة التاسعة والثمانون: لا يجوز لصاحب العمل فصل العامل بسبب تقدمه بشكوى. وإذا أثبت العامل أن الفصل كان تعسفياً وجب دفع تعويض عادل إضافة إلى مكافأة نهاية الخدمة.",
     "article_number":"المادة التاسعة والثمانون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    {"text":"المادة الخامسة والستون: لا يجوز تشغيل العامل أكثر من ثماني ساعات في اليوم أو ثماني وأربعين ساعة في الأسبوع. وفي شهر رمضان تنخفض ساعات العمل للمسلمين إلى ست ساعات يومياً.",
     "article_number":"المادة الخامسة والستون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    {"text":"المادة الثامنة والستون: يستحق العامل عن ساعات العمل الإضافية أجراً إضافياً لا يقل عن أجره الأصلي مضافاً إليه خمسون بالمئة من ذلك الأجر.",
     "article_number":"المادة الثامنة والستون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    {"text":"المادة الثالثة والستون: تُحدَّد مدة الإجازة السنوية بواحد وعشرين يوماً لكل سنة عمل، وتُزاد إلى ثلاثين يوماً إذا أمضى العامل خمس سنوات متواصلة.",
     "article_number":"المادة الثالثة والستون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    {"text":"المادة الثالثة والثلاثون: تستحق العاملة إجازة وضع بأجر كامل مدتها عشرة أسابيع. ويُحظر تشغيل المرأة في الأسابيع الستة التالية للوضع مباشرة.",
     "article_number":"المادة الثالثة والثلاثون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    {"text":"المادة السابعة والعشرون: عند انتهاء عقد العمل يلتزم صاحب العمل بإعطاء العامل شهادة خبرة مجانية تبين فيها مدة خدمته ونوع عمله ومقدار أجره.",
     "article_number":"المادة السابعة والعشرون","law_name":"نظام العمل","law_type":"نظام","source":"hrsd.gov.sa"},
    # الإجراءات الجزائية
    {"text":"المادة الثانية: لا يجوز القبض على أي إنسان أو توقيفه أو سجنه إلا في الأحوال المنصوص عليها نظاماً. ولا يُعامل المقبوض عليه إلا بما يحفظ كرامته الإنسانية.",
     "article_number":"المادة الثانية","law_name":"نظام الإجراءات الجزائية","law_type":"نظام","source":"boe.gov.sa"},
    {"text":"المادة التاسعة والستون: يحق للمتهم الاستعانة بمحامٍ خلال مراحل التحقيق والمحاكمة. ولا يجوز التحقيق مع المتهم في الجرائم الكبيرة إلا بحضور محامٍ.",
     "article_number":"المادة التاسعة والستون","law_name":"نظام الإجراءات الجزائية","law_type":"نظام","source":"boe.gov.sa"},
    {"text":"المادة الحادية والسبعون: يجب على المحقق إخبار المتهم بالتهمة المنسوبة إليه عند مثوله أمامه. وللمتهم رفض الإجابة حتى يحضر محاميه.",
     "article_number":"المادة الحادية والسبعون","law_name":"نظام الإجراءات الجزائية","law_type":"نظام","source":"boe.gov.sa"},
    {"text":"المادة الثانية والسبعون: للموقوف حق الاتصال بذويه وبمحاميه فور توقيفه. ويُعلَم بأسباب توقيفه كتابةً، وله حق الصمت وعدم الإجابة قبل حضور محاميه.",
     "article_number":"المادة الثانية والسبعون","law_name":"نظام الإجراءات الجزائية","law_type":"نظام","source":"boe.gov.sa"},
    {"text":"المادة الثالثة والسبعون: لا يجوز التوقيف في غير الأماكن المخصصة نظاماً. وللموقوف حق التظلم من قرار التوقيف أمام المحكمة المختصة.",
     "article_number":"المادة الثالثة والسبعون","law_name":"نظام الإجراءات الجزائية","law_type":"نظام","source":"boe.gov.sa"},
    # المحاماة
    {"text":"المادة الثالثة: يشترط فيمن يزاول مهنة المحاماة: أن يكون سعودي الجنسية بالأصل، حاصلاً على شهادة البكالوريوس في الشريعة أو الأنظمة، مقيداً في جدول المحامين الممارسين.",
     "article_number":"المادة الثالثة","law_name":"نظام المحاماة","law_type":"نظام","source":"boe.gov.sa"},
    # القضاء
    {"text":"المادة الثانية والثلاثون: يشترط فيمن يُعيَّن في وظائف القضاء أن يكون سعودي الجنسية بالأصل، مسلماً، حاصلاً على شهادة الأهلية في الشريعة الإسلامية أو ما يعادلها، وأن ينجح في الامتحان المقرر.",
     "article_number":"المادة الثانية والثلاثون","law_name":"نظام القضاء","law_type":"نظام","source":"boe.gov.sa"},
    # الجرائم المعلوماتية
    {"text":"المادة الثالثة: يعاقب بالسجن مدة لا تزيد على سنة وبغرامة لا تزيد على خمسمائة ألف ريال كل من يرتكب جريمة الدخول غير المشروع أو التشهير الإلكتروني.",
     "article_number":"المادة الثالثة","law_name":"نظام مكافحة الجرائم المعلوماتية","law_type":"نظام","source":"boe.gov.sa"},
    # حماية البيانات
    {"text":"المادة الرابعة: لا يجوز معالجة البيانات الشخصية إلا لتحقيق الغرض المشروع مع الحصول على موافقة صريحة من صاحب البيانات.",
     "article_number":"المادة الرابعة","law_name":"نظام حماية البيانات الشخصية","law_type":"نظام","source":"boe.gov.sa"},
    {"text":"المادة التاسعة والعشرون: يعاقب على الإفصاح عن البيانات الشخصية دون وجه حق بالسجن مدة لا تزيد على سنتين وبغرامة لا تزيد على ثلاثة ملايين ريال.",
     "article_number":"المادة التاسعة والعشرون","law_name":"نظام حماية البيانات الشخصية","law_type":"نظام","source":"boe.gov.sa"},
]

# ══════════════════════════════════════════════════════════
# Dictionaries
# ══════════════════════════════════════════════════════════
COLLOQUIAL = {
    "ايه":"ما","إيه":"ما","ايش":"ما","شو":"ما","وش":"ما",
    "اللي":"الذي","عشان":"لأن","ازاي":"كيف","امتى":"متى",
    "فين":"أين","مين":"من","ليه":"لماذا","عندي":"لدي",
    "اشتغلت":"عملت","فصلوني":"تم فصلي","طردوني":"تم فصلي",
    "اتطردت":"تم فصلي","مش":"لا","كمان":"أيضاً","زي":"مثل",
    "اللازم":"يجب","المفروض":"يجب",
    "مكافاه":"مكافأة","مكافأه":"مكافأة","مكافاة":"مكافأة",
    "رخصه":"رخصة","مزاوله":"مزاولة","المحاماه":"المحاماة",
    "ليا":"لي","هل ليا":"هل يحق لي",
    "فلوس":"تعويض مالي","الشغل":"العمل","شغل":"عمل",
    "علاوة":"زيادة في الراتب",
    "عاوز":"أريد","عايز":"أريد","عاوزة":"أريد","عايزة":"أريد",
    "ايه اللازم":"ما هي الإجراءات","ايه حقي":"ما هي حقوقي",
    "الحبس":"السجن","موقوف":"محتجز",
    "اتهمت":"تم اتهامي","قبضوا علي":"تم توقيفي",
}

ENGLISH_TO_ARABIC = {
    "conditions for lawyer license":  "شروط مزاولة مهنة المحاماة",
    "what is the penalty":            "ما هي عقوبة",
    "what are my rights":             "ما هي حقوقي",
    "how is calculated":              "كيف تحسب",
    "how to calculate":               "كيف تحسب",
    "i was dismissed":                "تم فصلي",
    "i was fired":                    "تم فصلي",
    "wrongful termination":           "فصل تعسفي",
    "unfair dismissal":               "فصل تعسفي",
    "am i entitled":                  "هل يحق لي",
    "end of service":                 "مكافأة نهاية الخدمة",
    "end of servise":                 "مكافأة نهاية الخدمة",
    "end of service gratuity":        "مكافأة نهاية الخدمة",
    "money laundering":               "غسل الأموال",
    "cybercrime":                     "الجرائم المعلوماتية",
    "working hours":                  "ساعات العمل",
    "annual leave":                   "الإجازة السنوية",
    "maternity leave":                "إجازة الأمومة",
    "overtime":                       "العمل الإضافي",
    "data protection":                "حماية البيانات الشخصية",
    "labor law":                      "نظام العمل",
    "labour law":                     "نظام العمل",
    "my rights":                      "حقوقي",
    "gratuity":                       "مكافأة نهاية الخدمة",
    "bankruptcy":                     "إفلاس",
    "arbitration":                    "تحكيم",
    "divorce":                        "طلاق",
    "custody":                        "حضانة",
    "alimony":                        "نفقة",
    "lawyer":                         "محامي",
    "attorney":                       "محامي",
    "judge":                          "قاضي",
    "penalty":                        "عقوبة",
    "penaly":                         "عقوبة",
    "fine":                           "غرامة",
    "imprisonment":                   "سجن",
    "rights":                         "حقوق",
    "rigths":                         "حقوق",
    "salary":                         "الأجر",
    "employee":                       "عامل",
    "employer":                       "صاحب عمل",
    "terminated":                     "فُصلت",
    "dismissed":                      "فُصلت",
    "dissmised":                      "فُصلت",
    "evidence":                       "إثبات",
    "notary":                         "توثيق",
    "if fired":                       "عند الفصل",
    "i worked":                       "عملت",
}

LAW_KEYWORDS = {
    "توثيق":"نظام التوثيق","كاتب عدل":"نظام التوثيق","موثق":"نظام التوثيق",
    "محاماة":"نظام المحاماة","محامي":"نظام المحاماة",
    "مزاولة مهنة المحاماة":"نظام المحاماة","مزاولة مهنة":"نظام المحاماة",
    "ترخيص محاماة":"نظام المحاماة",
    "إفلاس":"نظام الإفلاس","إعسار":"نظام الإفلاس",
    "تحكيم":"نظام التحكيم",
    "إثبات":"نظام الإثبات","شاهد":"نظام الإثبات",
    "تنفيذ حكم":"نظام التنفيذ","تنفيذ":"نظام التنفيذ",
    # الإجراءات الجزائية (أطول الـ keywords أولاً)
    "حقوق الموقوف":"نظام الإجراءات الجزائية",
    "حقوق المتهم":"نظام الإجراءات الجزائية",
    "جزائي":"نظام الإجراءات الجزائية",
    "متهم":"نظام الإجراءات الجزائية",
    "توقيف":"نظام الإجراءات الجزائية",
    "تهمة":"نظام الإجراءات الجزائية",
    "اتهمت":"نظام الإجراءات الجزائية",
    "اعتقلت":"نظام الإجراءات الجزائية",
    "محتجز":"نظام الإجراءات الجزائية",
    "جريمة":"نظام الإجراءات الجزائية",
    "موقوف":"نظام الإجراءات الجزائية",
    "قبض":"نظام الإجراءات الجزائية",
    "احتجاز":"نظام الإجراءات الجزائية",
    "قبضوا علي":"نظام الإجراءات الجزائية",
    "تم توقيفي":"نظام الإجراءات الجزائية",
    "تم اتهامي":"نظام الإجراءات الجزائية",
    # الأحوال الشخصية
    "زواج":"نظام الأحوال الشخصية","طلاق":"نظام الأحوال الشخصية",
    "نفقة":"نظام الأحوال الشخصية","حضانة":"نظام الأحوال الشخصية",
    "خلع":"نظام الأحوال الشخصية","مهر":"نظام الأحوال الشخصية",
    # العقار
    "عقار":"نظام التسجيل العيني للعقار",
    # القضاء
    "قضاء":"نظام القضاء","قاضي":"نظام القضاء",
    "تعيين القضاة":"نظام القضاء",
    # غسل الأموال
    "غسل الأموال":"نظام مكافحة غسل الأموال",
    "غسيل الأموال":"نظام مكافحة غسل الأموال",
    # المعاملات المدنية
    "أركان العقد":"نظام المعاملات المدنية",
    "عقد مدني":"نظام المعاملات المدنية",
    # نظام العمل
    "مكافأة نهاية الخدمة":"نظام العمل","نهاية الخدمة":"نظام العمل",
    "صاحب عمل":"نظام العمل","عقد عمل":"نظام العمل",
    "ساعات العمل":"نظام العمل","كم ساعة":"نظام العمل",
    "إجازة سنوية":"نظام العمل","إجازة أمومة":"نظام العمل",
    "فصل تعسفي":"نظام العمل","تم فصلي":"نظام العمل",
    "فُصلت":"نظام العمل","موظف":"نظام العمل","عامل":"نظام العمل",
    "مكافأة":"نظام العمل","اشتغلت":"نظام العمل","عملت":"نظام العمل",
    "فصلوني":"نظام العمل","طردوني":"نظام العمل","اتطردت":"نظام العمل",
    "راتب":"نظام العمل","أجر":"نظام العمل",
    "حقوقي":"نظام العمل","مستحقاتي":"نظام العمل",
    "فلوس":"نظام العمل","شهادة خبرة":"نظام العمل",
    "عمل إضافي":"نظام العمل","ساعات إضافية":"نظام العمل",
    # الجرائم المعلوماتية
    "جرائم معلوماتية":"نظام مكافحة الجرائم المعلوماتية",
    "اختراق":"نظام مكافحة الجرائم المعلوماتية",
    "قرصنة":"نظام مكافحة الجرائم المعلوماتية",
    # حماية البيانات
    "بيانات شخصية":"نظام حماية البيانات الشخصية",
    "حماية البيانات":"نظام حماية البيانات الشخصية",
    # المرافعات
    "مرافعات":"نظام المرافعات الشرعية",
    "دعوى":"نظام المرافعات الشرعية",
}

LEGAL_KEYWORDS = {
    "نظام","قانون","مادة","عقوبة","غرامة","سجن","محكمة","قاضي",
    "دعوى","متهم","عقد","زواج","طلاق","حضانة","إفلاس","تحكيم",
    "توثيق","محامي","عقار","إجراء","شرط","حق","رخصة","جريمة",
    "موظف","عامل","مكافأة","تعويض","إجازة","أجر","فصل","بيانات",
    "راتب","حقوق","مستحقات","ترخيص","إثبات","تنفيذ","لائحة",
    "موقوف","توقيف","محتجز","احتجاز",
}

NON_LEGAL = {
    "كورة","فريق","لاعب","مباراة","دوري","طبخ","أكل","وصفة",
    "فيلم","مسلسل","ممثل","غنية","مطرب","طقس","حرارة",
    "سعر الذهب","سعر الدولار","بورصة","سهم","عملة",
}

PRACTICAL_PATTERNS = {
    "أنا موظف","أنا عامل","اشتغلت","فصلوني","هل لي","هل يحق",
    "هل أستحق","صاحب العمل","طردوني","حقي","حقوقي","مستحقاتي",
    "تم فصلي","فُصلت","عملت","هل ليا","ليا فلوس",
    "اتطردت","كم ساعة","اتهمت","اتهمني","قبضوا علي",
    "موقوف","وقفوني","حقوق الموقوف",
}

QUERY_EXPANSION = {
    "مكافأة نهاية الخدمة": "يستحق العامل مكافأة عن مدة خدمته نصف شهر سنة",
    "فصلوني":              "إنهاء عقد العمل تعسفي تعويض فصل",
    "اتطردت":              "تم فصلي تعويض فصل تعسفي مكافأة نهاية الخدمة",
    "حقوقي":               "يستحق العامل مكافأة تعويض نظام العمل",
    "ساعات العمل":         "لا يجوز تشغيل العامل أكثر من ثماني ساعات يومياً",
    "كم ساعة":             "لا يجوز تشغيل العامل أكثر من ثماني ساعات",
    "شهادة خبرة":          "يلتزم صاحب العمل بإعطاء شهادة خبرة عند انتهاء العقد",
    "اتهمت":               "حق المتهم في محامٍ أثناء التحقيق الجزائي",
    "قبضوا علي":           "حقوق الموقوف التوقيف الإجراءات الجزائية",
    "حقوق الموقوف":        "يحق للموقوف الاتصال بمحاميه وذويه عند التوقيف",
    "موقوف":               "للموقوف حق الاتصال بذويه ومحاميه الإجراءات الجزائية",
    "شروط مزاولة المحاماة":"يشترط فيمن يزاول المحاماة سعودي جدول المحامين",
    "أركان العقد":         "أركان العقد الإيجاب والقبول والمحل",
    "عقوبات غسل الأموال":  "يعاقب على غسل الأموال سجن غرامة",
    "أحكام الطلاق":        "الطلاق رجعي بائن حل عقد الزواج",
}

LEGAL_TERMS_VOCAB = list(set(
    list(LAW_KEYWORDS.keys()) + list(COLLOQUIAL.keys()) +
    list(LEGAL_KEYWORDS) +
    ["مكافأة","نهاية","الخدمة","تعويض","فصل","تعسفي",
     "محكمة","عقوبة","أموال","عقار","صاحب","أجر","حقوق",
     "محاماة","مزاولة","رخصة","مهنة","موثق","قضاء","قاضي",
     "إجراءات","جزائية","متهم","توقيف","جريمة","موقوف"]
))

FUZZY_WHITELIST = {
    "شروط","حقوق","تعويض","مكافأة","إجازة","ساعات",
    "موظف","عامل","محامي","قاضي","زواج","طلاق","حضانة",
    "نظام","قانون","مادة","غرامة","سجن","مستحقات","راتب",
    "موقوف","محتجز","متهم","جريمة","توقيف",
}

STOP_WORDS = {
    "من","في","على","إلى","عن","مع","هي","هو","ما","لا","أن","إن",
    "كل","هذا","ذلك","التي","الذي","هذه","تلك","قد","إذا","إذ",
}

# ══════════════════════════════════════════════════════════
# Rate Limiter
# ══════════════════════════════════════════════════════════
class RateLimiter:
    def __init__(self, rpm=28):
        self._q = deque(); self._rpm = rpm
    def _clean(self):
        now = time.time()
        while self._q and now - self._q[0] > 60: self._q.popleft()
    def can_use(self): self._clean(); return len(self._q) < self._rpm
    def record(self):  self._q.append(time.time())
    def wait_secs(self):
        self._clean()
        if len(self._q) < self._rpm: return 0.0
        return max(0.0, 60 - (time.time() - self._q[0]))

groq_limiter = RateLimiter(rpm=28)

# ══════════════════════════════════════════════════════════
# Startup
# ══════════════════════════════════════════════════════════
@asynccontextmanager
async def lifespan(app: FastAPI):
    await startup()
    yield

async def startup():
    global vectorstore, bm25_index, bm25_texts, bm25_metas
    global embeddings, ACTIVE_MODELS, GEMINI_AVAILABLE, gemini_clients, groq_client

    print("🚀 Starting Saudi Legal AI v7...")

    # ── Gemini ────────────────────────────────────────────
    if GEMINI_KEY:
        genai.configure(api_key=GEMINI_KEY)
        for mn in GEMINI_PRIORITY:
            try:
                client = genai.GenerativeModel(mn)
                client.generate_content("hi", generation_config={"max_output_tokens": 3})
                gemini_clients[mn] = client
                GEMINI_AVAILABLE   = True
                print(f"  ✅ Gemini: {mn.split('/')[-1]}")
            except Exception as e:
                print(f"  ❌ Gemini: {mn.split('/')[-1]} — {str(e)[:40]}")

    # ── Groq ──────────────────────────────────────────────
    groq_client = Groq(api_key=GROQ_API_KEY)
    for m in [
        {"model": "llama-3.3-70b-versatile", "name": "Llama-3.3-70B"},
        {"model": "llama-3.1-8b-instant",    "name": "Llama-3.1-8B"},
    ]:
        try:
            groq_client.chat.completions.create(
                model=m["model"], messages=[{"role":"user","content":"hi"}],
                max_tokens=3, timeout=10)
            ACTIVE_MODELS.append(m)
            print(f"  ✅ Groq: {m['name']}")
        except Exception as e:
            print(f"  ❌ Groq: {m['name']} — {str(e)[:30]}")

    # ── Dataset ───────────────────────────────────────────
    print("📥 Loading dataset...")
    if HF_TOKEN: login(token=HF_TOKEN, add_to_git_credential=False)

    try:
        ds = load_dataset(HF_REPO_ID, token=HF_TOKEN, split="train")
        hf_items = [
            {"text":i.get("text",""),"article_number":i.get("article_number",""),
             "law_name":i.get("law_name",""),"law_type":i.get("law_type","نظام"),
             "source":"huggingface","_src":"hf","is_exec_reg":
             1 if any(p in i.get("law_name","") for p in ["اللائحة التنفيذية","اللوائح التنفيذية"]) else 0}
            for i in ds if len(i.get("text","")) > 30
        ]
        print(f"  ✅ HF: {len(hf_items)} articles")
    except Exception as e:
        print(f"  ⚠️ HF load failed: {e}")
        hf_items = []

    # Add manual supplements
    seen_keys = {re.sub(r"\s+","",i["text"])[:80] for i in hf_items}
    manual_added = 0
    for item in MANUAL:
        key = re.sub(r"\s+","",item["text"])[:80]
        if key not in seen_keys:
            seen_keys.add(key)
            hf_items.append({**item, "_src":"manual", "is_exec_reg":0})
            manual_added += 1

    print(f"  ✅ Manual: {manual_added} supplements added")
    print(f"  ✅ Total:  {len(hf_items)} articles")

    # ── Embeddings + ChromaDB ─────────────────────────────
    print("🔄 Building vector store...")
    embeddings = SentenceTransformerEmbeddings(
        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

    docs_lc = [
        Document(
            page_content=item["text"],
            metadata={
                "article_number": item.get("article_number",""),
                "law_name":       item.get("law_name",""),
                "law_type":       item.get("law_type","نظام"),
                "source":         item.get("source",""),
                "_src":           item.get("_src","hf"),
                "is_exec_reg":    item.get("is_exec_reg",0),
            }
        )
        for item in hf_items
    ]

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=1200, chunk_overlap=150,
        separators=["\nالمادة ","\n\n","\n","."," "])
    chunks = [d for d in docs_lc if len(d.page_content) <= 1200]
    chunks += splitter.split_documents([d for d in docs_lc if len(d.page_content) > 1200])

    chroma_client = chromadb.Client()
    vectorstore   = Chroma.from_documents(
        documents=chunks, embedding=embeddings,
        client=chroma_client, collection_name="saudi_legal_v7")
    print(f"  ✅ ChromaDB: {vectorstore._collection.count()} vectors")

    # ── BM25 ──────────────────────────────────────────────
    raw        = vectorstore.get()
    bm25_texts = raw["documents"]
    bm25_metas = raw["metadatas"]
    bm25_index = BM25Okapi([_tokenize(t) for t in bm25_texts])
    print(f"  ✅ BM25:    {len(bm25_texts)} docs")
    print("✅ Saudi Legal AI v7 Ready!")

# ══════════════════════════════════════════════════════════
# Pipeline Functions
# ══════════════════════════════════════════════════════════

def _tokenize(text): return [w for w in text.split() if len(w)>2 and w not in STOP_WORDS]
def _eng_ratio(text):
    alpha = [c for c in text if c.isalpha()]
    return sum(1 for c in alpha if c.isascii()) / max(len(alpha),1) if alpha else 0.0

def _dict_translate(q):
    lower = q.lower()
    for eng, ar in sorted(ENGLISH_TO_ARABIC.items(), key=lambda x: -len(x[0])):
        if eng.lower() in lower:
            q = re.sub(re.escape(eng), ar, q, flags=re.IGNORECASE); lower = q.lower()
    return q

def _llm_translate(q):
    if q in _trans_cache: return _trans_cache[q]
    if GEMINI_AVAILABLE:
        try:
            client = next(iter(gemini_clients.values()))
            r = client.generate_content(f"ترجم للعربية القانونية. أرجع الترجمة فقط:\n{q}")
            t = (r.text or "").strip()
            if sum(1 for c in t if "\u0600"<=c<="\u06ff") / max(len(t),1) > 0.5:
                _trans_cache[q] = t; return t
        except: pass
    if ACTIVE_MODELS and groq_limiter.can_use():
        try:
            r = groq_client.chat.completions.create(
                model=ACTIVE_MODELS[-1]["model"], max_tokens=100, temperature=0,
                messages=[{"role":"system","content":"ترجم للعربية القانونية. أرجع الترجمة فقط."},
                          {"role":"user","content":q}])
            t = r.choices[0].message.content.strip()
            if sum(1 for c in t if "\u0600"<=c<="\u06ff") / max(len(t),1) > 0.5:
                _trans_cache[q] = t; return t
        except: pass
    return _dict_translate(q)

def _fuzzy_correct(q):
    words, result, changed = q.split(), [], False
    for word in words:
        if len(word)<5 or word in FUZZY_WHITELIST or word.isdigit() or \
           not any("\u0600"<=c<="\u06ff" for c in word):
            result.append(word); continue
        match = fuzz_proc.extractOne(word, LEGAL_TERMS_VOCAB, scorer=fuzz.ratio, score_cutoff=82)
        if match and match[0] != word:
            result.append(match[0]); changed = True
        else:
            result.append(word)
    return " ".join(result)

def normalize_question(question):
    question = " ".join(question.split())
    r = _eng_ratio(question)
    if r > 0.40:
        t = _llm_translate(question)
        t = _dict_translate(t)
        if t != question: question = t
    elif r > 0.10:
        question = _dict_translate(question)
    question = _fuzzy_correct(question)
    for dial, formal in COLLOQUIAL.items():
        question = re.sub(rf"\b{dial}\b", formal, question, flags=re.IGNORECASE)
    question = re.sub(r"[؟?]+", "؟", question).strip()
    if not question.endswith("؟"): question += "؟"
    return question

def detect_law(question):
    # أطول keyword أولاً — 'حقوق الموقوف' يطغى على 'حقوقي'
    for kw, law in sorted(LAW_KEYWORDS.items(), key=lambda x: -len(x[0])):
        if kw in question: return law
    return None

def is_legal(question):
    if question in _legal_cache: return _legal_cache[question]
    if any(kw in question for kw in NON_LEGAL):
        _legal_cache[question] = False; return False
    if any(p in question for p in PRACTICAL_PATTERNS):
        _legal_cache[question] = True; return True
    if any(kw in question for kw in LEGAL_KEYWORDS | set(LAW_KEYWORDS.keys())):
        _legal_cache[question] = True; return True
    if GEMINI_AVAILABLE:
        try:
            client = next(iter(gemini_clients.values()))
            r = client.generate_content(
                f"هل السؤال متعلق بالقانون السعودي أو حقوق العمل؟ YES أو NO فقط:\n{question}")
            result = "YES" in (r.text or "").upper()
            _legal_cache[question] = result; return result
        except: pass
    _legal_cache[question] = False; return False

def expand_query(question):
    queries = [question]
    for pat, exp in QUERY_EXPANSION.items():
        if pat in question: queries.append(exp); break
    stop  = {"ما","هي","هو","هل","في","على","من","إلى","عن","لي","يجب","أريد","متى"}
    words = [w for w in question.split() if len(w)>3 and w not in stop]
    if words: queries.append(" ".join(words[:5]))
    return queries

def rewrite_query(question):
    if question in _rewrite_cache: return _rewrite_cache[question]
    triggers = {"فصلوني","طردوني","تم فصلي","اتطردت","حقوقي","اتهمت","قبضوا","موقوف","هل يحق","اشتغلت"}
    if not any(t in question for t in triggers): return expand_query(question)
    if GEMINI_AVAILABLE:
        try:
            client = next(iter(gemini_clients.values()))
            r = client.generate_content(
                f"أعد صياغة السؤال القانوني بـ 3 طرق بالمصطلحات الرسمية السعودية.\n"
                f"أرجع 3 أسئلة فقط، كل في سطر، بلا ترقيم:\n{question}")
            lines = [l.strip() for l in (r.text or "").strip().split("\n")
                     if l.strip() and len(l.strip()) > 10][:3]
            if lines:
                result = [question] + lines
                _rewrite_cache[question] = result; return result
        except: pass
    if ACTIVE_MODELS:
        try:
            r = groq_client.chat.completions.create(
                model=ACTIVE_MODELS[0]["model"], max_tokens=200, temperature=0.2,
                messages=[
                    {"role":"system","content":"أعد صياغة السؤال القانوني بـ 3 طرق. أرجع 3 أسئلة مفصولة بسطر جديد."},
                    {"role":"user","content":question}])
            lines = [l.strip() for l in r.choices[0].message.content.strip().split("\n")
                     if l.strip() and len(l.strip()) > 10][:3]
            result = [question] + lines
            _rewrite_cache[question] = result; return result
        except: pass
    return expand_query(question)

def bm25_search(query, k=5, target_law=None):
    scores = bm25_index.get_scores(_tokenize(query))
    docs   = []
    for idx in scores.argsort()[::-1]:
        if len(docs) >= k or scores[idx] < 0.05: break
        meta = bm25_metas[idx]
        if target_law and meta.get("law_name") != target_law: continue
        docs.append(Document(page_content=bm25_texts[idx], metadata=meta))
    return docs

def rerank_docs(docs, question, target_law=None):
    qwords = [w for w in question.split() if len(w) > 2]
    scored = []
    for doc in docs:
        s = 0
        if target_law and doc.metadata.get("law_name") == target_law: s += 10
        # اللائحة التنفيذية أقل أولوية من النظام الأصلي
        if doc.metadata.get("is_exec_reg") == 1: s -= 2
        s += sum(2 for w in qwords if w in doc.page_content)
        if "المادة" in doc.metadata.get("article_number", ""): s += 3
        s += min(len(doc.page_content) // 200, 3)
        if doc.metadata.get("_src") == "manual": s += 2
        scored.append((s, doc))
    return [d for _, d in sorted(scored, key=lambda x: x[0], reverse=True)]

def calc_coverage(question, docs):
    if not docs: return 0.0
    laws  = {d.metadata.get("law_name","") for d in docs}
    base  = 0.5 if any(law in laws for law in LAW_KEYWORDS.values()) else 0.0
    words = [w for w in question.split() if len(w) > 3]
    if not words: return base
    text  = " ".join(d.page_content for d in docs)
    m     = sum(1 for w in words if w in text or (len(w)>=4 and w[:4] in text))
    return max(base, m / len(words))

def build_context(docs):
    parts = []
    for i, doc in enumerate(docs):
        label = "الأكثر صلة" if i == 0 else f"مرجع {i+1}"
        law   = doc.metadata.get("law_name","")
        art   = doc.metadata.get("article_number","")
        parts.append(f"[{label}] {law} — {art}\n{doc.page_content}\n{'─'*40}")
    return "\n\n".join(parts)

def post_process(answer, docs):
    answer = answer.strip()
    if docs and "📋" not in answer and "لم أجد" not in answer:
        law = docs[0].metadata.get("law_name","")
        art = docs[0].metadata.get("article_number","")
        if law and art: answer = f"📋 المرجع: {law} — {art}\n\n{answer}"
    lines = answer.split("\n")
    clean = [l for l in lines if
             sum(1 for c in l if "\u0600"<=c<="\u06ff") / max(len(l.replace(" ","")),1) > 0.25
             or any(s in l for s in ["📋","✅","📝","⚠️","•","-","─"])]
    return "\n".join(clean).strip()

# ══════════════════════════════════════════════════════════
# System Prompts
# ══════════════════════════════════════════════════════════
SYSTEM_PROMPT = """أنت مساعد قانوني متخصص في الأنظمة السعودية.

نطاقك: وزارة العدل + نظام العمل + الإجراءات الجزائية + الجرائم المعلوماتية + حماية البيانات.

صيغة الإجابة الإلزامية:
📋 المرجع: [اسم النظام] — [رقم المادة]
✅ الإجابة: [نعم/لا للأسئلة العملية — إجابة مباشرة]
📝 التفاصيل: [شرح من النص القانوني]

قواعد:
١. العربية الفصحى حصراً
٢. اذكر النظام والمادة دائماً
٣. لا تخترع معلومات — استند للنصوص فقط
٤. ابدأ بـ 📋 مباشرة
⚠️ للاستئناس فقط وليست استشارة قانونية معتمدة."""

OUT_OF_SCOPE = """أنا مساعد قانوني متخصص في الأنظمة السعودية.

يمكنني مساعدتك في:
• نظام العمل (مكافأة نهاية الخدمة، ساعات العمل، الفصل...)
• أنظمة وزارة العدل (محاماة، توثيق، أحوال شخصية، إفلاس...)
• نظام الإجراءات الجزائية (حقوق المتهم، التوقيف...)
• نظام مكافحة الجرائم المعلوماتية
• نظام حماية البيانات الشخصية

⚠️ سؤالك خارج نطاق اختصاصي."""

# ══════════════════════════════════════════════════════════
# Generation
# ══════════════════════════════════════════════════════════
def _call_gemini(messages):
    if not GEMINI_AVAILABLE: return None, None
    sys_msg  = next((m["content"] for m in messages if m["role"]=="system"),"")
    user_msg = next((m["content"] for m in messages if m["role"]=="user"),"")
    prompt   = f"{sys_msg}\n\n{user_msg}"
    for model_name, client in gemini_clients.items():
        try:
            r = client.generate_content(
                prompt,
                safety_settings=GEMINI_SAFETY,
                generation_config={"max_output_tokens":1200,"temperature":0.1})
            answer = (r.text or "").strip()
            if not answer: continue
            ar = sum(1 for c in answer if "\u0600"<=c<="\u06ff")
            if ar / max(len(answer.replace(" ","")),1) < 0.35: continue
            return answer, model_name.split("/")[-1]
        except Exception as e:
            if "429" in str(e) or "quota" in str(e).lower(): continue
    return None, None

def _call_groq(messages):
    if not groq_limiter.can_use(): return None, None
    for m in sorted(ACTIVE_MODELS, key=lambda x: 0 if "70b" in x["model"] else 1):
        try:
            r = groq_client.chat.completions.create(
                model=m["model"], max_tokens=1000, temperature=0.1, messages=messages)
            answer = r.choices[0].message.content
            ar = sum(1 for c in answer if "\u0600"<=c<="\u06ff")
            if ar / max(len(answer.replace(" ","")),1) < 0.45: continue
            groq_limiter.record(); return answer, m["name"]
        except Exception as e:
            if "429" in str(e): continue
    return None, None

def generate_answer(messages):
    ans, model = _call_gemini(messages)
    if ans: return ans, model
    ans, model = _call_groq(messages)
    if ans: return ans, model
    wait = groq_limiter.wait_secs()
    if 0 < wait < 30:
        time.sleep(wait + 1)
        return generate_answer(messages)
    return None, None

# ══════════════════════════════════════════════════════════
# Main ask()
# ══════════════════════════════════════════════════════════
def ask_legal(question: str) -> dict:
    cache_key = hashlib.md5(question.strip().encode()).hexdigest()
    if cache_key in _answer_cache:
        return {**_answer_cache[cache_key], "cache_hit": True}

    original = question
    question = normalize_question(question)

    if not is_legal(question):
        return {"answer":OUT_OF_SCOPE,"sources":[],"coverage":0,"model":"out_of_scope","normalized":question}

    queries    = rewrite_query(question)
    target_law = detect_law(question) or detect_law(original)

    if not target_law and any(w in question or w in original for w in
       {"فصل","تعويض","مكافأة","راتب","أجر","ساعات","إجازة","موظف","عامل","اتطردت","فلوس"}):
        target_law = "نظام العمل"

    # Hybrid retrieval
    seen, result_docs = set(), []
    def add_docs(new_docs):
        for d in new_docs:
            key = d.page_content[:60]
            if key not in seen: seen.add(key); result_docs.append(d)

    for q in queries:
        filt  = {"law_name": target_law} if target_law else None
        vdocs = vectorstore.similarity_search(q, k=4, filter=filt) if filt \
                else vectorstore.similarity_search(q, k=4)
        add_docs(vdocs)
        add_docs(bm25_search(q, k=4, target_law=target_law))

    # Keyword scan
    if target_law and sum(1 for d in result_docs if d.metadata.get("law_name")==target_law) < 2:
        raw = vectorstore.get(where={"law_name": target_law})
        qw  = set(w for q in queries for w in q.split() if len(w) > 2)
        scored = sorted(
            [(sum(1 for w in qw if w in dt), dt, dm)
             for dt, dm in zip(raw["documents"], raw["metadatas"])
             if sum(1 for w in qw if w in dt) >= 1], reverse=True)
        add_docs([Document(page_content=dt, metadata=dm) for _,dt,dm in scored[:5]])

    final = rerank_docs(result_docs, question, target_law)[:8]
    cov   = calc_coverage(question, final)

    if cov < 0.20 or not final:
        return {"answer":OUT_OF_SCOPE,"sources":[],"coverage":0,"model":"low_coverage","normalized":question}

    context  = build_context(final)
    messages = [
        {"role":"system","content":SYSTEM_PROMPT},
        {"role":"user",  "content":f"المواد القانونية:\n\n{context}\n\nالسؤال: {question}"},
    ]

    ans, model_used = generate_answer(messages)
    if not ans:
        return {"answer":"الخدمة مشغولة — يرجى المحاولة مجدداً.",
                "sources":[],"coverage":0,"model":"unavailable","normalized":question}

    result = {
        "answer":    post_process(ans, final),
        "sources":   [{"law":d.metadata.get("law_name",""),"article":d.metadata.get("article_number","")} for d in final[:3]],
        "coverage":  round(cov*100),
        "model":     model_used,
        "normalized":question,
    }
    if result["model"] not in {"out_of_scope","low_coverage","unavailable",""}:
        _answer_cache[cache_key] = result
    return result

# ══════════════════════════════════════════════════════════
# Frontend HTML
# ══════════════════════════════════════════════════════════
FRONTEND_HTML = r"""<!DOCTYPE html>
<html lang="ar" dir="rtl">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>المساعد القانوني السعودي</title>
<link href="https://fonts.googleapis.com/css2?family=Noto+Kufi+Arabic:wght@300;400;500;700;900&family=Cormorant+Garamond:wght@400;600&display=swap" rel="stylesheet">
<style>
  :root {
    --ink:     #12100E;
    --parch:   #F5F0E8;
    --gold:    #B8860B;
    --gold-lt: #D4A843;
    --green:   #1B4332;
    --cream:   #FAF6EE;
    --border:  #D4C9A8;
    --muted:   #7A7060;
    --shadow:  rgba(18,16,14,.12);
  }

  *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }

  body {
    font-family: 'Noto Kufi Arabic', sans-serif;
    background: var(--parch);
    color: var(--ink);
    min-height: 100vh;
    display: flex;
    flex-direction: column;
    background-image:
      radial-gradient(ellipse at 20% 10%, rgba(184,134,11,.06) 0%, transparent 50%),
      radial-gradient(ellipse at 80% 90%, rgba(27,67,50,.05) 0%, transparent 50%);
  }

  /* ── Header ── */
  header {
    background: var(--green);
    padding: 0 2rem;
    display: flex;
    align-items: center;
    justify-content: space-between;
    height: 64px;
    box-shadow: 0 2px 16px var(--shadow);
    position: sticky;
    top: 0;
    z-index: 100;
  }
  .header-brand {
    display: flex;
    align-items: center;
    gap: .85rem;
  }
  .header-emblem {
    width: 38px; height: 38px;
    border: 2px solid var(--gold);
    border-radius: 50%;
    display: grid; place-items: center;
    font-size: 1.1rem;
    color: var(--gold);
    flex-shrink: 0;
  }
  .header-title {
    font-size: 1rem;
    font-weight: 700;
    color: var(--parch);
    letter-spacing: .01em;
  }
  .header-sub {
    font-size: .68rem;
    color: rgba(245,240,232,.55);
    margin-top: 1px;
  }
  .header-badge {
    font-size: .65rem;
    background: rgba(184,134,11,.2);
    border: 1px solid rgba(184,134,11,.4);
    color: var(--gold-lt);
    padding: 3px 10px;
    border-radius: 20px;
    letter-spacing: .03em;
  }

  /* ── Layout ── */
  .layout { display: flex; flex: 1; overflow: hidden; }

  /* ── Sidebar ── */
  .sidebar {
    width: 260px;
    background: var(--cream);
    border-left: 1px solid var(--border);
    display: flex;
    flex-direction: column;
    overflow-y: auto;
    flex-shrink: 0;
  }
  .sidebar-section {
    padding: 1.2rem 1rem;
    border-bottom: 1px solid var(--border);
  }
  .sidebar-label {
    font-size: .65rem;
    font-weight: 700;
    color: var(--muted);
    text-transform: uppercase;
    letter-spacing: .1em;
    margin-bottom: .75rem;
  }
  .sidebar-chip {
    display: block;
    width: 100%;
    text-align: right;
    background: transparent;
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: .55rem .8rem;
    font-family: 'Noto Kufi Arabic', sans-serif;
    font-size: .78rem;
    color: var(--ink);
    cursor: pointer;
    margin-bottom: .45rem;
    transition: all .18s;
    line-height: 1.4;
  }
  .sidebar-chip:hover {
    background: var(--green);
    color: var(--parch);
    border-color: var(--green);
  }
  .stat-row {
    display: flex;
    justify-content: space-between;
    font-size: .72rem;
    color: var(--muted);
    padding: .25rem 0;
    border-bottom: 1px solid rgba(212,201,168,.4);
  }
  .stat-val { font-weight: 700; color: var(--gold); }

  /* ── Chat area ── */
  .chat-wrap {
    flex: 1;
    display: flex;
    flex-direction: column;
    overflow: hidden;
    max-width: 760px;
    margin: 0 auto;
    width: 100%;
    padding: 0 1rem;
  }
  #chat {
    flex: 1;
    overflow-y: auto;
    padding: 1.5rem 0;
    display: flex;
    flex-direction: column;
    gap: 1.25rem;
    scrollbar-width: thin;
    scrollbar-color: var(--border) transparent;
  }

  /* ── Messages ── */
  .msg { display: flex; gap: .75rem; animation: rise .3s ease; }
  .msg.user  { flex-direction: row-reverse; }
  .msg.bot   { }

  @keyframes rise { from{opacity:0;transform:translateY(10px)} to{opacity:1;transform:none} }

  .avatar {
    width: 36px; height: 36px;
    border-radius: 50%;
    display: grid; place-items: center;
    font-size: .95rem;
    flex-shrink: 0;
    margin-top: 2px;
  }
  .msg.bot  .avatar { background: var(--green); color: var(--gold); border: 1.5px solid var(--gold); }
  .msg.user .avatar { background: var(--gold);  color: var(--green); }

  .bubble {
    max-width: 82%;
    padding: .85rem 1.1rem;
    border-radius: 16px;
    font-size: .88rem;
    line-height: 1.75;
  }
  .msg.bot .bubble {
    background: white;
    border: 1px solid var(--border);
    border-top-right-radius: 4px;
    box-shadow: 0 2px 8px var(--shadow);
  }
  .msg.user .bubble {
    background: var(--green);
    color: var(--parch);
    border-top-left-radius: 4px;
    box-shadow: 0 2px 8px rgba(27,67,50,.25);
  }

  .ref-line  { color: var(--gold); font-weight: 700; font-size: .78rem; margin-bottom: .5rem; }
  .ans-line  { margin: .3rem 0; }
  .meta-line { font-size: .65rem; color: var(--muted); margin-top: .6rem; border-top: 1px solid var(--border); padding-top: .5rem; }

  /* ── Sources ── */
  .sources {
    margin-top: .6rem;
    display: flex;
    flex-wrap: wrap;
    gap: .3rem;
  }
  .source-tag {
    font-size: .62rem;
    background: rgba(27,67,50,.06);
    border: 1px solid var(--border);
    color: var(--green);
    padding: 2px 8px;
    border-radius: 20px;
  }

  /* ── Loading dots ── */
  .dots span {
    display: inline-block;
    width: 7px; height: 7px;
    background: var(--gold);
    border-radius: 50%;
    margin: 0 2px;
    animation: bounce 1.1s infinite;
  }
  .dots span:nth-child(2) { animation-delay: .18s; }
  .dots span:nth-child(3) { animation-delay: .36s; }
  @keyframes bounce { 0%,80%,100%{transform:translateY(0)} 40%{transform:translateY(-7px)} }

  /* ── Input bar ── */
  .input-wrap {
    border-top: 1px solid var(--border);
    padding: .9rem 0;
    background: var(--parch);
  }
  .input-row {
    display: flex;
    gap: .65rem;
    align-items: flex-end;
  }
  #question {
    flex: 1;
    background: white;
    border: 1.5px solid var(--border);
    border-radius: 12px;
    padding: .75rem 1rem;
    font-family: 'Noto Kufi Arabic', sans-serif;
    font-size: .9rem;
    color: var(--ink);
    resize: none;
    outline: none;
    min-height: 48px;
    max-height: 130px;
    line-height: 1.6;
    transition: border-color .2s;
    box-shadow: inset 0 1px 4px var(--shadow);
  }
  #question:focus { border-color: var(--gold); }
  #question::placeholder { color: var(--border); }

  #send-btn {
    width: 48px; height: 48px;
    background: var(--green);
    border: none;
    border-radius: 12px;
    cursor: pointer;
    font-size: 1.15rem;
    color: var(--gold);
    flex-shrink: 0;
    transition: all .2s;
    box-shadow: 0 2px 8px rgba(27,67,50,.3);
    display: grid; place-items: center;
  }
  #send-btn:hover:not(:disabled) { background: #255c3f; transform: translateY(-1px); }
  #send-btn:disabled { opacity: .45; cursor: not-allowed; transform: none; }

  /* ── Disclaimer ── */
  .disclaimer {
    text-align: center;
    font-size: .63rem;
    color: var(--muted);
    padding: .35rem 0;
    letter-spacing: .01em;
  }

  /* ── Welcome card ── */
  .welcome {
    background: white;
    border: 1px solid var(--border);
    border-radius: 16px;
    padding: 1.5rem;
    text-align: center;
    box-shadow: 0 2px 12px var(--shadow);
  }
  .welcome-icon { font-size: 2.5rem; margin-bottom: .75rem; }
  .welcome h2  { font-size: 1.1rem; font-weight: 700; color: var(--green); margin-bottom: .4rem; }
  .welcome p   { font-size: .82rem; color: var(--muted); line-height: 1.65; }

  /* ── Responsive ── */
  @media (max-width: 700px) {
    .sidebar { display: none; }
    .chat-wrap { padding: 0 .75rem; }
    .bubble { max-width: 92%; }
  }
</style>
</head>
<body>

<header>
  <div class="header-brand">
    <div class="header-emblem">⚖</div>
    <div>
      <div class="header-title">المساعد القانوني السعودي</div>
      <div class="header-sub">أنظمة وزارة العدل · نظام العمل · الإجراءات الجزائية</div>
    </div>
  </div>
  <span class="header-badge" id="status-badge">جارٍ التحميل…</span>
</header>

<div class="layout">

  <!-- Sidebar -->
  <aside class="sidebar">
    <div class="sidebar-section">
      <div class="sidebar-label">أسئلة شائعة</div>
      <button class="sidebar-chip" onclick="send('ما هي شروط مزاولة مهنة المحاماة؟')">شروط المحاماة</button>
      <button class="sidebar-chip" onclick="send('كم ساعة العمل في اليوم نظاماً؟')">ساعات العمل</button>
      <button class="sidebar-chip" onclick="send('هل يحق للعامل تعويض إذا فُصل بدون سبب؟')">تعويض الفصل</button>
      <button class="sidebar-chip" onclick="send('كيف تحسب مكافأة نهاية الخدمة؟')">نهاية الخدمة</button>
      <button class="sidebar-chip" onclick="send('ما هي حقوق الموقوف أثناء التحقيق؟')">حقوق الموقوف</button>
      <button class="sidebar-chip" onclick="send('ما هي عقوبات غسل الأموال؟')">غسل الأموال</button>
      <button class="sidebar-chip" onclick="send('ما هي أحكام الطلاق في نظام الأحوال الشخصية؟')">أحكام الطلاق</button>
      <button class="sidebar-chip" onclick="send('ما هي إجراءات الإفلاس؟')">إجراءات الإفلاس</button>
    </div>
    <div class="sidebar-section">
      <div class="sidebar-label">الإحصائيات</div>
      <div class="stat-row"><span>إجمالي الأسئلة</span><span class="stat-val" id="stat-total">0</span></div>
      <div class="stat-row"><span>الأسئلة المجابة</span><span class="stat-val" id="stat-success">0</span></div>
      <div class="stat-row"><span>خارج النطاق</span><span class="stat-val" id="stat-blocked">0</span></div>
    </div>
    <div class="sidebar-section">
      <div class="sidebar-label">تنبيه قانوني</div>
      <p style="font-size:.7rem;color:var(--muted);line-height:1.65">
        المعلومات المقدمة للاستئناس فقط وليست استشارة قانونية معتمدة. يُنصح بمراجعة محامٍ مختص لأي إجراءات رسمية.
      </p>
    </div>
  </aside>

  <!-- Chat -->
  <div class="chat-wrap">
    <div id="chat">
      <div class="msg bot">
        <div class="avatar">⚖</div>
        <div class="bubble">
          <div class="welcome">
            <div class="welcome-icon">🏛️</div>
            <h2>مرحباً بك في المساعد القانوني السعودي</h2>
            <p>يمكنني الإجابة عن أسئلتك في أنظمة وزارة العدل، نظام العمل، الإجراءات الجزائية، وغيرها.<br>اكتب سؤالك أو اختر من الأسئلة الشائعة.</p>
          </div>
        </div>
      </div>
    </div>

    <div class="input-wrap">
      <div class="input-row">
        <textarea id="question" placeholder="اكتب سؤالك القانوني هنا…" rows="1"
          onkeydown="if(event.key==='Enter'&&!event.shiftKey){event.preventDefault();sendFromInput();}"></textarea>
        <button id="send-btn" onclick="sendFromInput()" title="إرسال">➤</button>
      </div>
      <div class="disclaimer">⚠️ للاستئناس فقط · ليست استشارة قانونية معتمدة · وزارة العدل السعودية</div>
    </div>
  </div>

</div>

<script>
const chat     = document.getElementById('chat');
const input    = document.getElementById('question');
const sendBtn  = document.getElementById('send-btn');
const badge    = document.getElementById('status-badge');
let totalQ=0, successQ=0, blockedQ=0;

// Check health on load
fetch('/health').then(r=>r.json()).then(d=>{
  badge.textContent = `${d.vectors?.toLocaleString()} مادة · مباشر`;
  badge.style.background='rgba(27,67,50,.3)';
  badge.style.borderColor='rgba(27,67,50,.6)';
  badge.style.color='#6fcf97';
}).catch(()=>{ badge.textContent='متصل'; });

function appendMsg(role, htmlContent) {
  const av   = role==='user' ? '👤' : '⚖';
  const wrap = document.createElement('div');
  wrap.className = `msg ${role}`;
  wrap.innerHTML = `<div class="avatar">${av}</div><div class="bubble">${htmlContent}</div>`;
  chat.appendChild(wrap);
  chat.scrollTop = chat.scrollHeight;
  return wrap;
}

function showLoader() { return appendMsg('bot','<div class="dots"><span></span><span></span><span></span></div>'); }

function formatAnswer(data) {
  let html = '';
  const lines = (data.answer||'').split('\n');
  lines.forEach(ln => {
    ln = ln.trim();
    if (!ln) return;
    if (ln.startsWith('📋'))      html += `<div class="ref-line">${ln}</div>`;
    else if (ln.startsWith('✅')) html += `<div class="ans-line"><strong>${ln}</strong></div>`;
    else                           html += `<div class="ans-line">${ln}</div>`;
  });
  // Sources
  const srcs = (data.sources||[]).filter(s=>s.law);
  if (srcs.length) {
    html += '<div class="sources">';
    srcs.forEach(s => {
      html += `<span class="source-tag">📚 ${s.law}</span>`;
    });
    html += '</div>';
  }
  if (data.model && !['out_of_scope','low_coverage','unavailable'].includes(data.model)) {
    html += `<div class="meta-line">🤖 ${data.model} · تغطية ${data.coverage}%</div>`;
  }
  return html;
}

async function send(q) {
  if (!q?.trim() || sendBtn.disabled) return;
  appendMsg('user', q.replace(/</g,'&lt;'));
  input.value = '';
  input.style.height = 'auto';
  sendBtn.disabled = true;
  const loader = showLoader();

  try {
    const res  = await fetch('/ask', {
      method: 'POST',
      headers: {'Content-Type':'application/json'},
      body: JSON.stringify({question:q, include_sources:true})
    });
    const data = await res.json();
    loader.remove();

    totalQ++;
    if (data.model&&['out_of_scope','low_coverage'].includes(data.model)) blockedQ++;
    else successQ++;
    document.getElementById('stat-total').textContent   = totalQ;
    document.getElementById('stat-success').textContent = successQ;
    document.getElementById('stat-blocked').textContent = blockedQ;

    if (res.ok) appendMsg('bot', formatAnswer(data));
    else appendMsg('bot', `<div style="color:#c0392b">⚠️ ${data.detail||'خطأ غير متوقع'}</div>`);
  } catch(e) {
    loader.remove();
    appendMsg('bot','<div style="color:#c0392b">❌ تعذر الاتصال بالخادم</div>');
  } finally {
    sendBtn.disabled = false;
    input.focus();
  }
}

function sendFromInput() { send(input.value.trim()); }

input.addEventListener('input', () => {
  input.style.height = 'auto';
  input.style.height = Math.min(input.scrollHeight, 130) + 'px';
});
</script>
</body>
</html>
"""

# ══════════════════════════════════════════════════════════
# FastAPI App
# ══════════════════════════════════════════════════════════
app = FastAPI(
    title="Saudi Legal AI",
    description="نظام الذكاء الاصطناعي للقانون السعودي v7",
    version="7.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"],
)

def _rate_limit(ip: str, max_rpm: int = 15) -> bool:
    now = time.time()
    user_ips[ip] = [t for t in user_ips[ip] if now - t < 60]
    if len(user_ips[ip]) >= max_rpm: return False
    user_ips[ip].append(now); return True

class QuestionRequest(BaseModel):
    question:        str
    include_sources: bool = True

class AnswerResponse(BaseModel):
    answer:      str
    model:       str
    coverage:    int
    sources:     list = []
    duration_ms: int  = 0
    disclaimer:  str  = "⚠️ للاستئناس فقط وليست استشارة قانونية معتمدة."

@app.get("/", response_class=HTMLResponse)
def root(): return FRONTEND_HTML

@app.get("/health")
def health():
    return {
        "status":  "healthy",
        "version": "7.0.0",
        "vectors": vectorstore._collection.count() if vectorstore else 0,
        "models":  {
            "primary":  [mn.split("/")[-1] for mn in gemini_clients],
            "fallback": [m["name"] for m in ACTIVE_MODELS],
        },
    }

@app.get("/stats")
def stats():
    return {
        "total":      api_stats["total"],
        "success":    api_stats["success"],
        "blocked":    api_stats["blocked"],
        "errors":     api_stats["errors"],
        "unique_ips": len(user_ips),
        "recent":     list(request_log)[:10],
    }

@app.get("/laws")
def get_laws():
    if not vectorstore: return {"laws":[], "total":0}
    names = sorted({m.get("law_name","") for m in vectorstore.get()["metadatas"] if m.get("law_name")})
    return {"laws": names, "total": len(names)}

@app.post("/ask", response_model=AnswerResponse)
async def ask_endpoint(req: QuestionRequest, request: Request):
    if not req.question.strip():
        raise HTTPException(400, "السؤال لا يمكن أن يكون فارغاً")
    if len(req.question) > 1000:
        raise HTTPException(400, "السؤال طويل جداً (الحد 1000 حرف)")

    ip = request.headers.get("X-Forwarded-For","unknown").split(",")[0].strip()
    if not _rate_limit(ip):
        raise HTTPException(429, "تجاوزت الحد — حاول بعد دقيقة")

    t0 = time.time()
    try:
        result = ask_legal(req.question.strip())
    except Exception as e:
        api_stats["errors"] += 1
        raise HTTPException(500, f"خطأ داخلي: {str(e)[:80]}")

    ms     = int((time.time() - t0) * 1000)
    status = "blocked" if result["model"] in {"out_of_scope","low_coverage"} else "success"
    api_stats["total"]  += 1
    api_stats[status]   += 1
    request_log.appendleft({
        "time": time.strftime("%H:%M:%S"),
        "ip":   ip,
        "q":    req.question[:60],
        "ms":   ms,
        "model":result.get("model",""),
    })

    return AnswerResponse(
        answer      = result["answer"],
        model       = result.get("model",""),
        coverage    = result.get("coverage",0),
        sources     = result.get("sources",[]) if req.include_sources else [],
        duration_ms = ms,
    )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)