"""
Saudi Legal AI API — v4.1
FastAPI + RAG + Groq + Gemini Multi-Model
"""
import os, gc, re, time, logging
from collections import deque, defaultdict
from datetime import datetime
from contextlib import asynccontextmanager
from dotenv import load_dotenv
from typing import Dict

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

from groq import Groq
import google.generativeai as genai
from huggingface_hub import login
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_process
import chromadb

load_dotenv()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ══════════════════════════════════════════════════════════
# 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')
CHROMA_PATH   = os.getenv('CHROMA_PATH', './chroma_db')

# ══════════════════════════════════════════════════════════
# Global State
# ══════════════════════════════════════════════════════════
vectorstore        = None
bm25_index         = None
bm25_texts         = []
bm25_metadatas     = []
embeddings         = None
ACTIVE_MODELS      = []
GEMINI_AVAILABLE   = False
gemini_clients: Dict = {}
groq_client        = None
request_log        = deque(maxlen=500)
stats              = {'total': 0, 'success': 0, 'blocked': 0, 'errors': 0}
active_ips         = {}
_expansion_cache   = {}
_rewrite_cache     = {}
_translation_cache = {}

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

# ══════════════════════════════════════════════════════════
# Extra Laws
# ══════════════════════════════════════════════════════════
EXTRA_LAWS = [
    {'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'},
]

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

NUMBERS_AR = {
    '1':'واحد','2':'اثنين','3':'ثلاث','4':'أربع','5':'خمس',
    '6':'ست','7':'سبع','8':'ثماني','9':'تسع','10':'عشر',
}

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

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

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

PRACTICAL_PATTERNS = [
    'أنا موظف','أنا عامل','اشتغلت','فصلوني','هل لي',
    'هل يحق','هل أستحق','صاحب العمل',
    'طردوني','حقي','حقوقي','مستحقاتي',
    'تم فصلي','فُصلت','عملت',
]

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

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

LEGAL_VOCABULARY = list(set(
    list(LAW_KEYWORDS.keys()) + list(COLLOQUIAL.keys()) +
    LEGAL_KEYWORDS + list(FUZZY_WHITELIST) +
    ['مكافأة','نهاية','الخدمة','تعويض','فصل','تعسفي',
     'محكمة','عقوبة','إثبات','تنفيذ','مرافعة','أموال',
     'عقار','تسجيل','صاحب','شركة','أجر','حقوق']
))

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

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

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

⚠️ تنبيه: المعلومات للاستئناس فقط وليست استشارة قانونية معتمدة.

قواعد:
1. العربية الفصحى فقط
2. اذكر رقم المادة واسم النظام دائماً
3. لو السؤال عملي أجب مباشرة بناءً على النصوص
4. لا تخترع معلومات
5. ابدأ بـ 📋 مباشرة"""

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

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

⚠️ تنبيه: المعلومات للاستئناس فقط وليست استشارة قانونية معتمدة.

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

# ══════════════════════════════════════════════════════════
# Rate Limiter
# ══════════════════════════════════════════════════════════
class SmartRateLimiter:
    def __init__(self):
        self.requests  = {'groq': deque()}
        self.limits    = {'groq': 28}
        self.last_used = {'groq': 0}

    def _clean(self, c):
        now = time.time()
        while self.requests[c] and now - self.requests[c][0] > 60:
            self.requests[c].popleft()

    def can_use(self, c):
        self._clean(c)
        if time.time() - self.last_used[c] < 0.5: return False
        return len(self.requests[c]) < self.limits[c]

    def record(self, c):
        self.requests[c].append(time.time())
        self.last_used[c] = time.time()

    def wait_time(self, c):
        self._clean(c)
        if len(self.requests[c]) < self.limits[c]: return 0
        return max(0, 60 - (time.time() - self.requests[c][0]))

rate_limiter = SmartRateLimiter()

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

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

    logger.info("Starting Saudi Legal AI v4.1...")

    # ── Groq ─────────────────────────────────────────────
    groq_client = Groq(api_key=GROQ_API_KEY)
    for m in [
        {'model': 'llama-3.3-70b-versatile', 'name': 'Groq llama-3.3'},
        {'model': 'llama-3.1-8b-instant',    'name': 'Groq 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)
            logger.info(f"✅ {m['name']}")
        except Exception as e:
            logger.warning(f"❌ {m['name']}: {str(e)[:30]}")

    # ── Gemini Multi-Model ────────────────────────────────
    if GEMINI_KEY:
        genai.configure(api_key=GEMINI_KEY)
        for model_name in GEMINI_MODELS:
            try:
                gemini_clients[model_name] = genai.GenerativeModel(model_name)
                GEMINI_AVAILABLE = True
                logger.info(f"✅ {model_name}")
            except Exception as e:
                logger.warning(f"❌ {model_name}: {str(e)[:40]}")

    # ── Embeddings ────────────────────────────────────────
    logger.info("Loading embeddings...")
    embeddings = SentenceTransformerEmbeddings(
        model_name='sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')

    # ── ChromaDB ──────────────────────────────────────────
    chroma_client_persist = chromadb.PersistentClient(path=CHROMA_PATH)
    try:
        collection = chroma_client_persist.get_collection('saudi_legal_v4')
        if collection.count() > 100:
            logger.info(f"✅ ChromaDB loaded: {collection.count()} chunks")
            vectorstore = Chroma(
                client=chroma_client_persist,
                collection_name='saudi_legal_v4',
                embedding_function=embeddings)
        else:
            raise Exception("Empty")
    except:
        logger.info("Loading dataset from HuggingFace...")
        login(token=HF_TOKEN, add_to_git_credential=False)
        dataset = load_dataset(HF_REPO_ID, token=HF_TOKEN, split='train')
        logger.info(f"✅ {len(dataset)} articles")

        docs = [
            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',''),
                })
            for item in dataset if len(item.get('text','')) > 30
        ]
        for a in EXTRA_LAWS:
            docs.append(Document(page_content=a['text'], metadata={
                'article_number': a['article_number'], 'law_name': a['law_name'],
                'law_type': a['law_type'], 'source': a['source'],
            }))

        splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
        chunks   = splitter.split_documents(docs)
        vectorstore = Chroma.from_documents(
            documents=chunks, embedding=embeddings,
            client=chroma_client_persist, collection_name='saudi_legal_v4')
        logger.info(f"✅ ChromaDB created: {vectorstore._collection.count()} chunks")

    # ── BM25 ──────────────────────────────────────────────
    stop_words = {'من','في','على','إلى','عن','مع','هي','هو','ما','لا','أن','إن'}
    def tokenize(text): return [w for w in text.split() if len(w) > 2 and w not in stop_words]
    all_chunks     = vectorstore.get()
    bm25_texts     = all_chunks['documents']
    bm25_metadatas = all_chunks['metadatas']
    bm25_index     = BM25Okapi([tokenize(t) for t in bm25_texts])
    logger.info(f"✅ BM25: {len(bm25_texts)} docs")
    logger.info("✅ Saudi Legal AI v4.1 Ready!")

# ══════════════════════════════════════════════════════════
# Pipeline Functions
# ══════════════════════════════════════════════════════════
def tokenize_arabic(text):
    stop_words = {'من','في','على','إلى','عن','مع','هي','هو','ما','لا','أن','إن'}
    return [w for w in text.split() if len(w) > 2 and w not in stop_words]

def detect_language_ratio(text):
    alpha = [c for c in text if c.isalpha()]
    if not alpha: return 0.0
    return len([c for c in alpha if c.isascii()]) / len(alpha)

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

def translate_with_llm(question):
    if question in _translation_cache: return _translation_cache[question]
    try:
        r = groq_client.chat.completions.create(
            model=ACTIVE_MODELS[0]['model'], max_tokens=200, temperature=0,
            messages=[
                {'role':'system','content':'أنت مترجم متخصص في القانون السعودي. ترجم للعربية الفصحى القانونية فقط بدون شرح.'},
                {'role':'user','content':question}
            ])
        translated = r.choices[0].message.content.strip()
        arabic_ratio = sum(1 for c in translated if '\u0600' <= c <= '\u06ff') / max(len(translated), 1)
        if arabic_ratio > 0.5:
            _translation_cache[question] = translated
            return translated
    except: pass
    return translate_with_dict(question)

def smart_translate(question):
    english_ratio = detect_language_ratio(question)
    if english_ratio > 0.40:
        translated = translate_with_llm(question)
        translated = translate_with_dict(translated)
        return translated, 'LLM'
    elif english_ratio > 0.10:
        translated = translate_with_dict(question)
        if translated != question: return translated, 'dict'
    return question, None

def fuzzy_correct_word(word, threshold=88):
    if len(word) < 5: return word
    if word in FUZZY_WHITELIST: return word
    result = fuzz_process.extractOne(word, LEGAL_VOCABULARY, scorer=fuzz.ratio, score_cutoff=threshold)
    if result:
        matched, score, _ = result
        if matched != word and matched not in FUZZY_WHITELIST: return matched
    return word

def fuzzy_normalize(question):
    clean = re.sub(r'[؟?،,.]', '', question)
    words, result = clean.split(), []
    for word in words:
        if len(word) < 5 or word.isdigit() or not any('\u0600' <= c <= '\u06ff' for c in word):
            result.append(word); continue
        result.append(fuzzy_correct_word(word))
    return ' '.join(result)

def normalize_question(question):
    log      = []
    question = ' '.join(question.split())
    for n, ar in NUMBERS_AR.items():
        question = re.sub(rf'\b{n}\s*سنين\b', f'{ar} سنوات', question)
        question = re.sub(rf'\b{n}\s*سنة\b',  f'{ar} سنوات', question)
    translated, method = smart_translate(question)
    if translated != question:
        log.append(f'🌐 {translated}')
        question = translated
    corrected = fuzzy_normalize(question)
    if corrected != question: question = corrected
    for col, formal in COLLOQUIAL.items():
        question = re.sub(rf'\b{col}\b', formal, question, flags=re.IGNORECASE)
    question = re.sub(r'[؟?]+', '؟', question).strip()
    if question and not question.endswith('؟'): question += '؟'
    return question, log

def is_legal_question(question):
    q_lower = question.lower()
    if any(p in question for p in PRACTICAL_PATTERNS): return True
    if any(kw in question for kw in LEGAL_KEYWORDS):   return True
    if any(kw in question for kw in LAW_KEYWORDS):     return True
    english_legal = [
        'law','legal','court','judge','penalty','fine','imprisonment',
        'contract','labor','labour','employment','bankruptcy','arbitration',
        'cybercrime','data protection','money laundering','salary',
        'employee','employer','end of service','gratuity','rights',
        'divorce','custody','marriage','notary','lawyer','attorney',
        'fraud','bribery','hacking','evidence','dismissed','termination','wrongful',
    ]
    if any(kw in q_lower for kw in english_legal): return True
    arabic_words = [w for w in question.split() if len(w) > 3 and any('\u0600' <= c <= '\u06ff' for c in w)]
    if len(arabic_words) >= 3:
        try:
            test_docs = vectorstore.similarity_search(question, k=3)
            if test_docs and calculate_coverage(question, test_docs) >= 0.20:
                return True
        except: pass
    return False

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

def expand_query(question):
    if question in _expansion_cache: return _expansion_cache[question]
    result = [question]
    for pattern, expansion in QUERY_EXPANSION.items():
        if pattern in question:
            result = [question, expansion]
            _expansion_cache[question] = result
            return result
    if any(p in question for p in PRACTICAL_PATTERNS):
        if any(w in question for w in ['تعويض','مكافأة','اشتغلت','عملت','سنوات']):
            result.append('مكافأة نهاية الخدمة يستحق العامل سنوات خدمة')
        elif any(w in question for w in ['فصل','فُصلت','فصلوني','طردوني','تم فصلي']):
            result.append('إنهاء عقد العمل تعويض فصل تعسفي')
        elif any(w in question for w in ['حقوق','حقوقي','مستحقات']):
            result.append('حقوق العامل يستحق مكافأة تعويض نظام العمل')
    cleaned = question
    for prefix in ['ما هي ','ما هو ','هل ','متى ','كيف ']: cleaned = cleaned.replace(prefix,'')
    cleaned = cleaned.replace('؟','').strip()
    if cleaned != question and len(cleaned) > 5: result.append(cleaned)
    words = [w for w in question.split() if len(w) > 3 and w not in
             {'هي','هو','ما','في','على','من','إلى','عن','هل','لي','يحق','يجب'}]
    if words: result.append(' '.join(words[:4]))
    _expansion_cache[question] = result
    return result

def rewrite_query(question):
    if question in _rewrite_cache: return _rewrite_cache[question]
    TRIGGER = ['متطلبات','ضوابط','وضح','اشرح','هل لي','هل يحق','أنا','عندي',
               'فصلوني','طردوني','تم فصلي','حقوقي','مستحقاتي']
    if not any(w in question for w in TRIGGER): return expand_query(question)
    try:
        if not ACTIVE_MODELS: return expand_query(question)
        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: return expand_query(question)

def bm25_search(query, k=5, target_law=None):
    tokens  = tokenize_arabic(query)
    scores  = bm25_index.get_scores(tokens)
    results = []
    for idx in scores.argsort()[::-1]:
        if len(results) >= k or scores[idx] < 0.1: break
        meta = bm25_metadatas[idx]
        if target_law and meta.get('law_name') != target_law: continue
        results.append(Document(page_content=bm25_texts[idx], metadata=meta))
    return results

def rerank_docs(docs, question, target_law=None):
    qwords = [w for w in question.split() if len(w) > 2]
    scored = []
    for doc in docs:
        score = 0
        if target_law and doc.metadata.get('law_name') == target_law: score += 8
        score += sum(2 for w in qwords if w in doc.page_content)
        if 'المادة' in doc.metadata.get('article_number',''): score += 3
        score += min(len(doc.page_content) // 200, 3)
        scored.append((score, doc))
    scored.sort(key=lambda x: x[0], reverse=True)
    return [d for _, d in scored]

def calculate_coverage(question, docs):
    if not docs: return 0.0
    words = [w for w in question.split() if len(w) > 3]
    if not words: return 1.0
    all_text = ' '.join(d.page_content for d in docs)
    return sum(1 for w in words if w in all_text or (len(w) >= 4 and w[:4] in all_text)) / len(words)

def build_context(docs):
    parts = []
    for i, doc in enumerate(docs):
        law     = doc.metadata.get('law_name','')
        article = doc.metadata.get('article_number','')
        label   = 'الأكثر صلة' if i == 0 else f'مرجع {i+1}'
        parts.append(f'[{label}] {law} — {article}\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','')
        article = docs[0].metadata.get('article_number','')
        if law and article: answer = f'📋 المرجع: {law} — {article}\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.3
             or any(s in l for s in ['📋','✅','📝','⚠️','•','-','─'])]
    return '\n'.join(clean).strip()

# ══════════════════════════════════════════════════════════
# Generation
# ══════════════════════════════════════════════════════════
def _call_groq(messages):
    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
            arabic = sum(1 for c in answer if '\u0600' <= c <= '\u06ff')
            if arabic / max(len([c for c in answer if c.strip()]),1) < 0.6: continue
            return answer, m['name']
        except Exception as e:
            if '429' in str(e): continue
            raise e
    return None, None

def _call_gemini_chain(messages):
    system = next((m['content'] for m in messages if m['role']=='system'), '')
    user   = next((m['content'] for m in messages if m['role']=='user'), '')
    prompt = f'{system}\n\n{user}'

    for model_name, client in gemini_clients.items():
        try:
            r      = client.generate_content(prompt)
            answer = r.text
            if not answer: continue
            arabic = sum(1 for c in answer if '\u0600' <= c <= '\u06ff')
            if arabic / max(len([c for c in answer if c.strip()]),1) < 0.3: continue
            logger.info(f"Gemini: {model_name.split('/')[-1]}")
            return answer, model_name.split('/')[-1]
        except Exception as e:
            if '429' in str(e): continue
            continue
    return None, None

def generate_with_fallback(messages):
    if rate_limiter.can_use('groq'):
        try:
            result, model = _call_groq(messages)
            if result:
                rate_limiter.record('groq')
                return result, model
        except Exception as e:
            if '429' in str(e) or 'rate' in str(e).lower():
                for _ in range(28): rate_limiter.requests['groq'].append(time.time())
    else:
        pass

    result, model = _call_gemini_chain(messages)
    if result: return result, model

    wait = rate_limiter.wait_time('groq')
    if wait > 0:
        time.sleep(min(wait + 1, 15))
        return generate_with_fallback(messages)
    return None, None

# ══════════════════════════════════════════════════════════
# ask_legal
# ══════════════════════════════════════════════════════════
def ask_legal_core(question: str) -> dict:
    original = question
    question, norm_log = normalize_question(question)

    if not is_legal_question(question):
        test_docs = vectorstore.similarity_search(question, k=3)
        if calculate_coverage(question, test_docs) < 0.20:
            return {'answer': OUT_OF_SCOPE_RESPONSE, 'sources': [], 'coverage': 0,
                    'model': 'out_of_scope', 'normalized': question}

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

    work_clues = ['فصل','تعويض','مكافأة','راتب','أجر','ساعات','إجازة',
                  'موظف','عامل','شركة','نهاية الخدمة','حقوقي','مستحقاتي']
    if not target_law and any(w in question or w in original for w in work_clues):
        target_law = 'نظام العمل'

    top_k       = 8 if len(question.split()) > 5 else 6
    k_per_query = max(3, top_k // len(queries))
    all_docs    = []

    for q in queries:
        if target_law:
            docs = vectorstore.similarity_search(q, k=k_per_query, filter={'law_name': target_law})
            if len(docs) < 2:
                extra = vectorstore.similarity_search(q, k=2)
                docs += [d for d in extra if d not in docs]
        else:
            docs = vectorstore.similarity_search(q, k=k_per_query)
        all_docs.extend(docs)

    bm25_docs, keyword_docs = [], []
    for q in queries: bm25_docs.extend(bm25_search(q, k=5, target_law=target_law))

    if target_law:
        all_in_law = vectorstore.get(where={'law_name': target_law})
        all_words  = set(w for q in queries for w in q.split() if len(w) > 2)
        scored_kw  = [(sum(1 for w in all_words if w in dt), dt, dm)
                      for dt, dm in zip(all_in_law['documents'], all_in_law['metadatas'])
                      if sum(1 for w in all_words if w in dt) >= 1]
        scored_kw.sort(reverse=True)
        keyword_docs = [Document(page_content=dt, metadata=dm) for _, dt, dm in scored_kw[:5]]

    seen, combined = set(), []
    for d in keyword_docs + bm25_docs + all_docs:
        key = d.page_content[:50]
        if key not in seen: seen.add(key); combined.append(d)

    final_docs = rerank_docs(combined, question, target_law)[:top_k]
    coverage   = calculate_coverage(question, final_docs)

    if target_law and not any(d.metadata.get('law_name') == target_law for d in final_docs):
        forced = vectorstore.similarity_search(question, k=top_k, filter={'law_name': target_law})
        if forced: final_docs = forced; coverage = calculate_coverage(question, final_docs)

    if coverage < 0.5 and target_law:
        for d in vectorstore.similarity_search(question, k=top_k, filter={'law_name': target_law}):
            key = d.page_content[:50]
            if key not in seen: seen.add(key); combined.append(d)
        final_docs = rerank_docs(combined, question, target_law)[:top_k]
        coverage   = calculate_coverage(question, final_docs)

    if len(final_docs) == 0 or coverage < 0.25:
        return {'answer': OUT_OF_SCOPE_RESPONSE, 'sources': [], 'coverage': 0,
                'model': 'quality_check', 'normalized': question}

    context = build_context(final_docs)
    answer, model_used = generate_with_fallback([
        {'role': 'system', 'content': SYSTEM_PROMPT},
        {'role': 'user',   'content': f'المواد القانونية:\n{context}\n\nالسؤال: {question}'}
    ])

    if not answer:
        return {'answer': 'كل الموديلات محجوزة حالياً.', 'sources': [], 'coverage': 0, 'model': ''}

    return {
        'answer':     post_process(answer, final_docs),
        'sources':    [{'law': d.metadata.get('law_name',''), 'article': d.metadata.get('article_number','')} for d in final_docs[:3]],
        'coverage':   round(coverage * 100),
        'model':      model_used,
        'normalized': question,
    }

# ══════════════════════════════════════════════════════════
# FastAPI
# ══════════════════════════════════════════════════════════
app = FastAPI(
    title='Saudi Legal AI',
    description='نظام الذكاء الاصطناعي للقانون السعودي — وزارة العدل',
    version='4.1.0',
    lifespan=lifespan
)

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

user_requests = defaultdict(list)

def check_rate_limit(ip: str, max_per_minute: int = 10) -> bool:
    now = time.time()
    user_requests[ip] = [t for t in user_requests[ip] if now - t < 60]
    if len(user_requests[ip]) >= max_per_minute: return False
    user_requests[ip].append(now)
    return True

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

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

@app.get('/')
def root():
    return {
        'name':    'Saudi Legal AI',
        'version': '4.1.0',
        'status':  'running',
        'docs':    '/docs'
    }

@app.get('/health')
def health():
    return {
        'status':        'healthy',
        'chunks':        vectorstore._collection.count() if vectorstore else 0,
        'groq_models':   [m['name'] for m in ACTIVE_MODELS],
        'gemini_models': list(gemini_clients.keys()),
        'version':       '4.1.0'
    }

@app.post('/ask', response_model=QuestionResponse)
async def ask(req: QuestionRequest, request: Request):
    if not req.question.strip():
        raise HTTPException(status_code=400, detail='السؤال فاضي!')
    if len(req.question) > 1000:
        raise HTTPException(status_code=400, detail='السؤال طويل جداً')

    ip = request.headers.get('X-Forwarded-For','unknown').split(',')[0].strip()
    if not check_rate_limit(ip):
        raise HTTPException(status_code=429, detail='حاول مرة أخرى بعد دقيقة')

    t0 = time.time()
    try:
        result = ask_legal_core(req.question)
    except Exception as e:
        logger.error(f"Error: {e}")
        stats['errors'] += 1
        raise HTTPException(status_code=500, detail='خطأ في المعالجة')

    ms     = int((time.time() - t0) * 1000)
    status = 'blocked' if result['model'] in ['out_of_scope','quality_check'] else 'success'

    request_log.appendleft({
        'time': datetime.now().strftime('%H:%M:%S'),
        'ip': ip, 'question': req.question[:60],
        'model': result['model'], 'status': status, 'ms': ms
    })
    stats['total']  += 1
    stats[status if status in stats else 'errors'] += 1
    active_ips[ip]  = active_ips.get(ip, 0) + 1

    return QuestionResponse(
        answer      = result['answer'],
        sources     = result['sources'] if req.include_sources else [],
        coverage    = result['coverage'],
        model       = result['model'],
        duration_ms = ms,
    )

@app.get('/stats')
def get_stats():
    return {
        'total':      stats['total'],
        'success':    stats['success'],
        'blocked':    stats['blocked'],
        'unique_ips': len(active_ips),
        'recent':     list(request_log)[:10]
    }

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

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