# app.py - Hybrid RAG (FAISS semantic + BM25, fused with RRF) for multi-format documents.
# Handles: PDFs (text + tables, OCR for scanned pages), DOCX (with tables), PPTX, XLSX/CSV, TXT/MD. Arabic + English.
import glob
import io
import os
import re
import shutil
import subprocess
import tempfile
import unicodedata
from collections import defaultdict

import faiss
import fitz  # PyMuPDF
import numpy as np
import pandas as pd
import streamlit as st
from docx import Document
from langchain_groq import ChatGroq
from PIL import Image, ImageOps
from pptx import Presentation
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

# ========================= CONSTANTS =========================
EMBED_MODEL_NAME = "paraphrase-multilingual-MiniLM-L12-v2"  # Arabic + English, 384-dim
LLM_MODEL_NAME = "openai/gpt-oss-120b"
CHUNK_WORDS = 80          # stays under the embedding model's 128-token limit
CHUNK_OVERLAP = 15
TABLE_CHUNK_WORDS = 60    # table/spreadsheet rows are token-heavy, so smaller groups
RRF_K = 60
OCR_DPI = 200             # Tesseract works best when text is ~30px tall; 400 hurt large poster text. Raise for tiny fax text.
OCR_PSMS = (3, 11)        # 3 = auto page layout, 11 = sparse text (finds scattered words on posters)
EXTRACT_VERSION = "ocr-cli-v3"  # bump to invalidate st.cache_data (it does not track helper functions)
OCR_LANGS = "ara+eng"
OCR_FALLBACK_LANGS = "eng"  # tried first (Arabic data can hurt English-only pages); also the fallback if Arabic data is missing
OCR_CONTENT_THRESHOLD = 200  # gray level below which a pixel counts as "content" (0-255)
OCR_CLIP_PAD = 8             # points of padding around detected content

# ========================= TEXT NORMALIZATION =========================
DIACRITICS_RE = re.compile(r"[ؐ-ًؚ-ٰٟۖ-ۭ]")
DIGIT_MAP = str.maketrans("٠١٢٣٤٥٦٧٨٩۰۱۲۳۴۵۶۷۸۹", "01234567890123456789")

def light_clean(text):
    text = unicodedata.normalize("NFKC", text)
    text = DIACRITICS_RE.sub("", text)
    return text.replace("ـ", "")

def normalize(text):
    text = light_clean(text).translate(DIGIT_MAP)
    text = re.sub("[إأآٱ]", "ا", text)
    text = text.replace("ى", "ي").replace("ة", "ه")
    return text.lower()

def tokenize(text):
    tokens = re.findall(r"[\w\+#]+", normalize(text))
    return [t[2:] if len(t) > 4 and t.startswith("ال") else t for t in tokens]

# ========================= EXTRACTION =========================
def make_block(text, source, loc, kind):
    return {"text": text, "source": source, "loc": loc, "kind": kind}

def pdf_tables(page):
    out = []
    try:
        for t in page.find_tables().tables:
            lines = []
            for r in t.extract():
                cells = [(c or "").replace("\n", " ").strip() for c in r]
                if any(cells):
                    lines.append(" | ".join(cells))
            if len(lines) >= 2:
                out.append("\n".join(lines))
    except Exception:
        pass
    return out

def _set_tessdata():
    """Point PyMuPDF at Tesseract's language data (path differs by distro/version)."""
    if os.environ.get("TESSDATA_PREFIX"):
        return
    for p in glob.glob("/usr/share/tesseract-ocr/*/tessdata") + ["/usr/share/tessdata"]:
        if os.path.isdir(p):
            os.environ["TESSDATA_PREFIX"] = p
            return

def _content_clip(page):
    """Bounding box of non-white content, so large blank margins don't shrink the text during OCR.
    Returns None if the page is blank, or the full page rect if the content already fills it."""
    try:
        pix = page.get_pixmap(dpi=72, colorspace=fitz.csGRAY)  # at 72 dpi, 1 pixel = 1 point
        arr = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width)
        ys, xs = np.where(arr < OCR_CONTENT_THRESHOLD)
        if xs.size == 0:
            return None
        r = page.rect
        clip = fitz.Rect(
            r.x0 + xs.min() - OCR_CLIP_PAD, r.y0 + ys.min() - OCR_CLIP_PAD,
            r.x0 + xs.max() + 1 + OCR_CLIP_PAD, r.y0 + ys.max() + 1 + OCR_CLIP_PAD,
        ) & r
        if clip.is_empty or clip.width < 20 or clip.height < 20:
            return r
        return clip
    except Exception:
        return page.rect

def _render_gray(page, clip):
    pix = page.get_pixmap(dpi=OCR_DPI, clip=clip, colorspace=fitz.csGRAY)
    return Image.frombytes("L", (pix.width, pix.height), pix.samples)

def _tesseract_cli(img, lang, psm):
    """Run the tesseract binary directly on a PIL image. Raises if it is missing or fails."""
    with tempfile.TemporaryDirectory() as tmp:
        path = os.path.join(tmp, "page.png")
        img.save(path)
        env = dict(os.environ)
        r = subprocess.run(["tesseract", path, "stdout", "-l", lang, "--psm", str(psm)],
                           capture_output=True, text=True, timeout=180, env=env)
        if r.returncode != 0:
            raise RuntimeError((r.stderr or "tesseract failed").strip()[:300])
        return unicodedata.normalize("NFKC", r.stdout).strip()

def _tesseract_pymupdf(img, lang):
    """Fallback when the tesseract binary is not on PATH: PyMuPDF's built-in OCR."""
    w, h = img.size
    pix = fitz.Pixmap(fitz.csGRAY, w, h, img.tobytes(), False)
    pix.set_dpi(OCR_DPI, OCR_DPI)
    ocr_doc = fitz.open("pdf", pix.pdfocr_tobytes(language=lang))
    try:
        return unicodedata.normalize("NFKC", ocr_doc[0].get_text("text") or "").strip()
    finally:
        ocr_doc.close()

def ocr_page(page):
    """OCR a scanned page. Returns (text, error).
    Tries English first, then Arabic+English, each on the raw render and a contrast-stretched
    inverted copy (for white-on-dark banners), with two page-segmentation modes, and merges the
    unique lines. error is None unless every attempt raised."""
    _set_tessdata()
    clip = _content_clip(page)
    if clip is None:  # truly blank page: nothing to OCR, and not an error
        return "", None, "page looked blank (no pixels darker than the content threshold)"
    try:
        raw = _render_gray(page, clip)
    except Exception as e:
        return "", f"could not render page: {e}", ""
    inv = ImageOps.invert(ImageOps.autocontrast(raw, cutoff=2))
    use_cli = shutil.which("tesseract") is not None
    errors, lines, seen = [], [], set()
    for lang in dict.fromkeys((OCR_FALLBACK_LANGS, OCR_LANGS)):  # de-duplicated, order kept
        for img in (raw, inv):
            for psm in (OCR_PSMS if use_cli else (None,)):
                try:
                    txt = _tesseract_cli(img, lang, psm) if use_cli else _tesseract_pymupdf(img, lang)
                except Exception as e:
                    errors.append(f"[{lang}] {e}")
                    continue
                for ln in txt.splitlines():
                    key = re.sub(r"\W+", "", ln.lower())
                    if len(key) >= 3 and key not in seen:
                        seen.add(key)
                        lines.append(ln.strip())
    info = f"engine={'tesseract CLI' if use_cli else 'PyMuPDF OCR'}, image={raw.size[0]}x{raw.size[1]}px, unique lines={len(lines)}"
    if lines:
        return "\n".join(lines), None, info
    return "", ("; ".join(errors) if errors else None), info

def extract_pdf(name, data):
    blocks = []
    stats = {"pages": 0, "tables": 0, "notes": []}
    doc = fitz.open(stream=data, filetype="pdf")
    empty_pages, ocr_pages, ocr_errors, ocr_infos = [], [], [], []
    for i, page in enumerate(doc, 1):
        stats["pages"] += 1
        loc = f"p.{i}"
        text = unicodedata.normalize("NFKC", page.get_text("text") or "").strip()
        for tbl in pdf_tables(page):
            blocks.append(make_block(tbl, name, loc, "table"))
            stats["tables"] += 1
        if len(text) < 50:
            ocr_text, ocr_err, ocr_info = ocr_page(page)
            ocr_infos.append((i, len(ocr_text), ocr_info))
            if ocr_err:
                ocr_errors.append((i, ocr_err))
            if len(ocr_text) > len(text):
                text = ocr_text
                ocr_pages.append(i)
        if len(text) < 50:
            empty_pages.append(i)
        if text:
            blocks.append(make_block(text, name, loc, "text"))
    if ocr_pages:
        shown = ", ".join(map(str, ocr_pages[:15])) + ("..." if len(ocr_pages) > 15 else "")
        stats["notes"].append(f"OCR used on page(s) {shown}. Text may contain recognition errors.")
    if ocr_errors:
        # Report the first error once; it usually repeats for every page.
        first_pages = ", ".join(str(p) for p, _ in ocr_errors[:15]) + ("..." if len(ocr_errors) > 15 else "")
        stats["notes"].append(
            f"OCR could not run on page(s) {first_pages}: {ocr_errors[0][1]}. "
            "Check that Tesseract and its language data (eng, ara) are installed."
        )
    if empty_pages:
        dbg = "; ".join(f"p.{p}: {n} chars, {inf}" for p, n, inf in ocr_infos if p in empty_pages[:3])
        if dbg:
            stats["notes"].append(f"OCR details: {dbg}")
        shown = ", ".join(map(str, empty_pages[:15])) + ("..." if len(empty_pages) > 15 else "")
        stats["notes"].append(
            f"Little or no readable text found on page(s) {shown}"
            + ("." if ocr_errors else ", even after OCR.")
        )
    return blocks, stats

def extract_docx(name, data):
    doc = Document(io.BytesIO(data))
    blocks, tables = [], 0
    text = "\n".join(p.text for p in doc.paragraphs if p.text.strip())
    if text.strip():
        blocks.append(make_block(text, name, None, "text"))
    for table in doc.tables:
        lines = []
        for row in table.rows:
            cells = [c.text.replace("\n", " ").strip() for c in row.cells]
            if any(cells):
                lines.append(" | ".join(cells))
        if len(lines) >= 2:
            blocks.append(make_block("\n".join(lines), name, None, "table"))
            tables += 1
    return blocks, {"pages": 1, "tables": tables,
                    "notes": ["Images embedded inside the DOCX are not read."]}

def extract_pptx(name, data):
    prs = Presentation(io.BytesIO(data))
    blocks, tables = [], 0
    for i, slide in enumerate(prs.slides, 1):
        loc = f"slide {i}"
        texts = []
        for shape in slide.shapes:
            if shape.has_text_frame and shape.text_frame.text.strip():
                texts.append(shape.text_frame.text.strip())
            if getattr(shape, "has_table", False) and shape.has_table:
                lines = [" | ".join(c.text.strip() for c in row.cells) for row in shape.table.rows]
                if len(lines) >= 2:
                    blocks.append(make_block("\n".join(lines), name, loc, "table"))
                    tables += 1
        if slide.has_notes_slide and slide.notes_slide.notes_text_frame.text.strip():
            texts.append("Notes: " + slide.notes_slide.notes_text_frame.text.strip())
        if texts:
            blocks.append(make_block("\n".join(texts), name, loc, "text"))
    return blocks, {"pages": len(prs.slides), "tables": tables,
                    "notes": ["Pictures and charts on slides are not read."]}

def df_to_rows(df):
    df = df.dropna(how="all").dropna(axis=1, how="all").fillna("")
    headers = [str(c) for c in df.columns]
    lines = []
    for _, row in df.iterrows():
        cells = []
        for h, v in zip(headers, row):
            v = str(v).strip()
            if v:
                cells.append(v if h.startswith("Unnamed") else f"{h}: {v}")
        if cells:
            lines.append(" | ".join(cells))
    return "\n".join(lines)

def extract_spreadsheet(name, data):
    blocks = []
    if name.lower().endswith(".csv"):
        df = None
        for enc in ("utf-8-sig", "cp1256", "latin-1"):
            try:
                df = pd.read_csv(io.BytesIO(data), encoding=enc)
                break
            except Exception:
                continue
        sheets = {"": df} if df is not None else {}
    else:
        sheets = pd.read_excel(io.BytesIO(data), sheet_name=None)
    for sheet, df in sheets.items():
        text = df_to_rows(df)
        if text:
            blocks.append(make_block(text, name, f"sheet {sheet}" if sheet else None, "rows"))
    return blocks, {"pages": len(sheets), "tables": len(blocks), "notes": []}

def extract_text_file(name, data):
    text = data.decode("utf-8", errors="ignore")
    blocks = [make_block(text, name, None, "text")] if text.strip() else []
    return blocks, {"pages": 1, "tables": 0, "notes": []}

@st.cache_data(show_spinner=False)
def extract_file(name, data, version=EXTRACT_VERSION):
    lower = name.lower()
    if lower.endswith(".pdf"):
        return extract_pdf(name, data)
    if lower.endswith(".docx"):
        return extract_docx(name, data)
    if lower.endswith(".pptx"):
        return extract_pptx(name, data)
    if lower.endswith((".xlsx", ".csv")):
        return extract_spreadsheet(name, data)
    if lower.endswith((".txt", ".md")):
        return extract_text_file(name, data)
    return [], {"pages": 0, "tables": 0, "notes": ["Unsupported file type."]}

# ========================= CHUNKING =========================
def make_chunk(block, text):
    return {"text": text, "source": block["source"], "loc": block["loc"], "kind": block["kind"]}

def chunk_blocks(blocks):
    chunks = []
    step = max(1, CHUNK_WORDS - CHUNK_OVERLAP)
    for b in blocks:
        if b["kind"] in ("table", "rows"):
            lines = [l for l in b["text"].split("\n") if l.strip()]
            header = lines[0] if b["kind"] == "table" else None
            body = lines[1:] if b["kind"] == "table" else lines
            group, wc = [], 0
            for line in body:
                lw = len(line.split())
                if group and wc + lw > TABLE_CHUNK_WORDS:
                    chunks.append(make_chunk(b, "\n".join(([header] if header else []) + group)))
                    group, wc = [], 0
                group.append(line)
                wc += lw
            if group:
                chunks.append(make_chunk(b, "\n".join(([header] if header else []) + group)))
        else:
            words = b["text"].split()
            for i in range(0, len(words), step):
                piece = " ".join(words[i:i + CHUNK_WORDS])
                if len(piece.strip()) > 20:
                    chunks.append(make_chunk(b, piece))
    return chunks

def cite_label(c):
    return f"{c['source']}, {c['loc']}" if c["loc"] else c["source"]

# ========================= INDEX + HYBRID SEARCH =========================
@st.cache_resource(show_spinner="Loading embedding model...")
def load_embedder():
    return SentenceTransformer(EMBED_MODEL_NAME, cache_folder="/tmp/model_cache", device="cpu")

def build_index(chunks, embedder):
    texts = [light_clean(c["text"]) for c in chunks]
    emb = embedder.encode(texts, batch_size=32, convert_to_numpy=True,
                          normalize_embeddings=True, show_progress_bar=False).astype("float32")
    faiss_index = faiss.IndexFlatIP(emb.shape[1])  # inner product on normalized vectors = cosine
    faiss_index.add(emb)
    bm25 = BM25Okapi([tokenize(c["text"]) for c in chunks])
    return {"chunks": chunks, "faiss": faiss_index, "bm25": bm25}

def hybrid_search(query, index, embedder, top_k, w_sem):
    """FAISS semantic + BM25, fused with Reciprocal Rank Fusion."""
    q_vec = embedder.encode([light_clean(query)], convert_to_numpy=True,
                            normalize_embeddings=True, show_progress_bar=False).astype("float32")
    n = len(index["chunks"])
    pool = min(n, top_k * 4)

    scores, ids = index["faiss"].search(q_vec, pool)
    sem_scores = {int(i): float(s) for i, s in zip(ids[0], scores[0]) if i >= 0}
    best_sem = float(scores[0][0])

    bm = index["bm25"].get_scores(tokenize(query))
    bm_rank = np.argsort(-bm)[:pool]

    fused = defaultdict(float)
    for r, i in enumerate(ids[0]):
        if i >= 0:
            fused[int(i)] += w_sem / (RRF_K + r + 1)
    for r, i in enumerate(bm_rank):
        if bm[i] > 0:
            fused[int(i)] += (1 - w_sem) / (RRF_K + r + 1)

    top = sorted(fused, key=fused.get, reverse=True)[:top_k]
    return top, best_sem, sem_scores, bm

# ========================= GENERATION =========================
def generate_answer(llm, query, retrieved):
    context = "\n\n---\n\n".join(f"[{cite_label(c)}]\n{c['text']}" for c in retrieved)
    prompt = f"""You are a document assistant. Answer ONLY from the context below.
- Answer in the same language as the question.
- After each fact, cite its source exactly as shown in the context, e.g. [file, p.N].
- Table rows use " | " between cells; the first line of a table chunk is its header.
- If the context does not contain the answer, say it was not found in the documents and
  suggest contacting the document owner. Do not guess.

Context:
{context}

Question: {query}
Answer:"""
    return llm.invoke(prompt).content

# ========================= STREAMLIT APP =========================
st.set_page_config(page_title="HR Policies / Benefits Assistant", page_icon="📚", layout="wide")
HEADER_IMAGE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "HR_policies.png")
if os.path.exists(HEADER_IMAGE):
    st.image(HEADER_IMAGE, use_container_width=True)
else:  # image missing from the repo: fall back to a text title
    st.title("HR Policies / Benefits Assistant")
    st.caption("Find information about company policies, benefits, leaves, and employee support.")

with st.sidebar:
    st.header("Settings")
    top_k = st.slider("Chunks sent to the LLM (Top-K)", 3, 10, 5)
    w_sem = st.slider("Semantic weight (BM25 gets the rest)", 0.0, 1.0, 0.6, 0.1)
    min_rel = st.slider("Minimum relevance to answer", 0.0, 0.8, 0.25, 0.05,
                        help="If the best semantic match is below this, the app says 'not found' instead of guessing.")

    with st.expander("OCR diagnostics"):
        tbin = shutil.which("tesseract")
        st.write(f"tesseract binary: `{tbin or 'NOT FOUND (using PyMuPDF OCR fallback)'}`")
        if tbin:
            try:
                _set_tessdata()
                out = subprocess.run(["tesseract", "--list-langs"], capture_output=True, text=True, timeout=20)
                st.code((out.stdout or out.stderr).strip())
            except Exception as e:
                st.write(f"Could not query Tesseract: {e}")
        st.write(f"TESSDATA_PREFIX: `{os.environ.get('TESSDATA_PREFIX', 'unset')}`")
        st.caption(f"Extraction version: {EXTRACT_VERSION}")
        if st.button("Clear extraction cache"):
            st.cache_data.clear()
            st.session_state.pop("signature", None)
            st.rerun()

embedder = load_embedder()

groq_key = None
try:
    groq_key = st.secrets.get("GROQ_API_KEY")
except Exception:
    pass
groq_key = groq_key or os.getenv("GROQ_API_KEY")
if not groq_key:
    st.error("GROQ_API_KEY not found. Add it to Streamlit secrets or environment variables.")
    st.stop()
llm = ChatGroq(temperature=0.1, model_name=LLM_MODEL_NAME, groq_api_key=groq_key)

uploaded_files = st.file_uploader(
    "Upload documents",
    type=["pdf", "docx", "pptx", "xlsx", "csv", "txt", "md"],
    accept_multiple_files=True,
)

# Rebuild the index whenever the set of uploaded files changes
if uploaded_files:
    signature = tuple(sorted((f.name, f.size) for f in uploaded_files))
    if st.session_state.get("signature") != signature:
        all_blocks, report = [], {}
        with st.spinner("Reading documents (scanned pages are OCR'd and may take a while)..."):
            for f in uploaded_files:
                try:
                    blocks, stats = extract_file(f.name, f.getvalue(), EXTRACT_VERSION)
                except Exception as e:
                    blocks, stats = [], {"pages": 0, "tables": 0, "notes": [f"Could not read file: {e}"]}
                all_blocks.extend(blocks)
                report[f.name] = stats
        chunks = chunk_blocks(all_blocks)
        if not chunks:
            st.error("No readable text found in the uploaded files.")
            st.session_state.pop("index", None)
            st.session_state.pop("signature", None)
            with st.expander("Extraction report", expanded=True):
                for fname, s in report.items():
                    for note in s["notes"]:
                        st.write(f"**{fname}**: {note}")
            st.stop()
        with st.spinner("Building FAISS + BM25 indexes..."):
            st.session_state.index = build_index(chunks, embedder)
        st.session_state.report = report
        st.session_state.signature = signature
    st.success(f"Indexed {len(st.session_state.index['chunks'])} chunks from {len(uploaded_files)} file(s)")

    with st.expander("📋 Extraction report"):
        for fname, s in st.session_state.report.items():
            st.write(f"**{fname}**: {s['pages']} page(s)/sheet(s), {s['tables']} table(s)")
            for note in s["notes"]:
                st.caption(f"⚠️ {note}")
else:
    st.session_state.pop("index", None)
    st.session_state.pop("signature", None)

# Query form: the LLM is only called when the user presses the button
with st.form("ask"):
    query = st.text_input("Ask a question about your documents (Arabic or English):")
    submitted = st.form_submit_button("Ask")

if submitted and query.strip():
    index = st.session_state.get("index")
    if not index:
        st.warning("Upload at least one document first.")
    else:
        top, best_sem, sem_scores, bm = hybrid_search(query, index, embedder, top_k, w_sem)
        retrieved = [index["chunks"][i] for i in top]

        st.subheader("Answer")
        if not retrieved or best_sem < min_rel:
            st.info("Not found in the uploaded documents.")
        else:
            with st.spinner("Generating answer..."):
                st.write(generate_answer(llm, query, retrieved))

        with st.expander("🔍 Retrieved chunks"):
            for rank, i in enumerate(top, 1):
                c = index["chunks"][i]
                sem_txt = f"{sem_scores[i]:.2f}" if i in sem_scores else "n/a"
                st.markdown(f"**{rank}. {cite_label(c)}** ({c['kind']})  ·  semantic {sem_txt}  ·  BM25 {bm[i]:.2f}")
                st.text(c["text"])
else:
    st.info("Upload documents, type a question, and press Ask.")