#!/usr/bin/env python3
"""PageHawk AI — 100% ON-DEVICE document intelligence via a local LLM (Ollama).

Nothing is uploaded, ever. Ask-your-document Q&A with page citations, document
summarize, and structured key-term extraction all run locally against Ollama
(qwen2.5-coder by default). This is the privacy wedge: a lawyer or diligence
analyst with a confidential contract can use it where every cloud AI is off-limits.

Degrades gracefully: if Ollama is not running, returns a clear, actionable error
instead of throwing, so the UI can tell the user exactly how to start it.

CLI:  python pagehawk_ai.py status
      python pagehawk_ai.py ask      --doc FILE --q "question"
      python pagehawk_ai.py summarize --doc FILE
      python pagehawk_ai.py terms    --doc FILE
"""
import os, re, json, math, time, subprocess, urllib.request
from collections import Counter

OLLAMA = os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")
# preference order; the first installed match wins (probed live, never hardcoded blindly)
PREFER = ["qwen2.5-coder:14b", "qwen2.5-coder", "qwen2.5", "llama3.1", "llama3.2",
          "mistral", "gemma2", "phi3"]


# ---------------------------------------------------------------- ollama plumbing
def _http(path, payload=None, timeout=300):
    url = OLLAMA + path
    if payload is None:
        req = urllib.request.Request(url)
    else:
        req = urllib.request.Request(url, data=json.dumps(payload).encode(),
                                     headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8", "ignore"))


def available():
    try:
        return [m["name"] for m in _http("/api/tags", timeout=5).get("models", [])]
    except Exception:
        return []


def pick_model():
    have = available()
    if not have:
        return None
    for p in PREFER:
        for h in have:
            if h == p or h.startswith(p):
                return h
    return have[0]


def _generate(prompt, model, num_predict=600, temperature=0.0, fmt=None, timeout=300):
    payload = {"model": model, "prompt": prompt, "stream": False,
               "options": {"temperature": temperature, "num_predict": num_predict}}
    if fmt:
        payload["format"] = fmt
    d = _http("/api/generate", payload, timeout=timeout)
    return d.get("response", "")


# ---------------------------------------------------------------- text extraction
def extract_pages(path):
    """List of page texts. PDFs split on form-feed (real page boundaries) so we can
    cite page numbers; scanned PDFs are OCR'd on-device; other formats are chunked
    into ~1800-word pseudo-pages so citations and retrieval still work."""
    ext = path.lower().rsplit(".", 1)[-1] if "." in path else ""
    if ext == "pdf":
        try:
            r = subprocess.run(["pdftotext", "-layout", "-q", path, "-"],
                               capture_output=True, timeout=120)
            txt = r.stdout.decode("utf-8", "ignore")
        except Exception:
            txt = ""
        if len(txt.strip()) < 80:                       # scanned / no text layer -> OCR
            try:
                import pagehawk_diff as d
                txt = d._ocr_pdf(path) or txt
            except Exception:
                pass
        pages = [p.strip() for p in txt.split("\f")]
        pages = [p for p in pages if p] or [txt.strip()]
        return pages or [""]
    # non-PDF: get plain text, then chunk
    try:
        import pagehawk_diff as d
        txt = d.extract(path)
    except Exception:
        txt = open(path, encoding="utf-8", errors="ignore").read() if os.path.exists(path) else ""
    words = txt.split()
    if len(words) <= 1800:
        return [txt.strip()] if txt.strip() else [""]
    return [" ".join(words[i:i + 1800]) for i in range(0, len(words), 1800)]


# ---------------------------------------------------------------- BM25 retrieval
_WORD = re.compile(r"[A-Za-z0-9']+")


def _tok(s):
    return [w.lower() for w in _WORD.findall(s)]


def retrieve(pages, query, k=4):
    """Return indices of the top-k pages most relevant to the query (BM25). On a query
    with no lexical overlap, falls back to the first k pages so the model still has context."""
    docs = [_tok(p) for p in pages]
    N = len(docs)
    avg = sum(len(d) for d in docs) / max(1, N)
    df = Counter()
    for d in docs:
        for t in set(d):
            df[t] += 1
    q = _tok(query)
    k1, b = 1.5, 0.75
    scores = []
    for i, d in enumerate(docs):
        tf = Counter(d)
        s = 0.0
        for t in q:
            if t not in tf:
                continue
            idf = math.log(1 + (N - df[t] + 0.5) / (df[t] + 0.5))
            s += idf * (tf[t] * (k1 + 1)) / (tf[t] + k1 * (1 - b + b * len(d) / max(1, avg)))
        scores.append((s, i))
    scores.sort(reverse=True)
    top = [i for s, i in scores[:k] if s > 0]
    if not top:
        top = list(range(min(k, N)))
    return sorted(top)


# ---------------------------------------------------------------- public API
_NORUN = ("Local AI is not running. Open Terminal and run  ollama serve  "
          "(and once:  ollama pull qwen2.5-coder:14b ). Everything stays on your Mac.")


def status():
    return {"running": bool(available()), "models": available(), "model": pick_model()}


_DOC_EXTS = (".pdf", ".doc", ".docx", ".rtf", ".txt", ".md", ".markdown", ".odt", ".html", ".htm")


def _gather(path, max_chunks=420):
    """Return (chunks, is_folder) where chunks = [(label, text), ...]. A folder flattens
    every document into per-(file, page) chunks so one question can search the whole set."""
    if os.path.isdir(path):
        files = sorted(f for f in os.listdir(path)
                       if f.lower().endswith(_DOC_EXTS) and not f.startswith("."))
        chunks = []
        for f in files:
            for i, pg in enumerate(extract_pages(os.path.join(path, f))):
                if pg.strip():
                    chunks.append(("%s p.%d" % (f, i + 1), pg))
                if len(chunks) >= max_chunks:
                    break
            if len(chunks) >= max_chunks:
                break
        return chunks, True
    return [("p.%d" % (i + 1), pg) for i, pg in enumerate(extract_pages(path))], False


def answer(path, question, k=4, timeout=300):
    """Ask-your-document(s) Q&A, fully on-device, with citations. Works on a single
    document (cite page) OR a whole folder (cite [file p.N] across every document)."""
    t0 = time.time()
    model = pick_model()
    if not model:
        return {"error": _NORUN, "answer": "", "pages": []}
    chunks, folder = _gather(path)
    if not chunks:
        return {"error": "No readable documents found here.", "answer": "", "pages": []}
    texts = [c[1] for c in chunks]
    idx = retrieve(texts, question, max(k, 6) if folder else k)
    ctx = "\n\n".join("[%s]\n%s" % (chunks[i][0], texts[i][:2400]) for i in idx)
    if folder:
        cite_rule = "cite the source in square brackets like [Acme NDA.pdf p.3] for every fact"
        scope = "DOCUMENTS (one folder, many files):"
    else:
        cite_rule = "cite the page in parentheses like (p. 3) for every fact"
        scope = "DOCUMENT PAGES:"
    prompt = (
        "You are a meticulous document analyst working OFFLINE on the user's Mac. "
        "Answer ONLY from the excerpts provided below and %s. If the answer is not in "
        "them, say so plainly. Be concise and precise.\n\n"
        "%s\n%s\n\nQUESTION: %s\n\nANSWER:" % (cite_rule, scope, ctx, question))
    try:
        ans = _generate(prompt, model, num_predict=650, temperature=0.0, timeout=timeout).strip()
    except Exception as e:
        return {"error": "Local AI error: %s" % e, "answer": "", "pages": []}
    out = {"answer": ans, "model": model, "folder": folder, "ms": int((time.time() - t0) * 1000)}
    if folder:
        out["sources"] = [chunks[i][0] for i in idx]
    else:
        out["pages"] = [i + 1 for i in idx]
    return out


def explain_passages(passages, query="", timeout=300):
    """v0.37 Search Workbench: synthesize what a set of SELECTED search hits collectively
    say, citing each source. `passages` = [{"name","anchor","snippet"}]. Fully on-device."""
    t0 = time.time()
    model = pick_model()
    if not model:
        return {"error": _NORUN, "explanation": ""}
    if not passages:
        return {"error": "No passages selected.", "explanation": ""}
    lines = []
    for i, p in enumerate(passages[:24]):
        tag = "%s%s" % (p.get("name", "?"), (" " + p.get("anchor", "")) if p.get("anchor") else "")
        lines.append("[%s] %s" % (tag, (p.get("snippet") or "").strip()))
    ctx = "\n".join(lines)
    ask = ("the search term “%s”" % query) if query else "these passages"
    prompt = (
        "You are a meticulous document analyst working OFFLINE on the user's Mac. The user "
        "searched across a folder and selected the matching passages below. Explain what they "
        "collectively say about %s: the throughline, any agreements or conflicts between them, "
        "and what a busy professional should take away. Cite each source in square brackets "
        "exactly as labeled, like [Acme NDA.pdf p.3], for every point. Answer ONLY from the "
        "passages; if they don't support a claim, say so.\n\nPASSAGES:\n%s\n\nEXPLANATION:"
        % (ask, ctx))
    try:
        ans = _generate(prompt, model, num_predict=700, temperature=0.1, timeout=timeout).strip()
    except Exception as e:
        return {"error": "Local AI error: %s" % e, "explanation": ""}
    return {"explanation": ans, "model": model, "count": len(passages),
            "sources": [p.get("name", "") for p in passages],
            "ms": int((time.time() - t0) * 1000)}


def summarize(path, timeout=300):
    """Whole-document summary as scannable bullets, fully on-device."""
    t0 = time.time()
    model = pick_model()
    if not model:
        return {"error": _NORUN, "summary": ""}
    if os.path.isdir(path):
        return {"error": "Summarize works on a single document. Ask a question to search across the whole folder.", "summary": ""}
    pages = extract_pages(path)
    words = "\n\n".join(pages).split()
    clipped = " ".join(words[:6500])
    prompt = (
        "You are a document analyst working offline on the user's Mac. Summarize the "
        "document below as 5-8 tight bullet points a busy professional can scan in "
        "seconds. Call out parties, dates, obligations, amounts, governing law, and "
        "anything unusual or risky. Output plain-text bullets each starting with '- '.\n\n"
        "DOCUMENT:\n%s\n\nSUMMARY:" % clipped)
    try:
        s = _generate(prompt, model, num_predict=750, temperature=0.1, timeout=timeout).strip()
    except Exception as e:
        return {"error": "Local AI error: %s" % e, "summary": ""}
    return {"summary": s, "model": model, "pages": len(pages),
            "ms": int((time.time() - t0) * 1000)}


def extract_terms(path, timeout=300):
    """Structured key-term / defined-term extraction (JSON), fully on-device."""
    t0 = time.time()
    model = pick_model()
    if not model:
        return {"error": _NORUN, "terms": []}
    if os.path.isdir(path):
        return {"error": "Key terms works on a single document. Ask a question to search across the whole folder.", "terms": []}
    pages = extract_pages(path)
    clipped = " ".join("\n\n".join(pages).split()[:6500])
    prompt = (
        "Extract the defined terms and key clauses from this document. Return STRICT "
        'JSON only, shape: {"terms":[{"term":"...","summary":"one line plain English",'
        '"page":N}]}. No prose outside the JSON.\n\nDOCUMENT:\n%s' % clipped)
    try:
        raw = _generate(prompt, model, num_predict=900, temperature=0.0, fmt="json", timeout=timeout)
        out = json.loads(raw or "{}")
        out["model"] = model
        out["ms"] = int((time.time() - t0) * 1000)
        return out
    except Exception as e:
        return {"error": "Local AI error: %s" % e, "terms": []}


# ---------------------------------------------------------------- AI table of contents (P1 #3 fallback)
# For PDFs with NO embedded bookmarks: read the document's own text on-device and
# propose a navigable outline [{title, page, level}]. Heuristics shortlist heading-
# looking lines per page (cheap, deterministic); the local model then picks the real
# section headings from that shortlist (few calls, since the 14b model is latency-
# bound per CALL, not per token of input). Pure fallback: never replaces a real TOC.
_TOC_MAX_PER_PAGE = 8
_TOC_BATCH_WORDS = 2000          # prompt budget per model call (keeps calls few + fast)
_WS_RE = re.compile(r"\s+")
_TOC_NUMBERED = re.compile(r"^(\d+(\.\d+)*[.)]?|[IVXLC]+[.)]|[A-Z][.)])\s+\S")
_TOC_KEYWORD = re.compile(r"^(chapter|section|part|appendix|article|exhibit|schedule|annex|"
                          r"introduction|background|summary|conclusion|overview|definitions|"
                          r"recitals|preamble|glossary|index|references|bibliography)\b", re.I)


def _toc_headingish(line):
    """Does this line LOOK like a section heading? Cheap shortlist filter only —
    the model makes the actual call from the shortlist."""
    t = line.strip()
    if not (3 <= len(t) <= 90) or not re.search(r"[A-Za-z]", t):
        return False
    words = t.split()
    if len(words) > 12:
        return False
    if t.endswith((",", ";", ":")) or (t.endswith(".") and not _TOC_NUMBERED.match(t)):
        return False
    if _TOC_NUMBERED.match(t) or _TOC_KEYWORD.match(t):
        return True
    letters = [w for w in words if w[:1].isalpha()]
    if not letters:
        return False
    caps = sum(1 for w in letters if w[:1].isupper())
    if t.upper() == t and len(letters) >= 1:      # ALL CAPS heading
        return True
    return len(letters) >= 2 and caps / len(letters) >= 0.7   # Title Case-ish


def _toc_candidates(pages):
    """[(page_no, line), ...] heading-looking lines, running headers/footers removed
    (a 'heading' repeating across many pages is chrome, not structure)."""
    per_page, freq = [], Counter()
    for i, pg in enumerate(pages):
        seen, lines = set(), []
        for ln in pg.splitlines():
            t = _WS_RE.sub(" ", ln).strip()
            if _toc_headingish(t) and t.lower() not in seen:
                seen.add(t.lower())
                lines.append(t)
            if len(lines) >= _TOC_MAX_PER_PAGE:
                break
        per_page.append(lines)
        for t in lines:
            freq[re.sub(r"\d+", "#", t.lower())] += 1
    cutoff = max(3, len(pages) // 4)
    out = []
    for i, lines in enumerate(per_page):
        for t in lines:
            if freq[re.sub(r"\d+", "#", t.lower())] <= cutoff:
                out.append((i + 1, t))
    if len(out) < 3:                              # thin shortlist → fall back to page-top lines
        for i, pg in enumerate(pages):
            for ln in pg.splitlines():
                t = _WS_RE.sub(" ", ln).strip()
                if 3 <= len(t) <= 90 and re.search(r"[A-Za-z]", t):
                    out.append((i + 1, t))
                    break
        out.sort()
    return out


def _toc_sanitize(items, npages):
    """Model output → a clean, navigable outline: valid pages, non-decreasing order,
    clamped levels, no dupes, bounded size."""
    clean, seen, prev = [], set(), 0
    for it in items or []:
        if not isinstance(it, dict):
            continue
        title = _WS_RE.sub(" ", str(it.get("title") or "")).strip()[:120]
        try:
            page = int(it.get("page"))
        except (TypeError, ValueError):
            continue
        if not title or not (1 <= page <= max(1, npages)) or page < prev:
            continue
        key = (title.lower(), page)
        if key in seen:
            continue
        seen.add(key)
        try:
            level = min(3, max(0, int(it.get("level") or 0)))
        except (TypeError, ValueError):
            level = 0
        clean.append({"title": title, "page": page, "level": level})
        prev = page
        if len(clean) >= 400:
            break
    return clean


def toc(path, timeout=300, log=None):
    """Generate a table of contents for a document with no embedded bookmarks.
    Fully on-device. Returns {ok, generated, items:[{title,page,level}], model, pages}."""
    def emit(m):
        if log:
            log(m)
    t0 = time.time()
    model = pick_model()
    if not model:
        return {"error": _NORUN, "items": []}
    if os.path.isdir(path):
        return {"error": "A table of contents works on a single document.", "items": []}
    emit("Reading the document…")
    pages = extract_pages(path)
    npages = len(pages)
    cands = _toc_candidates(pages)
    if not cands:
        return {"error": "No readable text to build a contents from (likely a scan with no OCR text).",
                "items": []}

    # batch the shortlist into few model calls (the model is latency-bound per call)
    batches, cur, words = [], [], 0
    for pg, line in cands:
        w = len(line.split()) + 2
        if cur and words + w > _TOC_BATCH_WORDS:
            batches.append(cur); cur, words = [], 0
        cur.append((pg, line)); words += w
    if cur:
        batches.append(cur)

    items = []
    for bi, batch in enumerate(batches):
        emit("Finding sections: pages %d-%d of %d…" % (batch[0][0], batch[-1][0], npages))
        listing = "\n".join("[page %d] %s" % (pg, line) for pg, line in batch)
        prompt = (
            "These lines were shortlisted from a %d-page document that has NO table of contents. "
            "Some are real section headings; others are noise (running headers, dates, names, "
            "table fragments).\n\n"
            "Build the table of contents. Return STRICT JSON only, exactly this shape:\n"
            '{"items":[{"title":"the heading, cleaned","page":N,"level":0}]}\n'
            "Rules: keep document order; use the [page N] the line came from; level 0 for top "
            "sections, 1-2 for subsections (numbering like 2.1 implies level 1); drop noise, "
            "duplicates, and body sentences; do not invent headings that are not in the list.\n\n"
            "SHORTLISTED LINES:\n%s\n\nJSON:" % (npages, listing))
        try:
            raw = _generate(prompt, model, num_predict=900, temperature=0.0, fmt="json", timeout=timeout)
            d = json.loads(raw or "{}")
            items.extend(d.get("items") or [])
        except Exception as e:
            if not items:
                return {"error": "Local AI error: %s" % e, "items": []}
    clean = _toc_sanitize(items, npages)
    if not clean:
        return {"error": "Couldn't find clear section headings in this document.", "items": []}
    return {"ok": True, "generated": True, "items": clean, "model": model,
            "pages": npages, "ms": int((time.time() - t0) * 1000)}


# ---------------------------------------------------------------- CLI
if __name__ == "__main__":
    import argparse, sys
    ap = argparse.ArgumentParser(description="PageHawk on-device document AI")
    ap.add_argument("cmd", choices=["ask", "summarize", "terms", "toc", "status", "explain"])
    ap.add_argument("--doc")
    ap.add_argument("--q", default="")
    ap.add_argument("--k", type=int, default=4)
    ap.add_argument("--passages")   # explain: path to a JSON file [{name,anchor,snippet}]
    ap.add_argument("--quiet", action="store_true")
    a = ap.parse_args()
    if a.cmd == "status":
        print(json.dumps(status())); sys.exit(0)
    if a.cmd == "explain":
        try:
            with open(a.passages, "r") as f:
                passages = json.load(f)
        except Exception as e:
            print(json.dumps({"error": "Could not read passages: %s" % e, "explanation": ""})); sys.exit(1)
        print(json.dumps(explain_passages(passages, a.q))); sys.exit(0)
    if not a.doc or not os.path.exists(a.doc):
        print(json.dumps({"error": "Document not found: %s" % a.doc})); sys.exit(1)
    if a.cmd == "ask":
        print(json.dumps(answer(a.doc, a.q, a.k)))
    elif a.cmd == "summarize":
        print(json.dumps(summarize(a.doc)))
    elif a.cmd == "terms":
        print(json.dumps(extract_terms(a.doc)))
    elif a.cmd == "toc":
        logger = (lambda m: None) if a.quiet else (lambda m: sys.stderr.write("· " + m + "\n"))
        print(json.dumps(toc(a.doc, log=logger)))
