from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
from typing import List, Optional
import mysql.connector
import numpy as np
import json
import time
import os
import re
import unicodedata
import threading

# Optional OCR / PDF libraries (safe imports for minimal environments)
try:
    import fitz  # PyMuPDF
except ImportError:
    fitz = None

try:
    import pytesseract
except ImportError:
    pytesseract = None

try:
    from pdf2image import convert_from_path
except ImportError:
    convert_from_path = None

try:
    import cv2
except ImportError:
    cv2 = None

from PIL import Image
from sklearn.feature_extraction.text import TfidfVectorizer

from config import (
    API_HOST,
    API_PORT,
    API_TITLE,
    DB,
    DENSE_SEARCH_MODE,
    DENSE_SEARCH_ENABLED,
    EMBEDDING_DEVICE,
    EMBEDDING_MODEL,
    PRF_MIN_SPARSE_SCORE,
    PRF_SEED_K,
    PRF_SEMANTIC_WEIGHT,
    RAG_DIR,
)

import zlib

try:
    import torch
except ImportError:
    torch = None

try:
    import regex
except ImportError:
    regex = re

if DENSE_SEARCH_MODE == "model" and EMBEDDING_MODEL and torch is not None:
    try:
        from sentence_transformers import SentenceTransformer
    except ImportError:
        SentenceTransformer = None
else:
    SentenceTransformer = None

app = FastAPI(title=API_TITLE)

def normalize_embedding_matrix(matrix):
    """Guarantee cosine-safe stored vectors, including legacy rows."""
    if matrix.size == 0:
        return matrix
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    return matrix / np.clip(norms, 1e-12, None)

def normalize_search_text(text):
    """Normalize Arabic orthography/digits consistently for lexical matching."""
    text = unicodedata.normalize("NFKC", str(text or ""))
    text = text.translate(str.maketrans("٠١٢٣٤٥٦٧٨٩", "0123456789"))
    text = re.sub(r"[\u064B-\u065F\u0670]", "", text)
    text = text.replace("ـ", "")
    text = text.replace("أ", "ا").replace("إ", "ا").replace("آ", "ا").replace("ٱ", "ا")
    text = text.replace("ى", "ي")
    text = text.replace("ة", "ه")
    text = text.replace("ؤ", "و").replace("ئ", "ي")
    return re.sub(r"\s+", " ", text).strip().lower()

def encode_deterministic_vector(text: str, dim: int = 1536) -> np.ndarray:
    """Generate deterministic L2-normalized feature vector matching OpenAIService."""
    vector = np.zeros(dim, dtype=np.float32)
    if not text:
        return vector
    tokens = regex.split(r'[\s\p{P}]+', text.lower()) if hasattr(regex, 'split') else re.split(r'\W+', text.lower())
    tokens = [t for t in tokens if t]
    if not tokens:
        tokens = ['empty']
    
    for token in tokens:
        b = token.encode('utf-8')
        h1 = zlib.crc32(b)
        h2 = zlib.crc32((token + '_seed2').encode('utf-8'))
        idx1 = abs(h1) % dim
        idx2 = abs(h2) % dim
        sign1 = 1.0 if (h1 & 1) else -1.0
        sign2 = 0.5 if (h2 & 1) else -0.5
        vector[idx1] += sign1
        vector[idx2] += sign2

    norm = float(np.linalg.norm(vector))
    if norm > 0:
        vector = np.round(vector / norm, 6)
    return vector

# Global variables for in-memory search
chunk_ids = np.array([])
source_ids = np.array([])
embeddings = np.array([])
texts = []
law_names = []
article_numbers = []
tfidf_matrix = None
vectorizer = None
embedding_dim = 0

def load_all_embeddings():
    """Load or reload all chunk embeddings and build TF-IDF sparse matrix from MySQL."""
    global chunk_ids, source_ids, embeddings, texts, law_names, article_numbers, tfidf_matrix, vectorizer, embedding_dim
    print(f"Loading embeddings from MySQL `{DB['database']}`...")
    start_time = time.time()
    rows = []
    try:
        cnx = mysql.connector.connect(**DB)
        cur = cnx.cursor(dictionary=True)
        cur.execute("SELECT id, rag_source_id, law_name, article_number, content, embedding FROM rag_chunks WHERE embedding IS NOT NULL")
        rows = cur.fetchall()
        cur.close()
        cnx.close()
    except Exception as db_err:
        print(f"WARNING: Could not connect to MySQL database `{DB.get('database')}`: {db_err}")
        rows = []

    if not rows:
        print("WARNING: No embeddings found in DB!")
        chunk_ids = np.array([])
        source_ids = np.array([])
        embeddings = np.array([])
        texts = []
        law_names = []
        article_numbers = []
        tfidf_matrix = None
        vectorizer = None
        embedding_dim = 0
    else:
        valid_rows = []
        raw_embeddings = []
        for r in rows:
            try:
                emb = json.loads(r["embedding"])
                if isinstance(emb, list) and len(emb) in (768, 1536):
                    valid_rows.append(r)
                    raw_embeddings.append(emb)
            except Exception:
                continue

        if not valid_rows:
            print("WARNING: No valid embeddings found!")
            chunk_ids = np.array([])
            source_ids = np.array([])
            embeddings = np.array([])
            texts = []
            law_names = []
            article_numbers = []
            tfidf_matrix = None
            vectorizer = None
            embedding_dim = 0
        else:
            chunk_ids = np.array([r["id"] for r in valid_rows])
            source_ids = np.array([r["rag_source_id"] for r in valid_rows])
            texts = [r["content"] for r in valid_rows]
            law_names = [r.get("law_name") or "" for r in valid_rows]
            article_numbers = [r.get("article_number") or "" for r in valid_rows]
            embeddings = normalize_embedding_matrix(np.array(raw_embeddings, dtype=np.float32))
            embedding_dim = embeddings.shape[1] if embeddings.size > 0 else 0

            print("Building TF-IDF Sparse Matrix for Lexical Search...")
            new_vectorizer = TfidfVectorizer(analyzer='word', ngram_range=(1, 1), max_features=30000, sublinear_tf=True)
            new_tfidf = new_vectorizer.fit_transform([normalize_search_text(text) for text in texts])
            vectorizer = new_vectorizer
            tfidf_matrix = new_tfidf

    print(f"Loaded {len(embeddings)} chunks (dim: {embedding_dim}) in {time.time() - start_time:.2f}s.")
    return len(embeddings)

_embeddings_loaded = False
_is_loading = False
_load_lock = threading.Lock()

model = None
_model_lock = threading.Lock()

def ensure_model_loaded(target_dim: int = 768):
    global model
    if model is not None or DENSE_SEARCH_MODE != "model" or SentenceTransformer is None or target_dim != 768:
        return model
    with _model_lock:
        if model is None:
            try:
                device = "mps" if (torch and hasattr(torch.backends, "mps") and torch.backends.mps.is_available()) else ("cuda" if (torch and torch.cuda.is_available()) else "cpu")
                print(f"Loading {EMBEDDING_MODEL} on {device}...")
                m = SentenceTransformer(EMBEDDING_MODEL, device=device)
                m.max_seq_length = 256
                model = m
                print("Model loaded successfully.")
            except Exception as e:
                print(f"Model load skipped: {e}")
                model = None
    return model

def ensure_embeddings_loaded():
    global _embeddings_loaded, _is_loading
    with _load_lock:
        if _embeddings_loaded or _is_loading:
            return
        _is_loading = True
    try:
        load_all_embeddings()
        if embedding_dim == 768:
            ensure_model_loaded(768)
        _embeddings_loaded = True
    except Exception as e:
        print(f"Error loading embeddings from DB: {e}")
    finally:
        _is_loading = False

# Background loading ensures instant startup so cPanel Passenger and /health never time out
threading.Thread(target=ensure_embeddings_loaded, daemon=True).start()

def encode_query_vector(query: str, target_dim: int) -> Optional[np.ndarray]:
    """Encode query to match stored chunk embeddings dimension."""
    if target_dim == 1536:
        return encode_deterministic_vector(query, dim=1536)
    active_model = ensure_model_loaded(target_dim)
    if active_model is not None:
        try:
            passage = f"query: {query}"
            with torch.inference_mode():
                return active_model.encode([passage], normalize_embeddings=True)[0]
        except Exception as e:
            print("Model encode error:", e)
    return None

class SearchRequest(BaseModel):
    query: str
    limit: int = 20
    filters: Optional[dict] = None  # e.g. {"source_ids": [1, 2]}
    semantic_weight: float = 1.0    # 1.0 = pure semantic (meaning-based), 0.0 = lexical
    mode: Optional[str] = "semantic" # "semantic" | "hybrid" | "exact"
    
def get_rrf_scores(ranked_indices, total_len, k=60):
    """Calculate Reciprocal Rank Fusion scores for a given ranked list."""
    rrf_scores = np.zeros(total_len)
    for rank, idx in enumerate(ranked_indices):
        rrf_scores[idx] = 1.0 / (k + rank + 1)
    return rrf_scores

@app.post("/search")
def search(req: SearchRequest):
    global chunk_ids, source_ids, embeddings, tfidf_matrix, vectorizer, texts
    if not _embeddings_loaded and not _is_loading:
        ensure_embeddings_loaded()

    # Snapshot references to avoid race conditions with concurrent indexing
    cur_chunk_ids = chunk_ids
    cur_source_ids = source_ids
    cur_embeddings = embeddings
    cur_tfidf_matrix = tfidf_matrix
    cur_vectorizer = vectorizer
    cur_texts = texts
    
    total_len = len(cur_chunk_ids)

    if total_len == 0 or cur_tfidf_matrix is None:
        return {"results": []}

    # Extract Filters safely
    valid_indices = np.arange(total_len)
    if req.filters and "source_ids" in req.filters and req.filters["source_ids"]:
        allowed_sources = set(req.filters["source_ids"])
        mask = np.array([sid in allowed_sources for sid in cur_source_ids])
        valid_indices = valid_indices[mask]
        
    if len(valid_indices) == 0:
        return {"results": []}
        
    # 1. Lexical Search (Sparse Vectors - TF-IDF)
    q_sparse = cur_vectorizer.transform([normalize_search_text(req.query)])
    sparse_scores = cur_tfidf_matrix.dot(q_sparse.T).toarray().flatten()
    filtered_sparse_scores = sparse_scores[valid_indices]
    sparse_ranked_local_idx = np.argsort(filtered_sparse_scores)[::-1]
    sparse_ranked_global_idx = valid_indices[sparse_ranked_local_idx]
    sparse_rrf = get_rrf_scores(sparse_ranked_global_idx, total_len)

    # 2. Semantic Search (Dense Query Vector vs Stored Chunks)
    target_dim = cur_embeddings.shape[1] if cur_embeddings.size > 0 else 0
    q_emb = None
    if target_dim > 0 and DENSE_SEARCH_ENABLED:
        q_emb = encode_query_vector(req.query, target_dim)

    # If pure semantic search mode is requested and query embedding succeeded:
    # Rank directly by cosine similarity (neural meaning matching)
    search_mode = req.mode or "semantic"
    if q_emb is not None and (search_mode == "semantic" or req.semantic_weight >= 0.95):
        dense_scores = np.dot(cur_embeddings, q_emb)
        filtered_dense_scores = dense_scores[valid_indices]
        dense_ranked_local_idx = np.argsort(filtered_dense_scores)[::-1]
        
        k = min(req.limit, len(filtered_dense_scores))
        top_k_idx = dense_ranked_local_idx[:k]
        
        results = []
        for idx in top_k_idx:
            global_idx = valid_indices[idx]
            d_score = float(filtered_dense_scores[idx])
            s_score = float(sparse_scores[global_idx]) if len(sparse_scores) > global_idx else 0.0
            
            # لا تظهر النتائج إلا إذا كانت درجة التطابق أكثر من 65%
            if d_score <= 0.65:
                continue

            results.append({
                "chunk_id": int(cur_chunk_ids[global_idx]),
                "source_id": int(cur_source_ids[global_idx]),
                "score": d_score,
                "dense_score": d_score,
                "sparse_score": s_score,
                "hybrid_match_score": d_score,
                "search_mode": "semantic",
                "text": cur_texts[global_idx],
                "content": cur_texts[global_idx],
                "law_name": law_names[global_idx] if law_names else None,
                "article_number": article_numbers[global_idx] if article_numbers else None,
            })
        return {"results": results}

    prf_seed_ids = []
    if q_emb is not None:
        dense_scores = np.dot(cur_embeddings, q_emb)
        filtered_dense_scores = dense_scores[valid_indices]
        dense_ranked_local_idx = np.argsort(filtered_dense_scores)[::-1]
        dense_ranked_global_idx = valid_indices[dense_ranked_local_idx]
        dense_rrf = get_rrf_scores(dense_ranked_global_idx, total_len)
        effective_semantic_weight = max(0.0, min(1.0, req.semantic_weight))
        search_mode = "hybrid"
        query_mode = "vector_1536" if target_dim == 1536 else "model"
    elif DENSE_SEARCH_MODE == "prf" and cur_embeddings.size > 0:
        positive_seed_indices = [
            int(idx)
            for idx in sparse_ranked_global_idx
            if sparse_scores[idx] >= PRF_MIN_SPARSE_SCORE
        ][:max(1, PRF_SEED_K)]

        if positive_seed_indices:
            prf_seed_ids = [int(cur_chunk_ids[i]) for i in positive_seed_indices]
            seed_weights = np.array(
                [sparse_scores[idx] for idx in positive_seed_indices],
                dtype=np.float32,
            )
            seed_vectors = cur_embeddings[positive_seed_indices]
            q_emb = np.average(seed_vectors, axis=0, weights=seed_weights)
            q_emb = q_emb / max(float(np.linalg.norm(q_emb)), 1e-12)

            dense_scores = np.dot(cur_embeddings, q_emb)
            filtered_dense_scores = dense_scores[valid_indices]
            dense_ranked_local_idx = np.argsort(filtered_dense_scores)[::-1]
            dense_ranked_global_idx = valid_indices[dense_ranked_local_idx]
            dense_rrf = get_rrf_scores(dense_ranked_global_idx, total_len)
            effective_semantic_weight = min(
                max(0.0, req.semantic_weight),
                max(0.0, min(1.0, PRF_SEMANTIC_WEIGHT)),
            )
            search_mode = "hybrid_prf"
            query_mode = "prf"
        else:
            dense_scores = np.zeros(total_len)
            dense_rrf = np.zeros(total_len)
            effective_semantic_weight = 0.0
            search_mode = "sparse"
            query_mode = "none"
    else:
        dense_scores = np.zeros(total_len)
        dense_rrf = np.zeros(total_len)
        effective_semantic_weight = 0.0
        search_mode = "sparse"
        query_mode = "none"

    # 3. Hybrid Fusion (Combine RRF Scores)
    hybrid_scores = (effective_semantic_weight * dense_rrf) + ((1.0 - effective_semantic_weight) * sparse_rrf)
    
    # Filter hybrid scores to valid indices
    filtered_hybrid_scores = hybrid_scores[valid_indices]
    filtered_ids = cur_chunk_ids[valid_indices]
    
    # Get Top K
    k = min(req.limit, len(filtered_hybrid_scores))
    top_k_idx = np.argpartition(filtered_hybrid_scores, -k)[-k:]
    # Sort descending
    top_k_idx = top_k_idx[np.argsort(filtered_hybrid_scores[top_k_idx])[::-1]]
    
    results = []
    for idx in top_k_idx:
        global_idx = valid_indices[idx]
        d_score = float(dense_scores[global_idx]) if len(dense_scores) > global_idx else 0.0
        s_score = float(sparse_scores[global_idx]) if len(sparse_scores) > global_idx else 0.0
        
        # When dense neural search is active, the true cosine similarity represents the neural match
        if search_mode in ("hybrid", "hybrid_prf") and d_score > 0:
            hybrid_match_score = min(1.0, d_score + (s_score * 0.05)) if s_score > 0 else d_score
        else:
            hybrid_match_score = s_score
            
        # لا تظهر النتائج إلا إذا كانت درجة التطابق أكثر من 65%
        if hybrid_match_score <= 0.65:
            continue

        # In hybrid_prf, guard against PRF centroid drift where unrelated laws get high dense scores
        if search_mode == "hybrid_prf" and s_score < 0.05:
            arabic_noise = {'في', 'من', 'على', 'عن', 'إلى', 'مع', 'هذا', 'هذه', 'ذلك', 'نظام', 'المملكة', 'سعودية', 'السعودية'}
            split_fn = getattr(regex, 'split', re.split)
            pattern = r'[\s\p{P}]+' if hasattr(regex, 'split') and regex is not re else r'[\s\W]+'
            q_tokens = [w for w in split_fn(pattern, req.query.lower()) if len(w) >= 3 and w not in arabic_noise]
            if q_tokens:
                chunk_txt = cur_texts[global_idx].lower()
                chunk_law = (law_names[global_idx] or '').lower() if law_names else ''
                has_substantive = any(tok in chunk_txt or tok in chunk_law for tok in q_tokens)
                if not has_substantive:
                    stems = [re.sub(r'^(ال|لل|بال)', '', tok) for tok in q_tokens if len(tok) >= 4]
                    has_substantive = any(stem in chunk_txt or stem in chunk_law for stem in stems)
                if not has_substantive:
                    continue

        results.append({
            "chunk_id": int(filtered_ids[idx]),
            "source_id": int(cur_source_ids[global_idx]),
            "score": float(filtered_hybrid_scores[idx]),
            "dense_score": d_score,
            "sparse_score": s_score,
            "hybrid_match_score": float(hybrid_match_score),
            "search_mode": search_mode,
            "text": cur_texts[global_idx],
            "content": cur_texts[global_idx],
            "law_name": law_names[global_idx] if law_names else None,
            "article_number": article_numbers[global_idx] if article_numbers else None,
        })
        
    return {
        "results": results,
        "meta": {
            "search_mode": search_mode,
            "dense_enabled": search_mode in ("hybrid", "hybrid_prf"),
            "query_embedding_mode": query_mode,
            "semantic_weight": effective_semantic_weight,
            "prf_seed_chunk_ids": prf_seed_ids,
            "total_chunks": total_len,
            "dimension": target_dim,
        },
    }

@app.get("/")
def root():
    target_dim = embeddings.shape[1] if embeddings.size > 0 else 0
    return {
        "service": "Meezan AI Saudi Legal Semantic Search API",
        "status": "online",
        "chunks_loaded": len(embeddings),
        "dimension": target_dim,
        "docs": "/docs",
        "health": "/health",
    }

@app.get("/health")
def health():
    target_dim = embeddings.shape[1] if embeddings.size > 0 else 0
    dense_active = embeddings.size > 0 and DENSE_SEARCH_ENABLED
    return {
        "status": "ok",
        "chunks": len(embeddings),
        "dimension": target_dim,
        "loading": _is_loading,
        "ready": _embeddings_loaded,
        "hybrid_enabled": dense_active,
        "dense_enabled": dense_active,
        "search_mode": "hybrid" if dense_active else "sparse",
        "query_embedding_mode": "vector_1536" if target_dim == 1536 else ("model" if model is not None else ("prf" if DENSE_SEARCH_MODE == "prf" else "none")),
        "embedding_model": "vector_1536" if target_dim == 1536 else EMBEDDING_MODEL,
    }


# ==========================================
# OCR & Document Ingestion Pipeline
# ==========================================

class ProcessDocumentRequest(BaseModel):
    file_path: str
    source_id: int
    law_name: str

def preprocess_image_for_ocr(img_pil: Image.Image) -> Image.Image:
    """Enhance image quality for Arabic OCR using OpenCV (Upscaling + Grayscale + Dilation)."""
    # 1. Convert PIL Image to OpenCV NumPy Array
    img_cv = np.array(img_pil)
    
    # Check if image is colored (RGB) or already grayscale
    if len(img_cv.shape) == 3 and img_cv.shape[2] == 3:
        # Convert RGB to BGR (OpenCV format), then to Grayscale
        img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
    else:
        gray = img_cv

    # 2. Upscale 2x to give Tesseract more pixel density for faint typewriter fonts
    scaled = cv2.resize(gray, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC)
    
    # 3. Morphological Dilation to thicken faint Arabic strokes and preserve dots (nuqtas)
    # We invert the image (white text on black) because dilate thickens white pixels
    inv = cv2.bitwise_not(scaled)
    kernel = np.ones((2,2), np.uint8)
    dilated_inv = cv2.dilate(inv, kernel, iterations=1)
    
    # Invert back to black text on white background
    final_img = cv2.bitwise_not(dilated_inv)
    
    # 4. Convert back to PIL Image so pytesseract can process it natively
    return Image.fromarray(final_img)

def extract_text_from_pdf(pdf_path: str) -> str:
    """Try native python extraction first. If the file is mostly images, fallback to OCR."""
    doc = fitz.open(pdf_path)
    full_text = ""
    run_ocr = False
    
    # Check first page to see if it contains native text
    if len(doc) > 0:
        page_text = doc[0].get_text("text").strip()
        if len(page_text) < 50: # Likely a scanned PDF
            run_ocr = True
            
    if not run_ocr:
        for page in doc:
            full_text += page.get_text("text") + "\n"
        # If the total extracted text is suspiciously small, fallback to OCR
        if len(full_text.strip()) < 100 * len(doc):
             run_ocr = True
             
    if run_ocr:
        print(f"Fallback to OCR for PDF: {pdf_path}")
        full_text = ""
        # Convert PDF to list of images (300 DPI for Arabic OCR)
        images = convert_from_path(pdf_path, dpi=300)
        for i, img in enumerate(images):
            # Preprocess the scanned page for better OCR
            enhanced_img = preprocess_image_for_ocr(img)
            
            # Arabic OCR config. psd 3 fully automatic page segmentation
            custom_oem_psm_config = r'--oem 3 --psm 3'
            text = pytesseract.image_to_string(enhanced_img, lang='ara', config=custom_oem_psm_config)
            full_text += text + "\n"
            
    return full_text

def extract_text_from_image(image_path: str) -> str:
    """Run Tesseract OCR on a single image file."""
    img = Image.open(image_path)
    enhanced_img = preprocess_image_for_ocr(img)
    
    custom_oem_psm_config = r'--oem 3 --psm 3'
    return pytesseract.image_to_string(enhanced_img, lang='ara', config=custom_oem_psm_config)

def fix_arabic_ocr_errors(text: str) -> str:
    """Fix common Arabic OCR mistakes produced by Tesseract and PDF extractors."""
    
    # 1. Fix mid-word line breaks (Arabic letter + \n + Arabic letter)
    # Arabic Unicode range: \u0600-\u06FF
    text = re.sub(r'([\u0600-\u06FF])\n([\u0600-\u06FF])', r'\1\2', text)
    
    # 2. Fix آ (Alef-Madda U+0622) misread as ك (Kaf) in common words
    # This is THE most frequent Tesseract Arabic error.
    # We use word-boundary-aware replacements for safety.
    kaf_fixes = {
        'آان': 'كان',
        'آانت': 'كانت',
        'آانوا': 'كانوا',
        'آل': 'كل',
        'آما': 'كما',
        'آذلك': 'كذلك',
        'آثر': 'كثر',        
        'آثير': 'كثير',
        'آبير': 'كبير',
        'آبيرة': 'كبيرة',
        'آامل': 'كامل',
        'آاملة': 'كاملة',
        'آتاب': 'كتاب',
        'آتابة': 'كتابة',
        'آافة': 'كافة',
        'آيفية': 'كيفية',
        'آيف': 'كيف',
        'آأن': 'كأن',
        'تفكيآ': 'تفكيك',
        'الشرآة': 'الشركة',
        'شرآة': 'شركة',
        'الشرآات': 'الشركات',
        'شرآات': 'شركات',
        'المشترآة': 'المشتركة',
        'مشترآة': 'مشتركة',
        'المشارآة': 'المشاركة',
        'مشارآة': 'مشاركة',
        'المرآز': 'المركز',
        'مرآز': 'مركز',
        'المرآزي': 'المركزي',
        'مرآزي': 'مركزي',
        'المتحرآة': 'المتحركة',
        'الحرآة': 'الحركة',
        'حرآة': 'حركة',
        'الترآيب': 'التركيب',
        'ترآيب': 'تركيب',
        'المالآ': 'المالك',
        'تمليآ': 'تمليك',
        'تملآ': 'تملك',
        'يملآ': 'يملك',
        'الملآية': 'الملكية',
        'ملآية': 'ملكية',
        'اآتتاب': 'اكتتاب',
        'الاآتتاب': 'الاكتتاب',
        'اآتساب': 'اكتساب',
        'الاآتساب': 'الاكتساب',
        'آقانون': 'كقانون',
        'آقاعدة': 'كقاعدة',
        'آاذبة': 'كاذبة',
        'التحكيآ': 'التحكيك',
        'المحكمآ': 'المحكمك',
        'المحاسآ': 'المحاسك',
        'آانون': 'كانون',
        'آريم': 'كريم',
        'آريمة': 'كريمة',
        'إشراآ': 'إشراك',
        'الإشتراآ': 'الاشتراك',
        'اشتراآ': 'اشتراك',
        'آتب': 'كتب',
        'مآتب': 'مكتب',
        'المآتب': 'المكتب',
        'آتابي': 'كتابي',
        'لشرآة': 'لشركة',
        'بشرآة': 'بشركة',
        'والشرآات': 'والشركات',
        'بالشرآة': 'بالشركة',
        'للشرآة': 'للشركة',
        'للشرآات': 'للشركات',        
        'آقانون': 'كقانون',
        'آثيراً': 'كثيراً',
        'آتابه': 'كتابه',
        'بذلآ': 'بذلك',
        'ذلآ': 'ذلك',
        'هنالآ': 'هنالك',
        'هناآ': 'هناك',
        'مسالآ': 'مسالك',
        'الأملاآ': 'الأملاك',
        'أملاآ': 'أملاك',
        'إمساآه': 'إمساكه',
        'المرآب': 'المركب',
        'مرآب': 'مركب',
        'الترآ': 'الترك',
        'المشآلة': 'المشكلة',
        'مشآلة': 'مشكلة',
    }
    
    for wrong, correct in kaf_fixes.items():
        text = text.replace(wrong, correct)
    
    # 3. Fix scattered Seen-Jeem artifacts from badly decoded PDF fonts  
    # ѧ (Cyrillic) appears in broken Arabic PDFs  
    text = text.replace('ѧ', '')
    
    # 4. Remove excessive whitespace
    text = re.sub(r'[ \t]+', ' ', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    
    # 5. Fix common diacritics/tashkeel errors
    # Remove zero-width characters that break word matching
    text = re.sub(r'[\u200b\u200c\u200d\u200e\u200f\ufeff]', '', text)
    
    # 6. Fix shattered Arabic list brackets and dashes (RTL inversion and newline artifacts)
    # Fixes: "-\n)\n1\n(" or "–\n) 1 (" -> "- (١)"
    def to_arabic_numeral_with_dash(match):
        num_str = match.group(1)
        arabic_nums = "٠١٢٣٤٥٦٧٨٩"
        english_nums = "0123456789"
        trans = str.maketrans(english_nums, arabic_nums)
        num_arabic = num_str.translate(trans)
        return f"- ({num_arabic})"
        
    def to_arabic_numeral_no_dash(match):
        num_str = match.group(1)
        arabic_nums = "٠١٢٣٤٥٦٧٨٩"
        english_nums = "0123456789"
        trans = str.maketrans(english_nums, arabic_nums)
        num_arabic = num_str.translate(trans)
        return f"({num_arabic})"
        
    text = re.sub(r'[-–_]\s*\)\s*(\d+)\s*\(', to_arabic_numeral_with_dash, text)
    text = re.sub(r'[-–_]\s*\(\s*(\d+)\s*\)', to_arabic_numeral_with_dash, text)
    # Fixes inverted RTL parenthesis: ") 1 (" -> "(١)"
    text = re.sub(r'\)\s*(\d+)\s*\(', to_arabic_numeral_no_dash, text)
    # Fixes spaced parenthesis: "( 1 )" -> "(١)"
    text = re.sub(r'\(\s*(\d+)\s*\)', to_arabic_numeral_no_dash, text)
    
    # 7. Fix orphaned punctuation: periods or commas that dropped to a new line
    # E.g. "Some text\n ." -> "Some text."
    text = re.sub(r'\n\s*([.،,])(?=\n|$)', r'\1', text)
    
    return text.strip()


def process_and_store_document(file_path: str, source_id: int, law_name: str):
    global embeddings, chunk_ids, source_ids
    print(f"Background task starting for {file_path} (Source ID: {source_id})")
    
    try:
        # 1. Extract Full Text
        ext = file_path.lower().split('.')[-1]
        if ext == 'pdf':
            full_text = extract_text_from_pdf(file_path)
        elif ext in ['jpg', 'jpeg', 'png']:
            full_text = extract_text_from_image(file_path)
        else:
            print(f"Unsupported file extension: {ext}")
            return
            
        # Clean up: Normalize Unicode, fix OCR errors, clean whitespace
        full_text = unicodedata.normalize('NFKC', full_text)
        full_text = fix_arabic_ocr_errors(full_text)
        full_text = full_text.replace("\n\n", "\n").replace("  ", " ")
        
        import re
        # 2. Chunking Logic (Split by "مادة (number)" or "المادة (number)")
        pattern = r"(?:^|\n)\s*((?:ال)?مادة\s*[\d١٢٣٤٥٦٧٨٩٠]+)"
        parts = re.split(pattern, full_text)
        
        chunks = []
        if len(parts) > 1:
             # parts[0] is preamble.
             if parts[0].strip():
                 chunks.append({
                     "article": "مقدمة",
                     "content": parts[0].strip()
                 })
             for i in range(1, len(parts), 2):
                 article_title = parts[i].strip() # e.g. "مادة 1"
                 article_content = parts[i+1].strip() if (i+1) < len(parts) else ""
                 
                 # Extract just the number for the DB field
                 num_match = re.search(r"[\d١٢٣٤٥٦٧٨٩٠]+", article_title)
                 article_num = num_match.group(0) if num_match else article_title
                 
                 chunks.append({
                     "article": article_num,
                     "content": article_title + "\n" + article_content
                 })
        else:
             # No articles found, just chunk by length or treat as single document
             chunks.append({
                 "article": "1",
                 "content": full_text.strip()
             })
             
        if not chunks:
            print("No text found to process.")
            return

        print(f"Extracted {len(chunks)} chunks from {file_path}")

        # 3. Save to TXT and JSON for manual review
        import re
        sanitized_title = re.sub(r'[^\w\u0600-\u06FF\s-]', '', law_name).strip().replace(' ', '_')
        base_filename = f"source_{source_id}_{sanitized_title}"
        json_filename = f"{base_filename}.json"
        txt_filename = f"{base_filename}.txt"
        
        STORAGE_DIR = RAG_DIR
        os.makedirs(STORAGE_DIR, exist_ok=True)
        json_path = os.path.join(STORAGE_DIR, json_filename)
        txt_path = os.path.join(STORAGE_DIR, txt_filename)
        
        result = {
            "source_id": source_id,
            "title": law_name,
            "type": "file",
            "total_articles": len(chunks),
            "articles": []
        }
        for i, c in enumerate(chunks):
            result["articles"].append({
                "index": i,
                "article_number": c["article"],
                "part": c.get("part"),
                "chapter": c.get("chapter"),
                "section": c.get("section"),
                "word_count": len(c["content"].split()),
                "content": c["content"]
            })
            
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=4)
            
        with open(txt_path, 'w', encoding='utf-8') as f:
            f.write(f"=== {law_name} ===\n")
            f.write(f"إجمالي المواد: {len(chunks)}\n")
            f.write("=" * 50 + "\n\n")
            for chunk in chunks:
                f.write(f"--- المادة {chunk['article']} ---\n")
                if chunk.get("part"): f.write(f"الباب: {chunk['part']}\n")
                if chunk.get("chapter"): f.write(f"الفصل: {chunk['chapter']}\n")
                if chunk.get("section"): f.write(f"الفرع: {chunk['section']}\n")
                f.write(chunk['content'] + "\n\n")
                
        print(f"Successfully fully processed {file_path} into {len(chunks)} chunks for manual review.")
        
        # Notify Laravel (Optional callback. Here we assume Laravel will just poll the chunks table or check RagSource status if we updated it).
        # We can update RagSource status here
        cnx = mysql.connector.connect(**DB)
        cur = cnx.cursor()
        cur.execute("UPDATE rag_sources SET status = 'pending_review' WHERE id = %s", (source_id,))
        cnx.commit()
        cur.close()
        cnx.close()
        
    except Exception as e:
        print(f"ERROR processing document {file_path}: {str(e)}")
        try:
             cnx = mysql.connector.connect(**DB)
             cur = cnx.cursor()
             cur.execute("UPDATE rag_sources SET status = 'failed' WHERE id = %s", (source_id,))
             cnx.commit()
             cur.close()
             cnx.close()
        except:
             pass

@app.post("/process-document")
def process_document(req: ProcessDocumentRequest, background_tasks: BackgroundTasks):
    if not os.path.exists(req.file_path):
        raise HTTPException(status_code=404, detail=f"File not found: {req.file_path}")
        
    # Kick off processing in the background so the HTTP request returns quickly
    background_tasks.add_task(process_and_store_document, req.file_path, req.source_id, req.law_name)
    return {"status": "processing_started", "message": "Document queued for OCR and embedding in background."}

# ==========================================
# Text Extraction & Chunking Only (No Embedding)
# ==========================================

class ExtractTextRequest(BaseModel):
    file_path: str

@app.post("/extract-text")
def extract_text(req: ExtractTextRequest):
    """Extract text from a PDF/Image and split into articles — NO embedding, NO DB save."""
    if not os.path.exists(req.file_path):
        raise HTTPException(status_code=404, detail=f"File not found: {req.file_path}")
    
    ext = req.file_path.lower().split('.')[-1]
    if ext == 'pdf':
        full_text = extract_text_from_pdf(req.file_path)
        method = "PyMuPDF/OCR"
    elif ext in ['jpg', 'jpeg', 'png']:
        full_text = extract_text_from_image(req.file_path)
        method = "Tesseract OCR"
    else:
        raise HTTPException(status_code=400, detail=f"Unsupported file type: {ext}")
    
    # Import chunk function (defined below)
    full_text_normalized = unicodedata.normalize('NFKC', full_text)
    full_text_cleaned = fix_arabic_ocr_errors(full_text_normalized)
    chunks = chunk_text_into_articles(full_text_cleaned)
    
    return {
        "file": os.path.basename(req.file_path),
        "method": method,
        "total_chars": len(full_text_cleaned),
        "total_words": len(full_text_cleaned.split()),
        "total_articles": len(chunks),
        "articles": [
            {
                "article_number": c["article"],
                "word_count": len(c["content"].split()),
                "content_preview": c["content"][:200]
            }
            for c in chunks
        ]
    }

@app.get("/extract-source/{source_id}")
def extract_source(source_id: int):
    """Extract text from a rag_source and split into articles — NO embedding, NO DB save."""
    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)
    cur.execute("SELECT id, title, type, content, file_path FROM rag_sources WHERE id = %s AND deleted_at IS NULL", (source_id,))
    src = cur.fetchone()
    cur.close()
    cnx.close()
    
    if not src:
        raise HTTPException(status_code=404, detail=f"Source {source_id} not found")
    
    full_text = ""
    method = ""
    
    if src['type'] == 'text' and src['content']:
        full_text = src['content']
        method = "Database text"
    elif src['type'] == 'file' and src['file_path']:
        abs_path = src['file_path']
        if not os.path.isabs(abs_path):
            for base in [os.path.dirname(RAG_DIR)]:
                candidate = os.path.join(base, abs_path)
                if os.path.exists(candidate):
                    abs_path = candidate
                    break
        if not os.path.exists(abs_path):
            raise HTTPException(status_code=404, detail=f"File not found: {abs_path}")
        ext = abs_path.lower().split('.')[-1]
        if ext == 'pdf':
            full_text = extract_text_from_pdf(abs_path)
        elif ext in ['jpg', 'jpeg', 'png']:
            full_text = extract_text_from_image(abs_path)
        method = "File OCR/Extract"
    else:
        raise HTTPException(status_code=400, detail="Source has no text content or file")
    
    full_text = unicodedata.normalize('NFKC', full_text)
    full_text = fix_arabic_ocr_errors(full_text)
    chunks = chunk_text_into_articles(full_text)
    
    return {
        "source_id": source_id,
        "title": src['title'],
        "type": src['type'],
        "method": method,
        "total_chars": len(full_text),
        "total_words": len(full_text.split()),
        "total_articles": len(chunks),
        "articles": [
            {
                "index": i,
                "article_number": c["article"],
                "word_count": len(c["content"].split()),
                "content": c["content"]
            }
            for i, c in enumerate(chunks)
        ]
    }

# ==========================================
# Universal Arabic Article Chunking
# ==========================================

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

ARABIC_TO_WESTERN = str.maketrans('٠١٢٣٤٥٦٧٨٩', '0123456789')

def normalize_article_number(raw: str) -> str:
    """تحويل أي معرف مادة إلى رقم."""
    raw = raw.strip()
    for word, num in ORDINAL_MAP.items():
        if word in raw:
            return num
    digits = re.findall(r'[\d٠-٩]+', raw)
    if digits:
        return digits[0].translate(ARABIC_TO_WESTERN)
    return raw

def update_hierarchy(headers_text: str, current: dict) -> dict:
    if not headers_text.strip():
        return current
        
    lines = [line.strip() for line in headers_text.split('\n') if line.strip()]
    last_level = None
    
    for line in lines:
        if 'باب' in line or 'الباب' in line or 'كتاب' in line or 'الكتاب' in line or 'قسم' in line or 'القسم' in line:
            current['part'] = line
            current['chapter'] = None
            current['section'] = None
            last_level = 'part'
        elif 'فصل' in line or 'الفصل' in line or 'مبحث' in line or 'المبحث' in line:
            current['chapter'] = line
            current['section'] = None
            last_level = 'chapter'
        elif re.match(r'^[\d٠-٩]+[\s\-\.–]', line) or 'أولا' in line or 'ثانيا' in line or 'فرع' in line or 'الفرع' in line or 'مطلب' in line:
            current['section'] = line
            last_level = 'section'
        else:
            if last_level == 'section' and current['section']:
                current['section'] += ' - ' + line
            elif last_level == 'chapter' and current['chapter']:
                current['chapter'] += ' - ' + line
            elif last_level == 'part' and current['part']:
                current['part'] += ' - ' + line
            else:
                if current['part']:
                     current['part'] += ' - ' + line
                else:
                     current['part'] = line
                     last_level = 'part'
                     
    return current

def extract_trailing_headers(text: str) -> tuple:
    """تفصل العناوين الهيكلية (مثل اسم الباب أو الفصل) الموجودة في نهاية المادة."""
    lines = text.split('\n')
    headers = []
    body_lines = []
    
    for i in range(len(lines)-1, -1, -1):
        line = lines[i]
        line_stripped = line.strip()
        
        if not line_stripped:
            headers.append(line)
            continue
            
        words = line_stripped.split()
        
        # Heuristic for structural header: short and doesn't end with standard punctuation
        if len(words) <= 12 and not line_stripped.endswith(('.', '،', ':', '؟', '!')):
            headers.append(line)
        else:
            body_lines = lines[:i+1]
            break
            
    headers.reverse()
    
    # Clean leading empty lines from headers
    while headers and not headers[0].strip():
        headers.pop(0)
        
    return '\n'.join(body_lines).strip(), '\n'.join(headers).strip()

def chunk_text_into_articles(full_text: str) -> list[dict]:
    """
    يقوم بتقسيم النص الكامل إلى مواد قانونية باستخدام regex ديناميكي
    يتعرف على 'مادة (1)', 'المادة الأولى', 'مادة 1 مكرر', إلخ.
    ويدمج العناوين الهيكلية في بداية المادة التالية.
    """
    full_text = unicodedata.normalize('NFKC', full_text)
    full_text = fix_arabic_ocr_errors(full_text)
    
    # تنسيق وتقليل المسافات الزائدة والسطور الفارغة
    lines_arr = []
    for line in full_text.split('\n'):
        line = re.sub(r' +', ' ', line.strip())
        lines_arr.append(line)
    
    full_text = '\n'.join(lines_arr)
    full_text = re.sub(r'\n{3,}', '\n\n', full_text)

    ordinals = '|'.join(ORDINAL_MAP.keys())

    pattern = (
        r'(?:^|\n)\s*((?:ال)?مادة'
        r'\s*[\(\[\-:\s]*'
        r'(?:[\d٠-٩]+|' + ordinals + r')'
        r'[\)\]\-:\s]*'
        r'(?:\s*(?:مكرر|مكرراً?)\s*(?:[\(\[]\s*[\d٠-٩]+\s*[\)\]])?)?'
        r')'
    )
    
    parts = re.split(pattern, full_text)
    chunks = []
    current_hierarchy = {'part': None, 'chapter': None, 'section': None}
    
    if len(parts) > 1:
        # المقدمة (ما قبل المادة الأولى)
        intro_text = parts[0].strip()
        if intro_text:
            intro_body, intro_headers = extract_trailing_headers(intro_text)
            if len(intro_body) > 20:
                chunks.append({
                    "article": "مقدمة",
                    "content": intro_body,
                    "part": None, "chapter": None, "section": None
                })
            current_hierarchy = update_hierarchy(intro_headers, current_hierarchy)
            
        # Loop through matched articles
        for i in range(1, len(parts), 2):
            title = parts[i].strip()
            raw_body = parts[i+1].strip() if (i+1) < len(parts) else ""
            
            body, trailing = extract_trailing_headers(raw_body)
            num = normalize_article_number(title)
            
            if 'مكرر' in title:
                m = re.search(r'مكرر(?:اً?)?\s*[\(\[]\s*([\d٠-٩]+)\s*[\)\]]', title)
                if m:
                    num = f"{num} مكرر {m.group(1).translate(ARABIC_TO_WESTERN)}"
                else:
                    num = f"{num} مكرر"

            # بناء المحتوى النهائي للمادة
            final_content = []
            final_content.append(title)
            if body:
                final_content.append(body)

            chunks.append({
                "article": num,
                "content": '\n\n'.join(final_content),
                "part": current_hierarchy['part'],
                "chapter": current_hierarchy['chapter'],
                "section": current_hierarchy['section']
            })

            current_hierarchy = update_hierarchy(trailing, current_hierarchy)
            
        # إلحاق أي عناوين نهائية متبقية بالمادة الأخيرة
        if trailing and chunks:
            chunks[-1]['content'] += '\n\n' + trailing

    else:
        text = full_text.strip()
        if len(text) > 2500:
            pos = 0
            idx = 1
            while pos < len(text):
                chunk = text[pos:pos+2000]
                if len(chunk.strip()) > 50:
                    chunks.append({"article": str(idx), "content": chunk.strip(), "part": None, "chapter": None, "section": None})
                    idx += 1
                pos += 1800
        else:
            chunks.append({"article": "1", "content": text, "part": None, "chapter": None, "section": None})
            
    return chunks


def reprocess_single_source(source_id: int, title: str, content_text: str, file_path: str, source_type: str):
    """Reprocess a single rag_source: extract text, chunk, embed, save."""
    global embeddings, chunk_ids, source_ids
    
    print(f"Reprocessing source {source_id}: {title} (type={source_type})")
    
    try:
        full_text = ""
        
        if source_type == 'text' and content_text:
            full_text = content_text
        elif source_type == 'file' and file_path:
            # Resolve full path
            abs_path = file_path
            if not os.path.isabs(file_path):
                # Try common base paths
                for base in [
                    os.path.dirname(RAG_DIR),
                    os.path.dirname(os.path.dirname(RAG_DIR)),
                ]:
                    candidate = os.path.join(base, file_path)
                    if os.path.exists(candidate):
                        abs_path = candidate
                        break
            
            if not os.path.exists(abs_path):
                print(f"  File not found: {abs_path}")
                return 0
            
            ext = abs_path.lower().split('.')[-1]
            if ext == 'pdf':
                full_text = extract_text_from_pdf(abs_path)
            elif ext in ['jpg', 'jpeg', 'png']:
                full_text = extract_text_from_image(abs_path)
            else:
                print(f"  Unsupported file type: {ext}")
                return 0
        elif source_type == 'url':
            print(f"  Skipping URL source (no text to process)")
            return 0
        else:
            print(f"  No content or file for source {source_id}")
            return 0
        
        if not full_text or len(full_text.strip()) < 10:
            print(f"  Empty content for source {source_id}")
            return 0
        
        # Chunk the text
        chunks = chunk_text_into_articles(full_text)
        
        if not chunks:
            print(f"  No chunks produced for source {source_id}")
            return 0
        
        print(f"  Extracted {len(chunks)} chunks")
        
        # Save JSON and TXT for manual review
        import re
        sanitized_title = re.sub(r'[^\w\u0600-\u06FF\s-]', '', title).strip().replace(' ', '_')
        base_filename = f"source_{source_id}_{sanitized_title}"
        json_filename = f"{base_filename}.json"
        txt_filename = f"{base_filename}.txt"
        
        STORAGE_DIR = RAG_DIR
        os.makedirs(STORAGE_DIR, exist_ok=True)
        json_path = os.path.join(STORAGE_DIR, json_filename)
        txt_path = os.path.join(STORAGE_DIR, txt_filename)
        
        result = {
            "source_id": source_id,
            "title": title,
            "type": source_type,
            "total_articles": len(chunks),
            "articles": []
        }
        for i, c in enumerate(chunks):
            result["articles"].append({
                "index": i,
                "article_number": c["article"],
                "part": c.get("part"),
                "chapter": c.get("chapter"),
                "section": c.get("section"),
                "word_count": len(c["content"].split()),
                "content": c["content"]
            })
            
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=4)
            
        with open(txt_path, 'w', encoding='utf-8') as f:
            f.write(f"=== {title} ===\n")
            f.write(f"إجمالي المواد: {len(chunks)}\n")
            f.write("=" * 50 + "\n\n")
            for chunk in chunks:
                f.write(f"--- المادة {chunk['article']} ---\n")
                if chunk.get("part"): f.write(f"الباب: {chunk['part']}\n")
                if chunk.get("chapter"): f.write(f"الفصل: {chunk['chapter']}\n")
                if chunk.get("section"): f.write(f"الفرع: {chunk['section']}\n")
                f.write(chunk['content'] + "\n\n")
        
        # Update source status to pending_review
        cnx = mysql.connector.connect(**DB)
        cur = cnx.cursor()
        cur.execute("UPDATE rag_sources SET status = 'pending_review' WHERE id = %s", (source_id,))
        cnx.commit()
        cur.close()
        cnx.close()
        
        print(f"  ✅ Source {source_id} extracted for review: {txt_filename}")
        return len(chunks)
        
    except Exception as e:
        print(f"  ❌ Error processing source {source_id}: {str(e)}")
        import traceback
        traceback.print_exc()
        return 0


def run_reprocess_all():
    """Background task: reprocess all sources."""
    global embeddings, chunk_ids, source_ids
    
    print("=" * 60)
    print("STARTING FULL REPROCESSING OF ALL RAG SOURCES")
    print("=" * 60)
    
    # Clear in-memory cache (will be rebuilt)
    embeddings = np.array([])
    chunk_ids = np.array([])
    source_ids = np.array([])
    
    # Fetch all sources
    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)
    cur.execute("SELECT id, title, type, content, file_path FROM rag_sources WHERE deleted_at IS NULL ORDER BY id")
    sources = cur.fetchall()
    cur.close()
    cnx.close()
    
    total_chunks = 0
    processed = 0
    failed = 0
    
    for src in sources:
        count = reprocess_single_source(
            source_id=src['id'],
            title=src['title'] or f"Source {src['id']}",
            content_text=src['content'] or '',
            file_path=src['file_path'] or '',
            source_type=src['type'] or 'text'
        )
        if count > 0:
            total_chunks += count
            processed += 1
        else:
            failed += 1
    
    # Reload all embeddings into memory
    print("\nReloading all embeddings into memory...")
    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)
    cur.execute("SELECT id, rag_source_id, embedding FROM rag_chunks WHERE embedding IS NOT NULL")
    rows = cur.fetchall()
    cur.close()
    cnx.close()
    
    if rows:
        chunk_ids = np.array([r["id"] for r in rows])
        source_ids = np.array([r["rag_source_id"] for r in rows])
        embeddings = normalize_embedding_matrix(
            np.array([json.loads(r["embedding"]) for r in rows], dtype=np.float32)
        )
    
    print("=" * 60)
    print(f"REPROCESSING COMPLETE")
    print(f"  Sources processed: {processed}")
    print(f"  Sources failed/skipped: {failed}")
    print(f"  Total chunks created: {total_chunks}")
    print(f"  Embeddings in memory: {len(embeddings)}")
    print("=" * 60)


@app.post("/reload")
def reload_embeddings():
    """Reload all chunk embeddings from MySQL without re-extracting."""
    count = load_all_embeddings()
    return {"status": "success", "message": f"Successfully reloaded {count} embeddings in memory", "count": count}


@app.post("/reprocess-all")
def reprocess_all(background_tasks: BackgroundTasks):
    """Reprocess ALL rag_sources: delete old chunks, re-extract, re-chunk, re-embed."""
    background_tasks.add_task(run_reprocess_all)
    return {
        "status": "reprocessing_started",
        "message": "Full reprocessing of all sources started in background. Check server logs for progress."
    }


@app.post("/reprocess-source/{source_id}")
def reprocess_source(source_id: int, background_tasks: BackgroundTasks):
    """Reprocess a single rag_source by ID."""
    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)
    cur.execute("SELECT id, title, type, content, file_path FROM rag_sources WHERE id = %s AND deleted_at IS NULL", (source_id,))
    src = cur.fetchone()
    cur.close()
    cnx.close()
    
    if not src:
        raise HTTPException(status_code=404, detail=f"Source {source_id} not found")
    
    def task():
        count = reprocess_single_source(
            source_id=src['id'],
            title=src['title'] or f"Source {src['id']}",
            content_text=src['content'] or '',
            file_path=src['file_path'] or '',
            source_type=src['type'] or 'text'
        )
        # Reload embeddings
        global embeddings, chunk_ids, source_ids
        cnx2 = mysql.connector.connect(**DB)
        cur2 = cnx2.cursor(dictionary=True)
        cur2.execute("SELECT id, rag_source_id, embedding FROM rag_chunks WHERE embedding IS NOT NULL")
        rows = cur2.fetchall()
        cur2.close()
        cnx2.close()
        if rows:
            chunk_ids = np.array([r["id"] for r in rows])
            source_ids = np.array([r["rag_source_id"] for r in rows])
            embeddings = normalize_embedding_matrix(
                np.array([json.loads(r["embedding"]) for r in rows], dtype=np.float32)
            )
    
    background_tasks.add_task(task)
    return {
        "status": "reprocessing_started",
        "source_id": source_id,
        "title": src['title'],
        "message": f"Reprocessing source {source_id} in background."
    }

@app.post("/index-source/{source_id}")
def index_source(source_id: int, background_tasks: BackgroundTasks):
    """Read the manually reviewed .txt file, embed it, and insert into DB."""
    if model is None:
        raise HTTPException(
            status_code=503,
            detail="Dense indexing is disabled. Enable DENSE_SEARCH_ENABLED and install requirements-dense.txt.",
        )

    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)
    cur.execute("SELECT id, title FROM rag_sources WHERE id = %s AND deleted_at IS NULL", (source_id,))
    src = cur.fetchone()
    cur.close()
    cnx.close()
    
    if not src:
        raise HTTPException(status_code=404, detail=f"Source {source_id} not found")
        
    title = src['title'] or f"Source {source_id}"
    
    import glob
    txt_pattern = os.path.join(RAG_DIR, f"source_{source_id}_*.txt")
    txt_files = glob.glob(txt_pattern)
    
    if not txt_files:
        raise HTTPException(status_code=404, detail=f"Reviewed TXT file not found for source {source_id}")
        
    txt_path = txt_files[0]
        
    def task():
        try:
            from import_chunks import parse_reviewed_txt
            _, parsed_title, articles = parse_reviewed_txt(txt_path)
            
            if not articles:
                print(f"  No articles found in {txt_filename}")
                return
                
            # Delete old chunks
            cnx = mysql.connector.connect(**DB)
            cur = cnx.cursor()
            cur.execute("DELETE FROM rag_chunks WHERE rag_source_id = %s", (source_id,))
            cnx.commit()
            
            new_embeddings = []
            new_ids = []
            new_source_ids = []
            
            for i, art in enumerate(articles):
                content = art['content']
                if not content.strip():
                    continue
                
                passage = f"passage: {title} - مادة {art['article_number']}\n{content}"
                emb = model.encode([passage], normalize_embeddings=True)[0]
                emb_json = json.dumps(emb.tolist())
                
                token_count = art.get('word_count', len(content.split()))
                meta_dict = {
                    "article_number": str(art['article_number']), 
                    "law_name": title,
                    "part": art.get('part'),
                    "chapter": art.get('chapter'),
                    "section": art.get('section')
                }
                meta = json.dumps(meta_dict)
                
                sql = """
                INSERT INTO rag_chunks (rag_source_id, law_name, article_number, part, chapter, section, content, meta_data, chunk_index, token_count, embedding, created_at, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), NOW())
                """
                val = (source_id, title, str(art['article_number']), art.get('part'), art.get('chapter'), art.get('section'), content, meta, i, token_count, emb_json)
                cur.execute(sql, val)
                
                new_embeddings.append(emb)
                new_ids.append(cur.lastrowid)
                new_source_ids.append(source_id)
            
            cnx.commit()
            cur.close()
            cnx.close()
            
            # Reload embeddings
            global embeddings, chunk_ids, source_ids
            cnx2 = mysql.connector.connect(**DB)
            cur2 = cnx2.cursor(dictionary=True)
            cur2.execute("SELECT id, rag_source_id, embedding FROM rag_chunks WHERE embedding IS NOT NULL")
            rows = cur2.fetchall()
            cur2.close()
            cnx2.close()
            if rows:
                chunk_ids = np.array([r["id"] for r in rows])
                source_ids = np.array([r["rag_source_id"] for r in rows])
                embeddings = normalize_embedding_matrix(
                    np.array([json.loads(r["embedding"]) for r in rows], dtype=np.float32)
                )
                
            # Update status
            cnx3 = mysql.connector.connect(**DB)
            cur3 = cnx3.cursor()
            cur3.execute("UPDATE rag_sources SET status = 'completed' WHERE id = %s", (source_id,))
            cnx3.commit()
            cur3.close()
            cnx3.close()
            
            print(f"✅ Source {source_id} successfully indexed from TXT.")
            
        except Exception as e:
            print(f"❌ Error indexing source {source_id}: {str(e)}")
            import traceback
            traceback.print_exc()
            
    background_tasks.add_task(task)
    return {
        "status": "indexing_started",
        "source_id": source_id,
        "message": f"Indexing source {source_id} in background."
    }


@app.post("/index-direct/{source_id}")
def index_direct(source_id: int):
    if model is None:
        raise HTTPException(
            status_code=503,
            detail="Dense indexing is disabled. Enable DENSE_SEARCH_ENABLED and install requirements-dense.txt.",
        )

    """Index a source directly from its text content in DB as a single chunk."""
    cnx = mysql.connector.connect(**DB)
    cur = cnx.cursor(dictionary=True)
    cur.execute("SELECT id, title, content, meta_data FROM rag_sources WHERE id = %s", (source_id,))
    src = cur.fetchone()
    
    if not src or not src['content']:
        cur.close()
        cnx.close()
        raise HTTPException(status_code=404, detail="Source or content not found")
        
    title = src['title'] or f"Source {source_id}"
    content = src['content'].strip()
    
    # Optional logic: if the content is extremely huge, chunk it. 
    # For rulings, they usually fit in one chunk and we want to retrieve the whole ruling.
    
    # Delete old chunks
    cur.execute("DELETE FROM rag_chunks WHERE rag_source_id = %s", (source_id,))
    cnx.commit()
    
    passage = f"passage: {title}\n{content}"
    # Truncate to reasonable token limit approximation if necessary, although e5 handles up to 512 normally
    emb = model.encode([passage], normalize_embeddings=True)[0]
    emb_json = json.dumps(emb.tolist())
    
    word_count = len(content.split())
    meta_dict = {}
    if src['meta_data']:
        try:
            meta_dict = json.loads(src['meta_data'])
        except:
            pass
            
    meta_json = json.dumps(meta_dict, ensure_ascii=False)
    
    cur.execute('''
        INSERT INTO rag_chunks (rag_source_id, article_number, content, chunk_index, token_count, embedding, meta_data, created_at, updated_at)
        VALUES (%s, %s, %s, %s, %s, %s, %s, NOW(), NOW())
    ''', (source_id, 'كامل الحكم', content, 0, word_count, emb_json, meta_json))
    
    chunk_id = cur.lastrowid
    
    cur.execute("UPDATE rag_sources SET status = 'completed' WHERE id = %s", (source_id,))
    cnx.commit()
    cur.close()
    cnx.close()
    
    # Update memory cache
    global chunk_ids, source_ids, embeddings, texts, tfidf_matrix, vectorizer
    if chunk_ids is not None and len(chunk_ids) > 0:
        new_chunk_ids = np.append(chunk_ids, chunk_id)
        new_source_ids = np.append(source_ids, source_id)
        new_embeddings = np.vstack([embeddings, emb])
        new_texts = list(texts)
        new_texts.append(content)
    else:
        new_chunk_ids = np.array([chunk_id])
        new_source_ids = np.array([source_id])
        new_embeddings = np.array([emb])
        new_texts = [content]
        
    # Rebuild TFIDF safely to avoid race conditions during concurrent searches
    new_vectorizer = TfidfVectorizer(analyzer='word', ngram_range=(1, 2))
    new_tfidf_matrix = new_vectorizer.fit_transform(new_texts)
    
    # ATOMIC ASSIGNMENT of all globals at once
    chunk_ids = new_chunk_ids
    source_ids = new_source_ids
    embeddings = new_embeddings
    texts = new_texts
    vectorizer = new_vectorizer
    tfidf_matrix = new_tfidf_matrix
    
    return {"status": "success", "chunk_id": chunk_id, "message": f"Source {source_id} indexed correctly."}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=API_HOST, port=API_PORT)