"""
Saudi Legal AI — Production Version v8.2
Gemini-first · Groq fallback · FastAPI + HTML frontend
"""
import os
# ── إيقاف رسائل التتبع المزعجة الخاصة بـ ChromaDB ──
os.environ["ANONYMIZED_TELEMETRY"] = "False"

import 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
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)
_answer_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
# ══════════════════════════════════════════════════════════
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"},
]

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

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

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

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

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

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

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 v8.2...")

    # ── 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 (HF + train.json) ───────────────────────────────────────────
    print("📥 Loading dataset...")
    if HF_TOKEN: login(token=HF_TOKEN, add_to_git_credential=False)

    hf_items = []
    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}")

    try:
        # حل مشكلة المسار المطلق لملف train.json
        base_dir = os.path.dirname(os.path.abspath(__file__))
        json_path = os.path.join(base_dir, "train.json")
        
        with open(json_path, "r", encoding="utf-8") as f:
            train_raw = json.load(f)
            
        added_from_json = 0
        for record in train_raw:
            try:
                conv = record.get("conversations", [])
                json_turn = next((c for c in conv if "output" in c.get("value", "")), None)
                if not json_turn: continue
                
                val = json.loads(json_turn["value"])
                output = val.get("output", val)
                text = (output.get("full_text", "") or "").strip()
                subj = (output.get("subject", "") or "وزارة العدل").strip()
                
                if len(text) > 50:
                    law_name = subj.split("/")[0].strip()
                    law_name = re.sub(r"\(.*?\)", "", law_name).strip()
                    hf_items.append({
                        "text": text, # أخذ المادة كاملة دون القص
                        "article_number": "نص قانوني",
                        "law_name": law_name, "law_type": "نظام", "_src": "train.json",
                        "is_exec_reg": 1 if "لائحة" in law_name else 0
                    })
                    added_from_json += 1
            except: pass
        print(f"  ✅ train.json: {added_from_json} articles added!")
    except Exception:
        print("  ℹ️ No train.json found. Skipping.")

    # Deduplicate & Add Manual
    seen_keys = {re.sub(r"\s+","",i["text"])[:80] for i in hf_items}
    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})

    # ── 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={k:v for k,v in item.items() if k!="text"})
        for item in hf_items if len(item.get("text","")) > 30
    ]

    # تكبير الحجم لعدم قص الكلمات في منتصفها
    splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=300, separators=["\nالمادة ","\n\n","\n","."," "])
    chunks = [d for d in docs_lc if len(d.page_content) <= 2000]
    chunks += splitter.split_documents([d for d in docs_lc if len(d.page_content) > 2000])

    chroma_client = chromadb.Client()
    vectorstore   = Chroma.from_documents(documents=chunks, embedding=embeddings, client=chroma_client, collection_name="saudi_legal_v8")

    # ── BM25 ──────────────────────────────────────────────
    raw        = vectorstore.get()
    bm25_texts = raw["documents"]
    bm25_metas = raw["metadatas"]
    bm25_index = BM25Okapi([[w for w in t.split() if len(w)>2 and w not in STOP_WORDS] for t in bm25_texts])
    print("✅ Saudi Legal AI v8.2 Ready!")

# ══════════════════════════════════════════════════════════
# Pipeline Functions (FAST MODE)
# ══════════════════════════════════════════════════════════

def clean_arabic_text(text: str) -> str:
    """✅ Normalize Text: Removes Hamzas & fixes Taa Marbuta"""
    text = " ".join(text.split())
    text = re.sub(r"ه\b", "ة", text)
    text = re.sub(r'[أإآ]', 'ا', text)
    return text

def normalize_question(question):
    question = clean_arabic_text(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 is_legal(question):
    q_norm = clean_arabic_text(question)
    if any(kw in q_norm for kw in NON_LEGAL): return False
    legal_kws_norm = [clean_arabic_text(kw) for kw in LEGAL_KEYWORDS | set(LAW_KEYWORDS.keys())]
    if any(kw in q_norm for kw in legal_kws_norm): return True
    if any(p in q_norm for p in ["اعمل اي", "اي العمل", "ماذا افعل", "حقي", "كيف"]): return True
    return False

def detect_law(question):
    for kw, law in sorted(LAW_KEYWORDS.items(), key=lambda x: -len(x[0])):
        if kw in question: return law
    return None

def fast_rewrite(question):
    queries = [question]
    for pat, exp in QUERY_EXPANSION.items():
        if pat in question: queries.append(exp); break
    words = [w for w in question.split() if len(w)>3 and w not in STOP_WORDS]
    if words: queries.append(" ".join(words[:5]))
    return list(set(queries))

def bm25_search(query, k=5, target_law=None):
    scores = bm25_index.get_scores([w for w in query.split() if len(w)>2 and w not in STOP_WORDS])
    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 clean_arabic_text(question).split() if len(w) > 2]
    scored = []
    PRIMARY_LAW_PREFIXES = [
        "نظام الإجراءات الجزائية", "نظام المحاماة", "نظام العمل", "نظام التوثيق",
        "نظام الأحوال الشخصية", "نظام الإفلاس", "نظام التحكيم", "نظام الإثبات",
        "نظام التنفيذ", "نظام المرافعات", "نظام القضاء", "نظام مكافحة",
        "نظام المعاملات", "نظام المحاكم", "نظام حماية"
    ]
    
    for doc in docs:
        s = 0
        law_name = doc.metadata.get("law_name", "")
        art = doc.metadata.get("article_number", "")
        content_norm = clean_arabic_text(doc.page_content)
        
        if target_law:
            if law_name == target_law: s += 10
            elif law_name.startswith(target_law): s += 8
            
        if any(law_name.startswith(p) for p in PRIMARY_LAW_PREFIXES): s += 4
        if doc.metadata.get("is_exec_reg") == 1: s -= 2
            
        s += sum(2 for w in qwords if w in content_norm)
        if "المادة" in art and "فقرة" not in art: 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}"
    return answer # بدون أي حذف للسطور

# ══════════════════════════════════════════════════════════
# 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    = fast_rewrite(question) # Fast rewrite
    target_law = detect_law(question) or detect_law(original)

    # ✅ Smart Detective (Strict Clues)
    strict_labor_clues = {"مكافأة", "نهاية الخدمة", "راتب", "أجر", "ساعات", "موظف", "عامل", "فصلوني", "اتطردت", "رفدي", "اترفدت", "اتفصلت", "إجازة", "امومة", "أمومة", "شهادة خبرة", "الشغل", "استقالة", "هستقيل", "استقيل"}
    anti_labor_clues   = {"محاماة", "محامي", "رخصة", "إفلاس", "قتل", "جريمة", "موقوف", "طلاق", "زواج", "نفقة", "عقار", "دين", "سلف"}
    
    if not target_law:
        has_labor = any(w in question or w in original for w in strict_labor_clues)
        has_anti  = any(w in question or w in original for w in anti_labor_clues)
        has_investigation = "تحقيق" in question or "التحقيق" in question
        
        if has_labor and not has_anti:
            target_law = "نظام العمل"
        elif has_investigation and has_labor:
            target_law = "نظام العمل"
        elif has_investigation and not has_labor:
            target_law = "نظام الإجراءات الجزائية"
        elif has_anti and any(w in question for w in ["فلوس", "دين", "سلف"]):
            target_law = "نظام المعاملات المدنية"

    # 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 w in clean_arabic_text(question).split() if len(w) > 2)
        scored = sorted(
            [(sum(1 for w in qw if w in clean_arabic_text(dt)), dt, dm)
             for dt, dm in zip(raw["documents"], raw["metadatas"])
             if sum(1 for w in qw if w in clean_arabic_text(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 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}

    ans = post_process(ans, final)

    result = {
        "answer":    ans,
        "sources":   [{"law":d.metadata.get("law_name",""),"article":d.metadata.get("article_number","")} for d in final[:3]],
        "coverage":  round(cov*100) if target_law else round(cov*50),
        "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 (100% Original Preserved)
# ══════════════════════════════════════════════════════════
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">

  <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>

  <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="نظام الذكاء الاصطناعي للقانون السعودي v8.2",
    version="8.2.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": "8.2.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)