#!/usr/bin/env python3
"""Rendering-regression harness for the PDFium engine.

Renders every page (capped) of every PDF in the committed corpus and compares it
to a frozen baseline using SSIM. The baseline is the engine's own output at the
moment it was approved — so a regression here means *this* build renders
differently from the blessed build, not that it disagrees with Acrobat.

    python tests/regression/run.py --update     # (re)generate the frozen baseline
    python tests/regression/run.py              # check current renders vs baseline

Exit code is non-zero if any page regresses below the SSIM threshold, so it
drops straight into CI.
"""

from __future__ import annotations

import argparse
import io
import sys
from pathlib import Path

import numpy as np
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parent))
from ssim import ssim  # noqa: E402

ROOT = Path(__file__).resolve().parents[2]
DEFAULT_CORPUS = ROOT / "corpus"
DEFAULT_BASELINE = Path(__file__).resolve().parent / "baseline"
# corpus subdirs that are *not* part of the frozen baseline (large/downloaded).
EXCLUDE_DIRS = {"fuzz"}


def _import_engine():
    # The compiled extension lives in gateway/ after scripts/build_cpp.ps1.
    sys.path.insert(0, str(ROOT / "gateway"))
    import pdfengine  # noqa: PLC0415

    return pdfengine


def _render_gray(page, dpi: int) -> np.ndarray:
    png = page.render(dpi).data
    img = Image.open(io.BytesIO(png)).convert("L")
    return np.asarray(img)


def _corpus_pdfs(corpus: Path):
    for p in sorted(corpus.rglob("*.pdf")):
        if any(part in EXCLUDE_DIRS for part in p.relative_to(corpus).parts):
            continue
        yield p


def _key(rel: Path, page_index: int) -> str:
    return f"{rel.as_posix().replace('/', '__')}__p{page_index}"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--update", action="store_true", help="write the frozen baseline instead of checking")
    ap.add_argument("--sweep", action="store_true",
                    help="render-stability sweep over a large corpus (no baseline); "
                         "passes as long as the engine never crashes")
    ap.add_argument("--corpus", type=Path, default=DEFAULT_CORPUS)
    ap.add_argument("--baseline", type=Path, default=DEFAULT_BASELINE)
    ap.add_argument("--dpi", type=int, default=72)
    ap.add_argument("--max-pages", type=int, default=2, help="pages rendered per document")
    ap.add_argument("--threshold", type=float, default=0.990, help="min SSIM to pass")
    args = ap.parse_args()

    pdfengine = _import_engine()
    args.baseline.mkdir(parents=True, exist_ok=True)

    failures: list[tuple[str, float]] = []
    checked = skipped = updated = rendered = 0

    for pdf in _corpus_pdfs(args.corpus):
        rel = pdf.relative_to(args.corpus)
        try:
            doc = pdfengine.PdfDocument.load_from_memory(pdf.read_bytes(), "")
            n = min(doc.page_count, args.max_pages)
        except Exception as exc:  # encrypted / intentionally-malformed fixtures
            print(f"  skip  {rel}  ({exc})")
            skipped += 1
            continue

        for i in range(n):
            try:
                cur = _render_gray(doc.get_page(i), args.dpi)
            except Exception as exc:
                print(f"  skip  {rel} p{i}  (render: {exc})")
                skipped += 1
                continue

            if args.sweep:
                # Surviving the render is the whole test; a crash kills the process.
                rendered += 1
                continue

            ref_path = args.baseline / f"{_key(rel, i)}.png"
            if args.update:
                Image.fromarray(cur).save(ref_path)
                updated += 1
                continue

            if not ref_path.exists():
                print(f"  NEW   {rel} p{i}  (no baseline — run --update)")
                failures.append((f"{rel} p{i}", -1.0))
                continue

            ref = np.asarray(Image.open(ref_path).convert("L"))
            score = ssim(cur, ref)
            checked += 1
            mark = "ok   " if score >= args.threshold else "FAIL "
            if score < args.threshold:
                failures.append((f"{rel} p{i}", score))
            print(f"  {mark} {rel} p{i}  SSIM={score:.4f}")

    print()
    if args.sweep:
        print(f"Sweep complete: {rendered} pages rendered, {skipped} skipped (graceful). "
              f"No crash — engine is render-stable over this corpus.")
        return 0
    if args.update:
        print(f"Baseline updated: {updated} images written to {args.baseline}")
        return 0

    print(f"Checked {checked} pages, skipped {skipped}, {len(failures)} regression(s).")
    for name, score in failures:
        tag = "missing baseline" if score < 0 else f"SSIM={score:.4f}"
        print(f"  REGRESSION: {name}  ({tag})")
    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(main())