#!/usr/bin/env python3
"""Struct-tree reading-order text extraction (file-remediation Stage 2/3 verification).

opendataloader-pdf's E2 re-tag path (remediate_pdf.py) can produce a
veraPDF/PDF-UA-compliant file that still gets reading order wrong — content
correctly marked /Artifact in the source (excluded from the tag tree, never
announced by a screen reader) can come back as real tagged content in the
re-tagged struct tree. Neither existing hard gate catches this: veraPDF
checks tag-tree WELL-FORMEDNESS, not WHICH content got tagged, and pixel
fidelity only checks rendered pixels, which never change (/Artifact vs a
/Tag around the same paint ops looks identical on the page).

struct_order_text()/struct_order_runs() give remediate_pdf.py an actual
signal: walk /StructTreeRoot depth-first (== reading order for a tagged
PDF, per PDF 32000-1 14.7.2) and resolve every MCID leaf to its real page
text via mcid_text.page_mcid_text. Comparing before vs. after struct order
surfaces exactly the leaked-artifact case this module exists to catch —
see remediate_pdf.py's post-retag artifact-preservation pass.

Usage: python reading_order.py FILE [N]
       python reading_order.py --diff BEFORE AFTER
Exit: 0 ok, 2 error.
"""

from __future__ import annotations

import argparse
import difflib
import re
import sys
from pathlib import Path

SCRIPTS_DIR = Path(__file__).resolve().parent
if str(SCRIPTS_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPTS_DIR))

import pikepdf

import mcid_text

_WS_RE = re.compile(r"\s+")


def struct_order_runs(pdf) -> list:
    """[(page_1based, mcid, text), ...] in struct-tree depth-first order —
    reading order for a tagged PDF. `pdf` may be a path/str (opened and
    closed here) or an already-open pikepdf.Pdf (e.g. the in-memory object
    remediate_pdf.py is mid-editing — no round-trip through disk needed).
    [] for an untagged PDF (no /StructTreeRoot). Never raises: an unopenable
    file also yields []. Cycle-guarded via a seen-objgen set. Text runs are
    whitespace-normalized; empty/whitespace-only runs are skipped."""
    owns = isinstance(pdf, (str, Path))
    try:
        doc = pikepdf.open(str(pdf)) if owns else pdf
    except Exception:  # noqa: BLE001 - best-effort report input, must not raise
        return []
    try:
        return _walk_struct_tree(doc)
    finally:
        if owns:
            doc.close()


def struct_order_text(pdf) -> list:
    """Whitespace-normalized text runs in struct-tree (reading) order.
    Convenience wrapper over struct_order_runs() for callers that only need
    the text, not the (page, mcid) identity."""
    return [text for _, _, text in struct_order_runs(pdf)]


def _walk_struct_tree(doc) -> list:
    root = doc.Root.get("/StructTreeRoot")
    if root is None:
        return []

    page_index = {}
    for i, page in enumerate(doc.pages):
        try:
            page_index[page.obj.objgen] = i + 1
        except Exception:  # noqa: BLE001
            continue

    mcid_cache: dict = {}

    def page_text(pg):
        if pg is None:
            return {}
        try:
            key = pg.objgen
        except Exception:  # noqa: BLE001
            return {}
        if key not in mcid_cache:
            mcid_cache[key] = mcid_text.page_mcid_text(pg)
        return mcid_cache[key]

    def page_number(pg):
        if pg is None:
            return None
        try:
            return page_index.get(pg.objgen)
        except Exception:  # noqa: BLE001
            return None

    out: list = []
    seen: set = set()

    def emit(pg, mcid):
        text = _WS_RE.sub(" ", page_text(pg).get(mcid, "")).strip()
        if text:
            out.append((page_number(pg), mcid, text))

    def walk(node, pg):
        if node is None:
            return
        if isinstance(node, pikepdf.Array):
            for item in node:
                walk(item, pg)
            return
        if not hasattr(node, "get"):
            return
        try:
            key = node.objgen if node.is_indirect else id(node)
        except Exception:  # noqa: BLE001
            key = id(node)
        if key in seen:
            return
        seen.add(key)
        pg = node.get("/Pg", pg)
        k = node.get("/K")
        if k is None:
            return
        items = k if isinstance(k, pikepdf.Array) else [k]
        for item in items:
            if isinstance(item, int):
                emit(pg, int(item))
            elif isinstance(item, pikepdf.Dictionary):
                if "/MCID" in item and "/S" not in item:
                    emit(item.get("/Pg", pg), int(item["/MCID"]))
                else:
                    walk(item, pg)
            elif isinstance(item, pikepdf.Array):
                walk(item, pg)

    walk(root, None)
    return out


def _diff(before_path, after_path, n=5) -> int:
    before = struct_order_text(before_path)
    after = struct_order_text(after_path)
    sm = difflib.SequenceMatcher(None, before, after)
    print(f"similarity ratio: {sm.ratio():.4f}  ({len(before)} -> {len(after)} runs)")
    opcodes = [op for op in sm.get_opcodes() if op[0] != "equal"]
    for tag, i1, i2, j1, j2 in opcodes[:n]:
        print(f"  {tag}: before[{i1}:{i2}]={before[i1:i2]!r} after[{j1}:{j2}]={after[j1:j2]!r}")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description="Struct-tree reading-order text extraction")
    parser.add_argument("file", nargs="?", help="PDF to inspect")
    parser.add_argument("n", nargs="?", type=int, default=8, help="number of runs to print")
    parser.add_argument("--diff", nargs=2, metavar=("BEFORE", "AFTER"))
    args = parser.parse_args()

    try:
        if args.diff:
            return _diff(args.diff[0], args.diff[1])
        if not args.file:
            print("ERROR: FILE required (or use --diff BEFORE AFTER)", file=sys.stderr)
            return 2
        runs = struct_order_text(args.file)
        print(f"{len(runs)} struct-order text runs")
        for text in runs[: args.n]:
            print("  ", text[:90])
        return 0
    except Exception as exc:  # noqa: BLE001
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    sys.exit(main())
