#!/usr/bin/env python3
"""Benchmark OCR engines on pages of a native-text PDF, using its text layer as ground truth.

    python3 scripts/bench_ocr.py --pdf source.pdf --pages 12,20,30 --engine tesseract \
        --engine "vision:http://localhost:8098|baidu/Unlimited-OCR" --engine docling --engine paddleocr

The PDF's own text layer is what a perfect transcription of the rendered page would say, so no
hand-labelled data is needed. Two scores per page, both on lower-cased word tokens:
  * bag F1   -- order-insensitive overlap (robust to differing column reading order)
  * seq sim  -- difflib ratio over the token sequence (penalises scrambled order)
Also wall-clock seconds per page and the peak resident memory of THIS process (a remote engine's own
footprint is not visible here; read it from `docker stats`).

Engines are optional: one whose library or server is missing is reported as unavailable, not an error.
Read-only: renders into a temp directory and writes nothing else unless --out is given.
"""
from __future__ import annotations

import argparse
import difflib
import json
import re
import resource
import statistics
import sys
import tempfile
import time
from collections import Counter
from pathlib import Path

_TOKEN = re.compile(r"[\wÄÖÜäöüß]+", re.UNICODE)


def tokens(text: str) -> list[str]:
    # Hyphenation at line ends ("zu-\nführen") is joined so it does not count against an engine
    # that correctly re-joins it (or for one that keeps the break).
    text = re.sub(r"-\s*\n\s*", "", text or "")
    return [t.lower() for t in _TOKEN.findall(text)]


def bag_f1(ref: list[str], hyp: list[str]) -> float:
    if not ref or not hyp:
        return 0.0
    overlap = sum((Counter(ref) & Counter(hyp)).values())
    if not overlap:
        return 0.0
    p, r = overlap / len(hyp), overlap / len(ref)
    return 2 * p * r / (p + r)


def seq_sim(ref: list[str], hyp: list[str]) -> float:
    if not ref or not hyp:
        return 0.0
    return difflib.SequenceMatcher(None, ref, hyp, autojunk=False).ratio()


def render(pdf: Path, page: int, dpi: int, out: Path) -> Path:
    import pymupdf
    with pymupdf.open(str(pdf)) as doc:
        doc[page - 1].get_pixmap(dpi=dpi).save(str(out))
    return out


def truth(pdf: Path, page: int) -> str:
    import pymupdf
    with pymupdf.open(str(pdf)) as doc:
        return doc[page - 1].get_text()


def pick_pages(pdf: Path, count: int, min_words: int) -> list[int]:
    """Evenly spaced pages with enough text to score meaningfully (adverts and photo pages drop out)."""
    import pymupdf
    with pymupdf.open(str(pdf)) as doc:
        eligible = [n + 1 for n in range(len(doc)) if len(tokens(doc[n].get_text())) >= min_words]
    if len(eligible) <= count:
        return eligible
    step = len(eligible) / count
    return [eligible[int(i * step)] for i in range(count)]


# ── engines: each is a callable (png_path) -> transcribed text, or raises Unavailable ─────────────

class Unavailable(Exception):
    pass


def engine_tesseract(lang: str = "deu+eng"):
    import shutil
    import subprocess
    if not shutil.which("tesseract"):
        raise Unavailable("tesseract is not installed here")

    def run(png: Path) -> str:
        return subprocess.run(["tesseract", str(png), "stdout", "-l", lang], capture_output=True,
                              text=True, timeout=180, check=True).stdout
    return run


def engine_vision(spec: str):
    """`vision:<base_url>|<model>` -- an OpenAI-compatible vision server, called exactly as the app calls it."""
    try:
        base, model = spec.split("|", 1)
    except ValueError:
        raise Unavailable("use vision:<url>|<model>")
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    try:
        from core import magazine as mag
    except Exception as e:  # the app's own module needs its environment
        raise Unavailable(f"cannot import the app's OCR client ({type(e).__name__})")
    import requests
    try:
        requests.get(base.rstrip("/") + ("/v1/models" if "/v1" not in base else "/models"), timeout=5).raise_for_status()
    except Exception as e:
        raise Unavailable(f"server not reachable ({type(e).__name__})")

    def run(png: Path) -> str:
        regions = mag.vision_ocr_page(png, base, model, timeout=300.0)
        if regions is None:
            raise RuntimeError("engine returned nothing")
        return "\n".join(str(r.get("text") or "") for r in regions)
    return run


def engine_docling():
    try:
        from docling.document_converter import DocumentConverter
    except Exception as e:
        raise Unavailable(f"docling not installed ({type(e).__name__})")
    conv = DocumentConverter()

    def run(png: Path) -> str:
        return conv.convert(str(png)).document.export_to_text()
    return run


def engine_paddleocr(lang: str = "de"):
    try:
        from paddleocr import PaddleOCR
    except Exception as e:
        raise Unavailable(f"paddleocr not installed ({type(e).__name__})")
    try:  # 3.x: page orientation/unwarping models are extra work a straight scan does not need
        ocr = PaddleOCR(lang=lang, use_doc_orientation_classify=False, use_doc_unwarping=False,
                        use_textline_orientation=False)
    except TypeError:  # 2.x
        ocr = PaddleOCR(lang=lang, use_angle_cls=False)

    def run(png: Path) -> str:
        lines: list[str] = []
        res = ocr.predict(str(png)) if hasattr(ocr, "predict") else ocr.ocr(str(png))
        for page in res or []:
            if isinstance(page, dict) or hasattr(page, "get"):          # 3.x result object
                lines.extend(str(t) for t in (page.get("rec_texts") or []))
                continue
            for item in page or []:                                      # 2.x [[box, (text, score)], ...]
                try:
                    lines.append(item[1][0])
                except (IndexError, TypeError):
                    pass
        return "\n".join(lines)
    return run


def build_engine(spec: str):
    if spec == "tesseract":
        return engine_tesseract()
    if spec.startswith("vision:"):
        return engine_vision(spec[len("vision:"):])
    if spec == "docling":
        return engine_docling()
    if spec == "paddleocr":
        return engine_paddleocr()
    raise Unavailable(f"unknown engine {spec!r}")


def peak_rss_mb() -> float:
    return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--pdf", required=True, type=Path)
    ap.add_argument("--pages", default="", help="comma-separated 1-based page numbers (default: evenly spaced)")
    ap.add_argument("--count", type=int, default=8, help="pages to sample when --pages is not given")
    ap.add_argument("--min-words", type=int, default=250)
    ap.add_argument("--dpi", type=int, default=110, help="110 matches the app's ingest default")
    ap.add_argument("--engine", action="append", required=True)
    ap.add_argument("--out", type=Path, help="write the full results as JSON")
    args = ap.parse_args()

    pages = [int(p) for p in args.pages.split(",") if p.strip()] or pick_pages(args.pdf, args.count, args.min_words)
    print(f"pages: {pages}")
    results = []
    with tempfile.TemporaryDirectory(prefix="bench-ocr-") as tmp:
        pngs = {p: render(args.pdf, p, args.dpi, Path(tmp) / f"p{p}.png") for p in pages}
        refs = {p: tokens(truth(args.pdf, p)) for p in pages}
        for spec in args.engine:
            entry = {"engine": spec, "pages": []}
            try:
                base_rss = peak_rss_mb()
                t0 = time.perf_counter()
                run = build_engine(spec)
                entry["load_sec"] = round(time.perf_counter() - t0, 1)
            except Unavailable as e:
                entry["unavailable"] = str(e)
                results.append(entry)
                print(f"{spec}: unavailable — {e}")
                continue
            for p in pages:
                t0 = time.perf_counter()
                try:
                    hyp = tokens(run(pngs[p]))
                    err = None
                except Exception as e:
                    hyp, err = [], f"{type(e).__name__}: {e}"[:200]
                sec = time.perf_counter() - t0
                entry["pages"].append({
                    "page": p, "sec": round(sec, 1), "ref_words": len(refs[p]), "hyp_words": len(hyp),
                    "f1": round(bag_f1(refs[p], hyp), 4), "seq": round(seq_sim(refs[p], hyp), 4), "error": err})
                print(f"  {spec} p{p}: f1={entry['pages'][-1]['f1']:.3f} seq={entry['pages'][-1]['seq']:.3f} {sec:.1f}s"
                      + (f" ERROR {err}" if err else ""))
            entry["process_rss_growth_mb"] = round(peak_rss_mb() - base_rss)
            results.append(entry)

    print("\n| engine | pages ok | mean F1 | mean seq | median s/page | RSS growth MB |")
    print("|---|---:|---:|---:|---:|---:|")
    for e in results:
        ok = [r for r in e.get("pages", []) if not r["error"]]
        if not ok:
            print(f"| {e['engine']} | 0/{len(e.get('pages', []))} | – | – | – | – | {e.get('unavailable', '')}")
            continue
        print(f"| {e['engine']} | {len(ok)}/{len(e['pages'])} | {statistics.mean(r['f1'] for r in ok):.3f} | "
              f"{statistics.mean(r['seq'] for r in ok):.3f} | {statistics.median(r['sec'] for r in ok):.1f} | "
              f"{e.get('process_rss_growth_mb', 0)} |")
    if args.out:
        args.out.write_text(json.dumps({"pdf": args.pdf.name, "pages": pages, "dpi": args.dpi, "results": results},
                                       ensure_ascii=False, indent=1), "utf-8")
    return 0


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