"""
PolicyQA — Document QA Agent
Upload a policy document or report (PDF or plain text).
Ask questions in natural language.
Get answers cited to exact sections.

Pipeline:
  Upload → PDF extract → section-aware chunking
  → BM25 + BGE-base dense + RRF fusion
  → BGE-reranker-large (cross-encoder)
  → flan-t5-base  generates answer  (250MB)
  → flan-t5-large judges (score >= 6) (770MB — heavier than generator)
  → Answer + "Section X · Para Y" citations
"""

# ── Patch gradio_client BEFORE importing gradio ────────────────────────────────
# gradio-client 1.3.0 has a bug where schema can be a bool and
# `if "X" in schema:` raises TypeError. Fix all occurrences first.
import pathlib, re as _re, importlib
_gc_path = None
try:
    import gradio_client.utils as _gcu
    _gc_path = pathlib.Path(_gcu.__file__)
    _src = _gc_path.read_text()
    _fixed = _re.sub(
        r'\bif\s+"(\w+)"\s+in\s+schema\s*:',
        r'if isinstance(schema, dict) and "\1" in schema:',
        _src
    )
    if _fixed != _src:
        _gc_path.write_text(_fixed)
        importlib.reload(_gcu)
        print(f"Patched {_gc_path.name} ✓")
    else:
        print("gradio_client already clean or pattern not found")
except Exception as _e:
    print(f"gradio_client patch error: {_e}")
# ──────────────────────────────────────────────────────────────────────────────

import re
import io
import fitz                          # pymupdf — PDF text extraction
import faiss
import numpy as np
import gradio as gr
from dataclasses import dataclass, field
from typing import Optional
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer, CrossEncoder
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# ── Models ─────────────────────────────────────────────────────────────────────
EMBED_MODEL    = "BAAI/bge-base-en-v1.5"
RERANKER_MODEL = "BAAI/bge-reranker-large"
GEN_MODEL      = "google/flan-t5-base"      # generator  ~250MB
JUDGE_MODEL    = "google/flan-t5-large"     # judge      ~770MB (heavier than generator)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32

print(f"Device: {DEVICE}")

print("Loading embedder …")
embedder = SentenceTransformer(EMBED_MODEL, device=DEVICE)

print("Loading reranker …")
reranker = CrossEncoder(RERANKER_MODEL, device=DEVICE)

print("Loading generator (flan-t5-base) …")
gen_tok   = AutoTokenizer.from_pretrained(GEN_MODEL)
gen_model = AutoModelForSeq2SeqLM.from_pretrained(GEN_MODEL, torch_dtype=DTYPE).to(DEVICE)

print("Loading judge (flan-t5-large) …")
judge_tok   = AutoTokenizer.from_pretrained(JUDGE_MODEL)
judge_model = AutoModelForSeq2SeqLM.from_pretrained(JUDGE_MODEL, torch_dtype=DTYPE).to(DEVICE)

print("All models ready.")

# ── Tunables ───────────────────────────────────────────────────────────────────
CHUNK_CHARS        = 600    # characters per chunk
CHUNK_OVERLAP      = 100    # overlap between consecutive chunks in same section
DENSE_TOPK         = 20
BM25_TOPK          = 20
RRF_K              = 60
RERANK_TOPK        = 5
RRF_GATE           = 0.30   # min normalised RRF score — below → escalate
JUDGE_PASS         = 6      # judge score threshold out of 10
GEN_MAX_TOKENS     = 350
JUDGE_MAX_TOKENS   = 80

# ── Data structures ────────────────────────────────────────────────────────────
@dataclass
class Chunk:
    text:          str
    section_path:  str   # e.g. "3 · Risk Assessment > 3.2 · Identification"
    section_index: int   # global section counter (for stable citation numbering)
    para_index:    int   # paragraph index within the section
    char_start:    int
    char_end:      int

@dataclass
class DocState:
    doc_name:    str           = ""
    chunks:      list          = field(default_factory=list)
    faiss_index: Optional[object] = None
    bm25_index:  Optional[object] = None
    toc:         list[str]     = field(default_factory=list)   # section headings in order
    history:     list          = field(default_factory=list)

state = DocState()

# ── PDF / text extraction ──────────────────────────────────────────────────────
def extract_text_from_pdf(path: str) -> str:
    doc  = fitz.open(path)
    pages = []
    for page in doc:
        pages.append(page.get_text("text"))
    doc.close()
    return "\n".join(pages)

def extract_text(file_obj) -> tuple[str, str]:
    """Returns (raw_text, doc_name). Gradio 4 passes a tempfile object with .name path."""
    path = file_obj.name if hasattr(file_obj, "name") else str(file_obj)
    name = path.split("/")[-1]
    if name.lower().endswith(".pdf") or path.lower().endswith(".pdf"):
        raw = extract_text_from_pdf(path)
    else:
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            raw = f.read()
    return raw, name

# ── Section-aware parser ───────────────────────────────────────────────────────
# Matches headings like: "1.", "1.2", "1.2.3", "ARTICLE IV", "Chapter 3",
# "Section 4", or ALL-CAPS lines (≤ 60 chars) used as headings in many reports.
HEADING_RE = re.compile(
    r"^("
    r"(?:chapter|section|article|part|annex)\s+[\dA-Z][\w\.]*"  # named headings
    r"|[\d]+(?:\.[\d]+)*\.?\s+\S"                                # numbered 1. / 1.2 / 1.2.3
    r"|[A-Z][A-Z\s\-]{3,59}$"                                    # ALL-CAPS titles ≤60 chars
    r")",
    re.IGNORECASE | re.MULTILINE,
)

def parse_sections(raw: str) -> list[tuple[str, str]]:
    """
    Returns list of (section_heading, section_body) pairs.
    Heading '' means text before the first heading (preamble).
    """
    lines   = raw.splitlines()
    sections: list[tuple[str, str]] = []
    current_heading = ""
    current_lines: list[str] = []

    for line in lines:
        stripped = line.strip()
        if not stripped:
            current_lines.append("")
            continue
        if HEADING_RE.match(stripped) and len(stripped) < 120:
            # flush previous section
            body = "\n".join(current_lines).strip()
            if body or current_heading:
                sections.append((current_heading, body))
            current_heading = stripped
            current_lines = []
        else:
            current_lines.append(stripped)

    # flush last
    body = "\n".join(current_lines).strip()
    if body or current_heading:
        sections.append((current_heading, body))

    # If we found zero headings, treat the whole doc as one section
    if len(sections) == 1 and sections[0][0] == "":
        return [("Document", sections[0][1])]

    return sections

def build_section_path(heading: str, idx: int) -> str:
    """Readable citation path, e.g. '§4 · Risk Assessment'."""
    if not heading:
        return "§Preamble"
    return f"§{idx + 1} · {heading[:80]}"

# ── Chunking ───────────────────────────────────────────────────────────────────
def chunk_section(
    text: str,
    section_path: str,
    section_index: int,
) -> list[Chunk]:
    """Split section body into overlapping character chunks."""
    chunks: list[Chunk] = []
    start = 0
    para  = 0
    while start < len(text):
        end = min(start + CHUNK_CHARS, len(text))
        chunks.append(Chunk(
            text          = text[start:end],
            section_path  = section_path,
            section_index = section_index,
            para_index    = para,
            char_start    = start,
            char_end      = end,
        ))
        if end == len(text):
            break
        start += CHUNK_CHARS - CHUNK_OVERLAP
        para  += 1
    return chunks

def ingest_document(raw: str) -> tuple[list[Chunk], object, object, list[str]]:
    sections = parse_sections(raw)
    all_chunks: list[Chunk] = []
    toc: list[str] = []

    for sec_idx, (heading, body) in enumerate(sections):
        if not body.strip():
            continue
        path = build_section_path(heading, sec_idx)
        toc.append(path)
        # prepend heading into each chunk so retrieval is heading-aware
        headed_body = f"{heading}\n{body}" if heading else body
        all_chunks += chunk_section(headed_body, path, sec_idx)

    if not all_chunks:
        return [], None, None, []

    texts = [c.text for c in all_chunks]

    # Dense index (BGE-base)
    embeddings = embedder.encode(
        texts, normalize_embeddings=True, batch_size=32, show_progress_bar=False
    ).astype("float32")
    faiss_index = faiss.IndexFlatIP(embeddings.shape[1])
    faiss_index.add(embeddings)

    # Sparse index (BM25)
    tokenised  = [t.lower().split() for t in texts]
    bm25_index = BM25Okapi(tokenised)

    return all_chunks, faiss_index, bm25_index, toc

# ── Hybrid retrieval + RRF ─────────────────────────────────────────────────────
def rrf_fuse(rank_dicts: list[dict[int, int]], k: int = RRF_K) -> dict[int, float]:
    scores: dict[int, float] = {}
    for rd in rank_dicts:
        for doc_id, rank in rd.items():
            scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank)
    return scores

def hybrid_retrieve(query: str) -> list[tuple[Chunk, float]]:
    chunks      = state.chunks
    faiss_index = state.faiss_index
    bm25_index  = state.bm25_index
    n           = len(chunks)

    # Dense
    q_emb = embedder.encode([query], normalize_embeddings=True).astype("float32")
    _, d_idx = faiss_index.search(q_emb, min(DENSE_TOPK, n))
    dense_ranks = {int(d_idx[0][r]): r for r in range(len(d_idx[0])) if d_idx[0][r] >= 0}

    # BM25
    bm25_scores  = bm25_index.get_scores(query.lower().split())
    bm25_top     = np.argsort(bm25_scores)[::-1][:BM25_TOPK]
    bm25_ranks   = {int(i): rank for rank, i in enumerate(bm25_top)}

    # RRF
    fused = rrf_fuse([dense_ranks, bm25_ranks])
    if not fused:
        return []

    max_score  = max(fused.values())
    sorted_ids = sorted(fused, key=lambda i: fused[i], reverse=True)
    return [(chunks[i], fused[i] / max_score) for i in sorted_ids if i < n]

def rerank_chunks(query: str, candidates: list[tuple[Chunk, float]]) -> list[tuple[Chunk, float]]:
    pairs     = [[query, c.text] for c, _ in candidates]
    ce_scores = reranker.predict(pairs, batch_size=16)
    ranked    = sorted(zip(candidates, ce_scores), key=lambda x: x[1], reverse=True)
    top       = ranked[:RERANK_TOPK]
    sigmoid   = lambda x: float(1.0 / (1.0 + np.exp(-x)))
    return [(chunk, sigmoid(s)) for (chunk, _), s in top]

# ── Generation ─────────────────────────────────────────────────────────────────
def build_context(top_chunks: list[tuple[Chunk, float]]) -> str:
    """Build numbered context block with inline section citations."""
    parts = []
    for i, (c, _) in enumerate(top_chunks, 1):
        parts.append(f"[{i}] ({c.section_path}, para {c.para_index + 1})\n{c.text}")
    return "\n\n".join(parts)

def generate_answer(query: str, context: str) -> str:
    prompt = (
        "You are a precise document analyst. "
        "Answer the question using ONLY the numbered context passages below. "
        "Reference passages by their section and paragraph number. "
        "Do not use outside knowledge or hallucinate.\n\n"
        f"Context:\n{context}\n\n"
        f"Question: {query}\n\n"
        "Answer (cite sections inline, e.g. 'According to §3 · Risk Assessment, para 2, ...'):"
    )
    inputs = gen_tok(prompt, return_tensors="pt", truncation=True, max_length=1200).to(DEVICE)
    with torch.no_grad():
        out = gen_model.generate(
            **inputs,
            max_new_tokens=GEN_MAX_TOKENS,
            num_beams=4,
            early_stopping=True,
            no_repeat_ngram_size=3,
            length_penalty=1.2,
        )
    return gen_tok.decode(out[0], skip_special_tokens=True).strip()

# ── Judge ──────────────────────────────────────────────────────────────────────
def judge_answer(query: str, context: str, answer: str) -> tuple[int, str]:
    """flan-t5-large (770MB) judges the flan-t5-base (250MB) answer. Returns (score, rationale)."""
    prompt = (
        "You are a strict QA evaluator for policy documents.\n"
        f"Question: {query}\n"
        f"Retrieved context:\n{context[:2000]}\n"
        f"Proposed answer: {answer}\n\n"
        "Score the answer 1–10 on:\n"
        "  • Faithfulness — no claims beyond the context\n"
        "  • Completeness — addresses the full question\n"
        "  • Section accuracy — cites correct sections\n\n"
        "Reply in EXACTLY this format:\n"
        "Score: <integer 1-10>\n"
        "Rationale: <one sentence>"
    )
    inputs = judge_tok(prompt, return_tensors="pt", truncation=True, max_length=1200).to(DEVICE)
    with torch.no_grad():
        out = judge_model.generate(**inputs, max_new_tokens=JUDGE_MAX_TOKENS, num_beams=2)
    raw = judge_tok.decode(out[0], skip_special_tokens=True).strip()

    m_score = re.search(r"Score:\s*(\d+)", raw)
    score   = max(1, min(10, int(m_score.group(1)))) if m_score else 1

    m_rat     = re.search(r"Rationale:\s*(.+)", raw, re.DOTALL)
    rationale = m_rat.group(1).strip() if m_rat else raw[:200]

    return score, rationale

# ── Escalation ─────────────────────────────────────────────────────────────────
ESCALATION = {
    "no_doc":          "No document loaded. Upload a PDF or paste text first.",
    "low_retrieval":   "No passages found with sufficient relevance to your question.",
    "gen_error":       "Answer generation failed — the question may be out of scope.",
    "judge_fail":      "The generated answer did not pass quality review.",
}

def escalate(key: str, query: str, detail: str = "") -> str:
    detail_block = f"\n\n*Judge rationale: {detail}*" if detail else ""
    return (
        f"⚠️ **Escalation — {ESCALATION.get(key, 'Unknown error.')}**"
        f"{detail_block}\n\n"
        f"**Try:**\n"
        f"- Rephrase your question using terms from the document\n"
        f"- Check the Table of Contents to find the relevant section\n"
        f"- Ask a narrower, more specific question\n\n"
        f"> *Your question: {query}*"
    )

# ── Citation formatter ─────────────────────────────────────────────────────────
def fmt_citation_block(top_chunks: list[tuple[Chunk, float]]) -> str:
    lines = []
    for i, (c, score) in enumerate(top_chunks, 1):
        lines.append(
            f"**[{i}] {c.section_path}** · para {c.para_index + 1} "
            f"*(relevance: {score:.2f})*\n"
            f"> {c.text[:220].strip()}…"
        )
    return "\n\n".join(lines)

def suggest_followups(top_chunks: list[tuple[Chunk, float]]) -> str:
    seen_sections = list(dict.fromkeys(c.section_path for c, _ in top_chunks))[:3]
    questions = [
        f"What are the key requirements in {seen_sections[0]}?" if seen_sections else
        "What are the key requirements in this section?",
        "What exceptions or conditions apply here?",
        "What are the penalties or consequences described?",
        "What definitions are given for the terms used?",
    ]
    return "\n".join(f"- {q}" for q in questions[:3])

# ── Core agent ─────────────────────────────────────────────────────────────────
def answer_question(query: str) -> str:
    query = query.strip()
    if not query:
        return "Please enter a question."
    if not state.chunks or state.faiss_index is None:
        return escalate("no_doc", query)

    # 1. Hybrid retrieval
    candidates = hybrid_retrieve(query)
    if not candidates or candidates[0][1] < RRF_GATE:
        return escalate("low_retrieval", query)

    # 2. Rerank top candidates
    top_chunks = rerank_chunks(query, candidates)

    # 3. Build context (numbered, with section labels)
    context = build_context(top_chunks)

    # 4. Generate answer — flan-t5-base
    try:
        answer = generate_answer(query, context)
    except Exception:
        return escalate("gen_error", query)

    # 5. Judge answer — flan-t5-large
    try:
        score, rationale = judge_answer(query, context, answer)
    except Exception:
        score, rationale = 5, "Judge unavailable."

    # 6. Escalate if judge rejects
    if score < JUDGE_PASS:
        return escalate("judge_fail", query, detail=rationale)

    # 7. Return cited answer
    cite_block = fmt_citation_block(top_chunks[:3])
    followups  = suggest_followups(top_chunks)

    state.history.append({"q": query, "a": answer, "score": score})

    return (
        f"### Answer\n{answer}\n\n"
        f"*Judge score: {score}/10 · {rationale}*\n\n"
        f"---\n\n"
        f"**Cited passages**\n\n{cite_block}\n\n"
        f"---\n\n"
        f"**Suggested follow-ups**\n{followups}"
    )

# ── Document ingestion handler ─────────────────────────────────────────────────
def load_document(file_obj, pasted_text: str) -> str:
    # Prefer uploaded file; fall back to pasted text
    if file_obj is not None:
        try:
            raw, name = extract_text(file_obj)
        except Exception as e:
            return f"❌ Could not read file: {e}"
    elif pasted_text and pasted_text.strip():
        raw  = pasted_text.strip()
        name = "pasted-document.txt"
    else:
        return "Upload a file or paste text to begin."

    if len(raw.strip()) < 100:
        return "❌ Document appears empty or too short."

    chunks, faiss_index, bm25_index, toc = ingest_document(raw)
    if not chunks:
        return "❌ Could not extract any text from the document."

    state.doc_name    = name
    state.chunks      = chunks
    state.faiss_index = faiss_index
    state.bm25_index  = bm25_index
    state.toc         = toc
    state.history     = []

    toc_block = "\n".join(f"  {t}" for t in toc[:30])
    overflow  = f"\n  … and {len(toc) - 30} more sections" if len(toc) > 30 else ""

    return (
        f"✅ **{name}** indexed\n"
        f"  {len(chunks)} chunks · {len(toc)} sections\n\n"
        f"**Table of contents (detected)**\n{toc_block}{overflow}"
    )

def clear_state() -> tuple:
    state.doc_name = ""; state.chunks = []; state.faiss_index = None
    state.bm25_index = None; state.toc = []; state.history = []
    return None, "", "", "*Document cleared.*"

# ── Gradio UI ──────────────────────────────────────────────────────────────────
CSS = """
body { font-family: 'IBM Plex Mono', monospace; }
#hdr { background:#0d0d0d; color:#e8e8e8; padding:1.6rem 2.4rem 1.3rem; border-bottom:2px solid #1e1e1e; }
#hdr h1 { font-size:1.45rem; font-weight:600; margin:0 0 .2rem; color:#fff; letter-spacing:-.3px; }
#hdr p  { font-size:.78rem; color:#666; margin:0; }
#load-btn { background:#111 !important; color:#e8e8e8 !important; border:1px solid #2a2a2a !important; }
#ask-btn  { background:#111 !important; color:#e8e8e8 !important; border:1px solid #2a2a2a !important; }
#clr-btn  { background:transparent !important; color:#555 !important; border:1px solid #222 !important; }
"""

PIPELINE_INFO = (
    "**Pipeline**\n"
    "1. BGE-base → FAISS (dense, top-20)\n"
    "2. BM25Okapi (sparse, top-20)\n"
    "3. RRF fusion k=60\n"
    "4. BGE-reranker-large (top-5)\n"
    "5. flan-t5-base → answer (250MB)\n"
    "6. flan-t5-large → judge (770MB, ≥ 6/10)\n\n"
    "*Escalates on: no doc · low retrieval · judge < 6*"
)

with gr.Blocks(css=CSS, title="PolicyQA") as demo:

    gr.HTML("""
    <div id="hdr">
      <h1>📋 PolicyQA — Document QA Agent</h1>
      <p>Upload a policy document or report → ask questions → get answers cited to exact sections</p>
    </div>""")

    with gr.Row():

        # ── Left: document loader ──────────────────────────────────────────────
        with gr.Column(scale=1, min_width=300):
            gr.Markdown("**01 · Load Document**")

            file_upload = gr.File(
                label="Upload PDF or .txt",
                file_types=[".pdf", ".txt"],
            )
            paste_box = gr.Textbox(
                label="Or paste document text",
                placeholder="Paste the full text of your policy document here…",
                lines=6,
            )
            load_btn  = gr.Button("Index Document", elem_id="load-btn")
            load_out  = gr.Markdown("*No document loaded.*")

            gr.Markdown("---")
            gr.Markdown(PIPELINE_INFO)
            clr_btn = gr.Button("Clear & Reset", elem_id="clr-btn")

        # ── Right: QA interface ────────────────────────────────────────────────
        with gr.Column(scale=2):
            gr.Markdown("**02 · Ask the Agent**")

            q_box   = gr.Textbox(
                label="Your question",
                placeholder="What are the reporting obligations under Section 4?",
                lines=3,
            )
            ask_btn = gr.Button("Ask →", elem_id="ask-btn")
            ans_box = gr.Markdown("*Index a document first, then ask a question.*")

    # ── Events ─────────────────────────────────────────────────────────────────
    load_btn.click(
        fn=load_document,
        inputs=[file_upload, paste_box],
        outputs=load_out,
    )
    ask_btn.click(
        fn=answer_question,
        inputs=[q_box],
        outputs=ans_box,
    )
    q_box.submit(
        fn=answer_question,
        inputs=[q_box],
        outputs=ans_box,
    )
    clr_btn.click(
        fn=clear_state,
        inputs=[],
        outputs=[file_upload, paste_box, q_box, load_out],
    )

# Fix Gradio 4.x localhost reachability check
try:
    import gradio.networking as _gn
    _gn.is_localhost = lambda *a, **kw: False
except Exception:
    pass

demo.launch(server_name="0.0.0.0", server_port=7860)