#!/usr/bin/env python3
"""Invisible OCR text layer for scanned PDF pages (file-remediation Stage 2,
scanned/hybrid route — runs before Engine E2 tagging).

Per page with < 100 extractable chars (same threshold as remediate_assess's
scanned-page heuristic): render via PyMuPDF, OCR via macOS Vision
(VNRecognizeTextRequest, accurate mode), insert invisible text (render_mode=3)
at each recognized line's position. Vision returns normalized bottom-up
bboxes (origin bottom-left, like PDF's own coordinate space); converted here
to fitz's top-left page-point space.

Vision is macOS-only (pyobjc-framework-Vision) and lazy-imported: importing
this module never requires it; add_text_layer raises RuntimeError with a
clear message if it's unavailable.

Usage: python ocr_text_layer.py IN.pdf OUT.pdf [--dpi 200]
"""

from __future__ import annotations

import argparse
import json
import sys
import tempfile
from pathlib import Path

import fitz  # PyMuPDF

MIN_TEXT_CHARS = 100


def _import_vision():
    try:
        import Vision
        import Quartz
        from Foundation import NSURL
    except ImportError as exc:
        raise RuntimeError(
            "macOS Vision framework unavailable (pyobjc-framework-Vision not "
            "installed) — run `python scripts/setup_remediation.py --install`"
        ) from exc
    return Vision, Quartz, NSURL


def _ocr_page_image(png_path):
    """Run Vision text recognition on a PNG file. Returns
    [(text, x0, y0, w, h)] in normalized bottom-up fractions of the image."""
    Vision, Quartz, NSURL = _import_vision()
    url = NSURL.fileURLWithPath_(str(png_path))
    src = Quartz.CGImageSourceCreateWithURL(url, None)
    cgimg = Quartz.CGImageSourceCreateImageAtIndex(src, 0, None)

    request = Vision.VNRecognizeTextRequest.alloc().initWithCompletionHandler_(None)
    request.setRecognitionLevel_(Vision.VNRequestTextRecognitionLevelAccurate)
    handler = Vision.VNImageRequestHandler.alloc().initWithCGImage_options_(cgimg, None)
    ok, err = handler.performRequests_error_([request], None)
    if not ok:
        raise RuntimeError(f"Vision OCR failed: {err}")

    lines = []
    for obs in request.results() or []:
        top = obs.topCandidates_(1)
        if not top:
            continue
        text = str(top[0].string())
        box = obs.boundingBox()
        lines.append((text, box.origin.x, box.origin.y, box.size.width, box.size.height))
    return lines


def add_text_layer(pdf_path, out_path, dpi=200):
    """Add an invisible OCR text layer to every under-text page. Returns
    {"pages_ocred": [page_num...], "lines_added": N}."""
    doc = fitz.open(str(pdf_path))
    pages_ocred, lines_added = [], 0

    for i, page in enumerate(doc):
        if len(page.get_text() or "") >= MIN_TEXT_CHARS:
            continue

        pix = page.get_pixmap(dpi=dpi)
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            tmp_path = tmp.name
        try:
            pix.save(tmp_path)
            lines = _ocr_page_image(tmp_path)
        finally:
            Path(tmp_path).unlink(missing_ok=True)

        page_w, page_h = page.rect.width, page.rect.height
        for text, bx, by, bw, bh in lines:
            if not text.strip():
                continue
            x0 = bx * page_w
            box_h_pt = bh * page_h
            top_pt = (1 - by - bh) * page_h  # top edge, top-down
            baseline_y = top_pt + box_h_pt * 0.8
            fontsize = max(box_h_pt * 0.85, 1.0)
            page.insert_text((x0, baseline_y), text, fontsize=fontsize, render_mode=3)
            lines_added += 1
        pages_ocred.append(i + 1)

    doc.save(str(out_path))
    doc.close()
    return {"pages_ocred": pages_ocred, "lines_added": lines_added}


def main():
    parser = argparse.ArgumentParser(description="Add an invisible OCR text layer to a PDF")
    parser.add_argument("input")
    parser.add_argument("output")
    parser.add_argument("--dpi", type=int, default=200)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    try:
        result = add_text_layer(args.input, args.output, dpi=args.dpi)
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        sys.exit(2)

    print(json.dumps(result, indent=2) if args.json else
          f"OCR'd {len(result['pages_ocred'])} page(s), {result['lines_added']} line(s)")


if __name__ == "__main__":
    main()
