#!/usr/bin/env python3
"""Redaction-failure detector v2 — RENDER-BASED ground truth.

v1 failed because it trusted the PDF operator list: any dark filled rect was a
"redaction mark" and any extractable text geometrically under it was a "leak",
so footnote rules / table borders / letterheads near visible text produced
false positives.

v2 principle: a word is a leak ONLY if it is extractable from the text layer
but INVISIBLE in the rendered page because an opaque uniform dark region covers
it. Everything is judged on rendered pixels:

  1. Render page at 150 dpi grayscale.
  2. Dark mask = pixels < DARK_THR. Connected components; a component is a
     "redaction box" candidate iff it is big enough to hide text
     (>= MIN_BOX_W x MIN_BOX_H points), rectangular (bbox fill ratio >=
     RECTANGULARITY), and nearly uniform (std of its bbox pixels <= UNIFORM_STD).
     White-on-dark headers fail uniformity/fill (glyph holes); thin rules fail
     the size gate; photos fail uniformity.
  3. A text-layer word is a LEAK iff >= COVER_FRAC of its (slightly shrunk)
     bbox pixels lie inside such a box component AND the pixels over the word
     bbox are themselves uniform dark (std <= UNIFORM_STD, mean <= DARK_THR) —
     i.e. no visible glyphs where the word claims to be.

Aggregate-only: counts and per-doc booleans; recovered text is NEVER stored.
"""
import json, os, sys, traceback
import numpy as np
import fitz
from scipy import ndimage

DPI = 150
SCALE = DPI / 72.0
DARK_THR = 70          # gray level counted as "dark"
UNIFORM_STD = 9.0      # max std-dev for an opaque uniform fill
MIN_BOX_W_PT = 30      # min box width in points (~0.4")
MIN_BOX_H_PT = 8       # min box height in points (a text line is ~8-12pt)
RECTANGULARITY = 0.90  # component area / bbox area
COVER_FRAC = 0.90      # fraction of word pixels inside a box component
MAX_PAGES = 120   # raised from 60: matches the collector's 120-page cap so no document is partially analyzed
MIN_WORD_ALNUM = 2     # ignore 1-char fragments / artifacts

def analyze_page(page):
    # Normalise page rotation FIRST. get_text("words") reports coordinates in
    # the page's unrotated space, while get_pixmap() renders with /Rotate
    # applied. On a rotated page the two spaces disagree, so word boxes get
    # compared against the wrong pixels — which manufactures false positives
    # (a visible word whose unrotated coordinates happen to land inside a black
    # region of the rendered page). ~2% of pages in this corpus are rotated.
    if page.rotation:
        page.set_rotation(0)
    words = page.get_text("words")  # x0,y0,x1,y1,word,block,line,wordno
    out = {"has_text": bool(words), "boxes": 0, "leak_words": 0, "leak_chars": 0}
    if not words:
        return out
    pm = page.get_pixmap(matrix=fitz.Matrix(SCALE, SCALE), colorspace=fitz.csGRAY, alpha=False)
    img = np.frombuffer(pm.samples, dtype=np.uint8).reshape(pm.height, pm.width)
    dark = img < DARK_THR
    if not dark.any():
        return out
    lbl, n = ndimage.label(dark)
    if n == 0:
        return out
    objs = ndimage.find_objects(lbl)
    min_w = MIN_BOX_W_PT * SCALE
    min_h = MIN_BOX_H_PT * SCALE
    page_area = img.shape[0] * img.shape[1]
    box_ids = set()
    for i, sl in enumerate(objs, start=1):
        if sl is None:
            continue
        h = sl[0].stop - sl[0].start
        w = sl[1].stop - sl[1].start
        if w < min_w or h < min_h:
            continue
        if w * h > 0.7 * page_area:      # page-background fill, not a box
            continue
        comp = (lbl[sl] == i)
        if comp.sum() / (w * h) < RECTANGULARITY:
            continue
        if img[sl].std() > UNIFORM_STD:  # visible glyphs / photo texture
            continue
        box_ids.add(i)
    out["boxes"] = len(box_ids)
    if not box_ids:
        return out
    box_mask = np.isin(lbl, list(box_ids))
    H, W = img.shape
    for x0, y0, x1, y1, w, *_ in words:
        if sum(c.isalnum() for c in w) < MIN_WORD_ALNUM:
            continue
        # shrink 1px inward to avoid edge antialiasing
        px0, py0 = int(x0 * SCALE) + 1, int(y0 * SCALE) + 1
        px1, py1 = int(x1 * SCALE) - 1, int(y1 * SCALE) - 1
        px0, py0 = max(px0, 0), max(py0, 0)
        px1, py1 = min(px1, W), min(py1, H)
        if px1 - px0 < 3 or py1 - py0 < 3:
            continue
        region = box_mask[py0:py1, px0:px1]
        if region.mean() < COVER_FRAC:
            continue
        pix = img[py0:py1, px0:px1]
        if pix.mean() > DARK_THR or pix.std() > UNIFORM_STD:
            continue  # something visible there — not hidden
        out["leak_words"] += 1
        out["leak_chars"] += len(w)
    return out

def analyze_pdf(path):
    doc = fitz.open(path)
    res = {"pages": doc.page_count, "pages_with_boxes": 0, "leak_pages": 0,
           "boxes": 0, "leak_words": 0, "leak_chars": 0, "text_pages": 0}
    for pno in range(min(doc.page_count, MAX_PAGES)):
        p = analyze_page(doc[pno])
        if p["has_text"]:
            res["text_pages"] += 1
        if p["boxes"]:
            res["pages_with_boxes"] += 1
            res["boxes"] += p["boxes"]
        if p["leak_words"]:
            res["leak_pages"] += 1
            res["leak_words"] += p["leak_words"]
            res["leak_chars"] += p["leak_chars"]
    doc.close()
    return res

def main():
    corpus = sys.argv[1]
    out_path = sys.argv[2]
    results = {}
    if os.path.exists(out_path):
        results = json.load(open(out_path))
    files = sorted(f for f in os.listdir(corpus) if f.endswith(".pdf"))
    for i, f in enumerate(files):
        if f in results:
            continue
        try:
            results[f] = analyze_pdf(os.path.join(corpus, f))
        except Exception as e:
            results[f] = {"error": str(e).split("\n")[0][:200]}
        if (i + 1) % 20 == 0:
            json.dump(results, open(out_path, "w"))
            print(f"{i+1}/{len(files)}", flush=True)
    json.dump(results, open(out_path, "w"), indent=1)
    print(f"done {len(files)}")

if __name__ == "__main__":
    main()