import pickle
import os
import re
import json
import torch
import numpy as np
from sentence_transformers import SentenceTransformer, util
import gradio as gr
from fastapi import FastAPI
from collections import defaultdict
import math

# =========================
# Arabic Text Normalization
# =========================

def normalize_arabic(text: str) -> str:
    """Normalize Arabic text for consistent matching."""
    if not text:
        return ""
    # Normalize Alef variants → ا
    text = re.sub(r'[أإآاٱ]', 'ا', text)
    # Normalize Yeh variants → ي
    text = re.sub(r'[يىئ]', 'ي', text)
    # Normalize Waw variants
    text = re.sub(r'[ؤو]', 'و', text)
    # Normalize Heh variants
    text = re.sub(r'[ةه]', 'ه', text)
    # Remove Tatweel (kashida)
    text = re.sub(r'ـ', '', text)
    # Remove Tashkeel (diacritics/harakat)
    text = re.sub(r'[\u064B-\u065F\u0610-\u061A\u06D6-\u06DC\u06DF-\u06E4\u06E7\u06E8\u06EA-\u06ED]', '', text)
    # Normalize whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    return text


def tokenize_arabic(text: str) -> list:
    """Tokenize Arabic text into meaningful tokens."""
    normalized = normalize_arabic(text)
    # Split on whitespace and punctuation
    tokens = re.findall(r'[\u0600-\u06FF]+', normalized)
    # Filter out very short tokens (1-2 chars) unless they're legal terms
    legal_short_terms = {'لا', 'في', 'من', 'مع', 'إن', 'ان', 'قد', 'بل', 'لم', 'لن'}
    tokens = [t for t in tokens if len(t) > 2 or t in legal_short_terms]
    return tokens


# =========================
# Arabic Stemmer (Light Stemming)
# =========================

class ArabicLightStemmer:
    """
    Light Arabic stemmer that strips common prefixes and suffixes
    to improve recall in keyword matching.
    """
    PREFIXES = ['ال', 'وال', 'بال', 'كال', 'فال', 'للـ', 'لل', 'و', 'ف', 'ب', 'ك', 'ل']
    SUFFIXES = ['ين', 'ون', 'ان', 'ات', 'ية', 'يه', 'ها', 'هم', 'هن', 'كم',
                'كن', 'نا', 'تم', 'تن', 'وا', 'ه', 'ا', 'ة', 'ي', 'ن']

    def stem(self, word: str) -> str:
        word = normalize_arabic(word)
        # Strip prefixes (longest first)
        for prefix in sorted(self.PREFIXES, key=len, reverse=True):
            if word.startswith(prefix) and len(word) - len(prefix) >= 3:
                word = word[len(prefix):]
                break
        # Strip suffixes (longest first)
        for suffix in sorted(self.SUFFIXES, key=len, reverse=True):
            if word.endswith(suffix) and len(word) - len(suffix) >= 3:
                word = word[:-len(suffix)]
                break
        return word

stemmer = ArabicLightStemmer()


# =========================
# BM25 Implementation
# =========================

class BM25:
    """
    BM25 ranking algorithm adapted for Arabic text with light stemming.
    k1 and b tuned for short Arabic legal snippets.
    """
    def __init__(self, corpus: list, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.corpus = corpus
        self.N = len(corpus)
        self.tokenized_corpus = []
        self.doc_lengths = []
        self.df = defaultdict(int)      # document frequency
        self.idf = {}
        self.avgdl = 0.0

        self._build_index()

    def _tokenize(self, text: str) -> list:
        tokens = tokenize_arabic(text)
        return [stemmer.stem(t) for t in tokens]

    def _build_index(self):
        total_len = 0
        for doc in self.corpus:
            tokens = self._tokenize(doc)
            self.tokenized_corpus.append(tokens)
            self.doc_lengths.append(len(tokens))
            total_len += len(tokens)
            for token in set(tokens):
                self.df[token] += 1

        self.avgdl = total_len / self.N if self.N > 0 else 1.0

        # IDF with smoothing
        for term, df in self.df.items():
            self.idf[term] = math.log((self.N - df + 0.5) / (df + 0.5) + 1)

    def get_scores(self, query: str) -> np.ndarray:
        query_tokens = self._tokenize(query)
        scores = np.zeros(self.N)

        for term in query_tokens:
            if term not in self.idf:
                continue
            idf_val = self.idf[term]
            for i, doc_tokens in enumerate(self.tokenized_corpus):
                tf = doc_tokens.count(term)
                if tf == 0:
                    continue
                dl = self.doc_lengths[i]
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avgdl)
                scores[i] += idf_val * (numerator / denominator)

        return scores

    def get_top_n(self, query: str, n: int = 20) -> list:
        scores = self.get_scores(query)
        top_indices = np.argsort(scores)[::-1][:n]
        return [(int(idx), float(scores[idx])) for idx in top_indices]


# =========================
# Load JSON Data
# =========================

with open('/home/user/app/alldec.json', 'r') as f:
    datal = json.load(f)

def extract_text_by_number(num_to_find):
    if not str(num_to_find).isdigit():
        return "ادخل فقط رقم القرار"
    target_num = str(num_to_find).strip()
    for data in datal:
        nested_data = data.get('data', {})
        num = nested_data.get('num')
        if num is not None and str(num).strip() == target_num:
            return data.get('text', 'لا يوجد نص لهذا القرار')
    return "لم نتمكن من ايجاد القرار"


# =========================
# Load Embeddings & Build BM25
# =========================

with open("/home/user/app/embmmn7.obj", "rb") as fileobj:
    corpus_embeddings, corpus = pickle.load(fileobj)

embedder = SentenceTransformer("ramdane/jurimodel")

print("Building BM25 index for Arabic corpus...")
bm25_index = BM25(corpus)
print(f"BM25 index built for {len(corpus)} documents.")


# =========================
# Hybrid Search Core
# =========================

def hybrid_search(
    query: str,
    corpus_subset: list,
    embeddings_subset,
    subset_indices: list,
    top_k: int = 20,
    semantic_weight: float = 0.65,
    bm25_weight: float = 0.35,
    min_threshold: float = 0.0
) -> list:
    """
    Hybrid retrieval combining BM25 keyword scores with semantic similarity.
    Strategy:
    - BM25 captures exact legal term matches (critical in Arabic law).
    - Semantic search captures paraphrase / concept similarity.
    - Scores are min-max normalized before weighted fusion (Reciprocal Rank
      Fusion is also optionally applied for robustness).
    Returns list of (subset_index, fused_score) sorted descending.
    """
    n = len(corpus_subset)
    if n == 0:
        return []

    # --- Semantic scores ---
    query_embedding = embedder.encode(
        normalize_arabic(query), convert_to_tensor=True
    )
    if isinstance(embeddings_subset, list):
        if isinstance(embeddings_subset[0], torch.Tensor):
            emb_tensor = torch.stack(embeddings_subset)
        else:
            emb_tensor = torch.tensor(embeddings_subset)
    else:
        emb_tensor = embeddings_subset

    sem_hits = util.semantic_search(query_embedding, emb_tensor, top_k=n)[0]
    sem_scores = np.zeros(n)
    for hit in sem_hits:
        sem_scores[hit['corpus_id']] = hit['score']

    # --- BM25 scores on subset ---
    bm25_scores_full = bm25_index.get_scores(query)
    bm25_scores = np.array([bm25_scores_full[idx] for idx in subset_indices])

    # --- Min-max normalization ---
    def minmax(arr):
        mn, mx = arr.min(), arr.max()
        if mx - mn < 1e-9:
            return np.zeros_like(arr)
        return (arr - mn) / (mx - mn)

    sem_norm = minmax(sem_scores)
    bm25_norm = minmax(bm25_scores)

    # --- Adaptive weighting ---
    # If query is very short (1-2 words), boost BM25 weight for exact matching.
    query_words = query.split()
    if len(query_words) <= 2:
        effective_sem = 0.45
        effective_bm25 = 0.55
    elif len(query_words) <= 5:
        effective_sem = semantic_weight
        effective_bm25 = bm25_weight
    else:
        # Longer queries benefit more from semantic search
        effective_sem = 0.75
        effective_bm25 = 0.25

    fused = effective_sem * sem_norm + effective_bm25 * bm25_norm

    # Filter by threshold
    results = [
        (i, float(fused[i]))
        for i in range(n)
        if fused[i] >= min_threshold
    ]
    results.sort(key=lambda x: x[1], reverse=True)
    return results[:top_k]


# =========================
# Chamber Definitions
# =========================

CHAMBRES = [
    {
        "id": 0,
        "name": "العقارية",
        "phrases": ["العقارية", "الغرفة العقارية", "عقاري", "العقار"]
    },
    {
        "id": 1,
        "name": "المدنية",
        "phrases": ["الغرفة المدنية", "غرفة مدنية", "المدنية"]
    },
    {
        "id": 2,
        "name": "الاجتماعية",
        "phrases": ["الغرفة الاجتماعية", "الإجتماعية", "الاجتماعية"]
    },
    {
        "id": 3,
        "name": "الجنح",
        "phrases": ["الجنح", "جنح", "مخالفات", "جنح ومخالفات"]
    },
    {
        "id": 4,
        "name": "الجنائية",
        "phrases": ["جنائية", "الجنائية", "الغرفة الجنائية"]
    },
    {
        "id": 5,
        "name": "شؤون الأسرة",
        "phrases": ["شؤون الأسرة", "شؤون الاسرة", "الاحوال الشخصية",
                    "الأحوال الشخصية", "المواريث", "الأسرة"]
    },
    {
        "id": 6,
        "name": "التجارية",
        "phrases": ["التجارية", "تجارية", "البحرية", "الغرفة التجارية"]
    },
    {
        "id": 7,
        "name": "الإدارية / مجلس الدولة",
        "phrases": ["محكمة التنازع", "غرفة ادارية", "مجلس الدولة",
                    "الغرفة الادارية", "إدارية"]
    },
]


def get_chamber_subset(chamber_id: int):
    """
    Filter corpus to documents matching the selected chamber.
    Returns (subset_corpus, subset_embeddings, subset_original_indices).
    """
    if chamber_id < 0 or chamber_id >= len(CHAMBRES):
        return corpus, corpus_embeddings, list(range(len(corpus)))

    phrases = CHAMBRES[chamber_id]["phrases"]
    normalized_phrases = [normalize_arabic(p) for p in phrases]

    subset_corpus = []
    subset_embeddings = []
    subset_indices = []

    for i, text in enumerate(corpus):
        norm_text = normalize_arabic(text)
        if any(
            re.search(r'\b' + re.escape(p) + r'\b', norm_text, re.IGNORECASE)
            for p in normalized_phrases
        ):
            subset_corpus.append(text)
            subset_embeddings.append(corpus_embeddings[i])
            subset_indices.append(i)

    return subset_corpus, subset_embeddings, subset_indices


# =========================
# Main Search Function
# =========================

def preprocess_query(query: str) -> str:
    """Clean and normalize the search query."""
    query = re.sub(r'\s+', ' ', query).strip()
    return query


def showrs(query: str, rank: int, chamber_id: int = -1) -> str:
    """
    Perform hybrid search and return the result at the requested rank.
    Args:
        query:      The Arabic search query.
        rank:       Zero-based rank of the desired result (0 = best match).
        chamber_id: -1 = all chambers; 0-7 = specific chamber.
    """
    clean_query = preprocess_query(query)
    if not clean_query:
        return "الرجاء إدخال كلمات البحث"

    # ----- Subset selection -----
    if chamber_id == -1:
        # All documents
        sub_corpus = corpus
        sub_embeddings = corpus_embeddings
        sub_indices = list(range(len(corpus)))
    elif 0 <= chamber_id < len(CHAMBRES):
        sub_corpus, sub_embeddings, sub_indices = get_chamber_subset(chamber_id)
        if not sub_corpus:
            return "لم نتمكن من العثور على وثائق في هذه الغرفة"
    else:
        return "خطأ في اختيار الغرفة"

    # ----- Hybrid retrieval -----
    top_k = max(rank + 1, 20)
    results = hybrid_search(
        query=clean_query,
        corpus_subset=sub_corpus,
        embeddings_subset=sub_embeddings,
        subset_indices=sub_indices,
        top_k=top_k,
        min_threshold=0.0
    )

    if rank < len(results):
        local_idx, score = results[rank]
        return sub_corpus[local_idx]

    return (
        "لم نتمكن من إيجاد النتيجة — "
        "إما لعدم وجود الاجتهاد أو لعدم كتابة جملة بحث مناسبة"
    )


# =========================
# Gradio Interface Handler
# =========================

def greet_user(query: str, rank, chamber):
    try:
        rank = int(rank)
        chamber = int(chamber)

        if chamber == -2:
            return extract_text_by_number(query)

        return showrs(query, rank, chamber)
    except Exception as e:
        return f"خطأ في المدخلات: {str(e)}"


# =========================
# Gradio UI
# =========================

with gr.Blocks(
    title="البحث في الاجتهادات القضائية",
    theme=gr.themes.Base(
        primary_hue="blue",
        font=[gr.themes.GoogleFont("Tajawal"), "sans-serif"]
    ),
    css="""
        .rtl-text { direction: rtl; text-align: right; }
        textarea, input { direction: rtl !important; text-align: right !important; }
    """
) as demo:

    gr.Markdown(
        """
        <div style="direction:rtl; text-align:center; padding:10px">
        <h2>🔍 البحث في الاجتهادات القضائية</h2>
        <p>يستخدم هذا النظام البحث الهجين: <b>BM25</b> (تطابق الكلمات المفتاحية) + <b>البحث الدلالي</b> (تشابه المعنى)</p>
        </div>
        """,
        elem_classes=["rtl-text"]
    )

    with gr.Row():
        with gr.Column(scale=3):
            query = gr.Textbox(
                label="كلمات البحث",
                placeholder="مثال: فسخ عقد البيع لعدم الدفع...",
                lines=2,
                elem_classes=["rtl-text"]
            )
        with gr.Column(scale=1):
            order = gr.Number(label="الترتيب (ابدأ من 0)", value=0, precision=0)
            chamber = gr.Dropdown(
                label="الغرفة",
                choices=[
                    ("-1 — كل الغرف", -1),
                    ("0 — العقارية", 0),
                    ("1 — المدنية", 1),
                    ("2 — الاجتماعية", 2),
                    ("3 — الجنح", 3),
                    ("4 — الجنائية", 4),
                    ("5 — شؤون الأسرة", 5),
                    ("6 — التجارية", 6),
                    ("7 — الإدارية / مجلس الدولة", 7),
                    ("-2 — البحث برقم القرار", -2),
                ],
                value=-1,
                type="value"
            )

    btn = gr.Button("🔎 بحث", variant="primary")
    result = gr.TextArea(
        label="النتيجة",
        lines=15,
        elem_classes=["rtl-text"]
    )

    gr.Markdown(
        """
        <div style="direction:rtl; font-size:0.85em; color:#666; margin-top:8px">
        💡 <b>تلميح:</b> لرؤية نتائج متعددة، غيّر رقم الترتيب (0 = أفضل نتيجة، 1 = الثانية…).
        للبحث بواسطة رقم القرار، اختر الخيار <b>-2</b> وأدخل رقم القرار.
        </div>
        """
    )

    btn.click(
        fn=greet_user,
        inputs=[query, order, chamber],
        outputs=result,
        api_name="predict"
    )
    query.submit(
        fn=greet_user,
        inputs=[query, order, chamber],
        outputs=result
    )


# =========================
# FastAPI + Compatibility Endpoint
# =========================

fastapi_app = FastAPI(title="Arabic Legal Search API")


@fastapi_app.post("/run/predict")
async def old_api(data: dict):
    """Backward-compatible endpoint for existing clients."""
    try:
        inputs = data["data"]
        output = greet_user(inputs[0], inputs[1], inputs[2])
        return {"data": [output]}
    except Exception as e:
        return {"error": str(e)}


@fastapi_app.get("/health")
async def health():
    return {
        "status": "ok",
        "corpus_size": len(corpus),
        "bm25_indexed": True
    }


# Mount Gradio
app = gr.mount_gradio_app(fastapi_app, demo, path="/")

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