#!/usr/bin/env python3
"""Losslessly shrink scanned PDFs by re-encoding CCITT G4 page images as
JBIG2 generic regions.

Patent documents, file wrappers and many journal scans are 300 dpi 1-bit
page images stored with CCITT Group 4 compression. JBIG2 generic-region
coding stores the very same pixels in roughly half the space; on a mixed
prior-art collection files came out about 45% smaller on average.

Safety properties, in order of importance:
  * Generic-region coding only. JBIG2 symbol/text mode (which can substitute
    look-alike glyphs, the cause of the well-known scanner digit-swap bug)
    is never used.
  * Every output is render-verified against its input, page by page, at
    300 dpi. A result is kept only if the page count matches and every page
    renders byte-identical.
  * A result is kept only if it is actually smaller.
  * By default the original is untouched and a new file is written next to
    it (NAME.shrunk.pdf) or under --output-dir. Overwriting needs --in-place,
    and --in-place needs either --backup-dir or an explicit --no-backup.
  * Modification times are preserved.

Usage
-----
    pdf-shrink scan.pdf                          # writes scan.shrunk.pdf
    pdf-shrink -r docs/ --output-dir smaller/    # mirror a tree
    pdf-shrink -r docs/ --dry-run --report r.csv # measure only
    pdf-shrink -r docs/ --in-place --backup-dir ~/pdf-originals --jobs 4
"""
import argparse
import csv
import hashlib
import os
import random
import shutil
import subprocess
import sys
import tempfile
import time

# Folders that sync to a cloud service. Backups must never be written inside
# one: the backup has to survive a bad sync, and a copy inside the synced tree
# would also double the upload.
SYNC_DIRS = ["~/Dropbox", "~/Google Drive", "~/GoogleDrive", "~/OneDrive",
             "~/Nextcloud", "~/ownCloud", "~/Sync", "~/iCloudDrive", "~/pCloudDrive"]


def find_jbig2(explicit=None):
    """Path to the jbig2enc command-line encoder, or None."""
    for cand in (explicit, os.environ.get("PDF_SHRINK_JBIG2")):
        if cand:
            return cand if os.access(cand, os.X_OK) else None
    here = os.path.dirname(os.path.abspath(__file__))
    for name in ("jbig2", "jbig2enc"):
        local = os.path.join(here, name)
        if os.path.isfile(local) and os.access(local, os.X_OK):
            return local
        found = shutil.which(name)
        if found:
            return found
    return None


def is_within(path, parent):
    """True if path is parent or lies underneath it."""
    try:
        path = os.path.abspath(path)
        parent = os.path.abspath(parent)
        return os.path.commonpath([path, parent]) == parent
    except ValueError:
        return False


def jbig2_repack(src, dst, jbig2):
    """Rewrite src to dst with CCITT G4 images re-encoded as JBIG2. Returns count."""
    import pikepdf
    pdf = pikepdf.open(src)
    tmp = tempfile.mkdtemp(prefix="pdf-shrink-")
    changed = 0
    try:
        for page in pdf.pages:
            xo = page.get("/Resources", {}).get("/XObject", {})
            for name, obj in list(xo.items()):
                try:
                    if obj.get("/Subtype") != "/Image":
                        continue
                    filt = obj.get("/Filter")
                    fl = ([str(f) for f in (filt if isinstance(filt, pikepdf.Array) else [filt])]
                          if filt is not None else [])
                    if "/CCITTFaxDecode" not in fl or int(obj.get("/BitsPerComponent", 8)) != 1:
                        continue
                    pil = pikepdf.PdfImage(obj).as_pil_image().convert("1")
                    p = os.path.join(tmp, "i.pbm")
                    pil.save(p)
                    # -d duplicate-line removal, -p PDF-ready. No -s: never symbol mode.
                    r = subprocess.run([jbig2, "-d", "-p", p], capture_output=True)
                    if r.returncode != 0 or not r.stdout:
                        continue
                    if len(r.stdout) >= len(bytes(obj.get_raw_stream_buffer())):
                        continue
                    w, h = pil.size
                    new = pikepdf.Stream(pdf, r.stdout)
                    # Carry over every key from the original image dictionary
                    # except the ones describing the old encoding. Rebuilding the
                    # dict from scratch silently drops rendering-relevant flags,
                    # /Interpolate in particular, which changes how a small scan
                    # is upscaled onto a large page.
                    for k in obj.keys():
                        # /Decode is deliberately NOT carried over: as_pil_image()
                        # already applied it when producing the bitmap above, so
                        # copying it would invert the image a second time.
                        if k in ("/Filter", "/DecodeParms", "/Length", "/Decode"):
                            continue
                        try:
                            new[k] = obj[k]
                        except Exception:
                            pass
                    new.Type = pikepdf.Name("/XObject")
                    new.Subtype = pikepdf.Name("/Image")
                    new.Width, new.Height, new.BitsPerComponent = w, h, 1
                    if new.get("/ColorSpace") is None and not new.get("/ImageMask"):
                        new.ColorSpace = pikepdf.Name("/DeviceGray")
                    new.Filter = pikepdf.Name("/JBIG2Decode")
                    xo[name] = pdf.make_indirect(new)
                    changed += 1
                except Exception:
                    continue
        pdf.save(dst, object_stream_mode=pikepdf.ObjectStreamMode.generate,
                 compress_streams=True, recompress_flate=True)
    finally:
        shutil.rmtree(tmp, ignore_errors=True)
        pdf.close()
    return changed


def page_count(path):
    try:
        r = subprocess.run(["pdfinfo", path], capture_output=True, timeout=120, text=True)
        for line in r.stdout.splitlines():
            if line.startswith("Pages:"):
                return int(line.split()[1])
    except Exception:
        pass
    return None


def max_page_megapixels(path, dpi):
    """Largest page area in megapixels once rendered at `dpi`, or None."""
    try:
        r = subprocess.run(["pdfinfo", "-f", "1", "-l", "100000", path],
                           capture_output=True, timeout=180, text=True)
    except Exception:
        return None
    best = 0.0
    for line in r.stdout.splitlines():
        # "Page N size: 612 x 792 pts" (or "Page size:" for single-page output)
        if "size:" not in line or " x " not in line:
            continue
        try:
            wh = line.split("size:", 1)[1].strip().split(" x ")
            w = float(wh[0])
            h = float(wh[1].split()[0])
        except (ValueError, IndexError):
            continue
        best = max(best, (w / 72 * dpi) * (h / 72 * dpi) / 1e6)
    return best or None


def render_page_digest(path, page, dpi, workdir, scale_mode=False):
    """Render exactly one page and return its SHA-256, or None.

    One page at a time keeps peak temp usage bounded. Rendering a whole
    document up front can mean tens of GB for large-format drawing sheets
    (a 10667x13750 page is ~147 MB uncompressed).
    """
    out = os.path.join(workdir, "p")
    geom = ["-scale-to", str(dpi)] if scale_mode else ["-r", str(dpi)]
    try:
        r = subprocess.run(["pdftoppm"] + geom + ["-gray",
                            "-f", str(page), "-l", str(page), path, out],
                           capture_output=True, timeout=900)
    except subprocess.TimeoutExpired:
        return None
    if r.returncode != 0:
        return None
    files = [f for f in os.listdir(workdir) if f.startswith("p")]
    if len(files) != 1:
        for f in files:
            os.unlink(os.path.join(workdir, f))
        return None
    fp = os.path.join(workdir, files[0])
    h = hashlib.sha256()
    with open(fp, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    os.unlink(fp)
    return h.hexdigest()


def verify_identical(a, b, dpi, scale_mode=False):
    """Compare two PDFs page by page, holding at most two pages on disk."""
    na, nb = page_count(a), page_count(b)
    if na is None or nb is None:
        return False, "page count unreadable"
    if na != nb:
        return False, f"page count {na} vs {nb}"
    if na == 0:
        return False, "no pages"
    with tempfile.TemporaryDirectory() as ta, tempfile.TemporaryDirectory() as tb:
        for pg in range(1, na + 1):
            ha = render_page_digest(a, pg, dpi, ta, scale_mode)
            hb = render_page_digest(b, pg, dpi, tb, scale_mode)
            if ha is None or hb is None:
                return False, f"render failed on page {pg}"
            if ha != hb:
                return False, f"pixel mismatch on page {pg}"
    return True, f"identical ({na} pages)"


def destination(path, args):
    """Where the shrunk copy of `path` is written."""
    if args.in_place:
        return path
    if args.output:
        return os.path.abspath(args.output)
    if args.output_dir:
        return os.path.join(args.output_dir, os.path.relpath(path, args.root))
    stem, ext = os.path.splitext(path)
    return stem + args.suffix + ext


def process(path, args):
    size_in = os.path.getsize(path)
    st = os.stat(path)
    if args.max_page_mp:
        mp = max_page_megapixels(path, args.verify_dpi)
        if mp and mp > args.max_page_mp:
            return (path, size_in, size_in, "SKIP", f"page {mp:.0f} MP > {args.max_page_mp:.0f} MP cap")
    with tempfile.TemporaryDirectory(prefix="pdf-shrink-") as td:
        cand = os.path.join(td, "out.pdf")
        try:
            n = jbig2_repack(path, cand, args.jbig2)
        except Exception as e:
            return (path, size_in, size_in, "SKIP", f"repack error: {type(e).__name__}")
        if not os.path.exists(cand):
            return (path, size_in, size_in, "SKIP", "no output")
        size_out = os.path.getsize(cand)
        if n == 0:
            return (path, size_in, size_in, "SKIP", "no CCITT G4 images")
        if size_out >= size_in:
            return (path, size_in, size_in, "SKIP", f"not smaller ({n} imgs)")
        ok, why = verify_identical(path, cand, args.verify_scale or args.verify_dpi,
                                   scale_mode=bool(args.verify_scale))
        if not ok:
            return (path, size_in, size_in, "FAIL", why)
        if args.dry_run:
            return (path, size_in, size_out, "WOULD-SHRINK", f"{n} imgs verified")
        dst = destination(path, args)
        if args.in_place and args.backup_dir:
            bak = os.path.join(args.backup_dir, os.path.relpath(path, args.root))
            os.makedirs(os.path.dirname(bak), exist_ok=True)
            shutil.copy2(path, bak)
        os.makedirs(os.path.dirname(dst) or ".", exist_ok=True)
        tmp = dst + ".pdf-shrink-tmp"
        shutil.copyfile(cand, tmp)
        os.replace(tmp, dst)
        os.utime(dst, (st.st_atime, st.st_mtime))
        return (path, size_in, size_out, "SHRUNK", f"{n} imgs verified -> {dst}")


def collect(inputs, recursive, suffix):
    targets = []
    for item in inputs:
        item = os.path.abspath(item)
        if os.path.isdir(item):
            if not recursive:
                sys.exit(f"{item} is a directory; pass -r/--recursive to process it")
            for dirpath, _dirs, files in os.walk(item):
                for fn in files:
                    if fn.lower().endswith(".pdf") and not fn.lower().endswith(suffix.lower() + ".pdf"):
                        targets.append(os.path.join(dirpath, fn))
        elif os.path.isfile(item):
            targets.append(item)
        else:
            sys.exit(f"not found: {item}")
    return sorted(set(targets))


def build_parser():
    ap = argparse.ArgumentParser(
        prog="pdf-shrink",
        description="Losslessly shrink scanned PDFs (CCITT G4 -> JBIG2 generic region), "
                    "render-verifying every page before keeping a result.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""examples:
  pdf-shrink scan.pdf
  pdf-shrink -r docs/ --output-dir smaller/
  pdf-shrink -r docs/ --dry-run --report measure.csv
  pdf-shrink -r docs/ --in-place --backup-dir ~/pdf-originals --jobs 4""")
    ap.add_argument("inputs", nargs="+", help="PDF files, or directories with -r")
    ap.add_argument("-r", "--recursive", action="store_true",
                    help="descend into directories given as inputs")
    out = ap.add_argument_group("output (default: NAME.shrunk.pdf next to each input)")
    out.add_argument("-o", "--output", help="output file (single input file only)")
    out.add_argument("--output-dir", help="write results here, mirroring the input tree")
    out.add_argument("--suffix", default=".shrunk",
                     help="suffix for default output names (default: .shrunk)")
    out.add_argument("--in-place", action="store_true",
                     help="replace the input files; requires --backup-dir or --no-backup")
    out.add_argument("--backup-dir",
                     help="with --in-place: copy each original here first. Refused if it "
                          "resolves inside a cloud-synced folder or inside the input tree")
    out.add_argument("--no-backup", action="store_true",
                     help="with --in-place: keep no copy of the originals "
                          "(verification still runs)")
    ap.add_argument("--dry-run", action="store_true",
                    help="encode and verify, but write nothing")
    ap.add_argument("--jbig2", help="path to the jbig2enc encoder (default: $PDF_SHRINK_JBIG2, "
                                    "then jbig2 on PATH)")
    ap.add_argument("--verify-dpi", type=int, default=300,
                    help="verify by rendering at this fixed dpi (default: 300)")
    ap.add_argument("--verify-scale", type=int, default=0,
                    help="OFF BY DEFAULT, and it weakens the check. Renders with the long side "
                         "scaled to N px instead of a fixed dpi: much faster on large-format "
                         "pages, but blind to a dropped /Interpolate flag. Fixed --verify-dpi "
                         "is the trustworthy mode.")
    ap.add_argument("--limit", type=int, default=0,
                    help="process only N files, chosen at random, and extrapolate the result")
    ap.add_argument("--seed", type=int, default=1, help="random seed for --limit")
    ap.add_argument("--first", action="store_true",
                    help="with --limit, take the first N in path order instead of a random sample")
    ap.add_argument("--max-page-mp", type=float, default=0.0,
                    help="skip files whose largest page exceeds this many megapixels at "
                         "the verify dpi (0 = no limit)")
    ap.add_argument("--report", help="write a CSV row per file (path, bytes in/out, status, note)")
    ap.add_argument("--resume", action="store_true",
                    help="append to an existing --report, skipping files already recorded")
    ap.add_argument("--jobs", type=int, default=1,
                    help="worker processes. Each can hold a rendered page in memory.")
    return ap


def main(argv=None):
    ap = build_parser()
    args = ap.parse_args(argv)

    args.jbig2 = find_jbig2(args.jbig2)
    if not args.jbig2:
        sys.exit("pdf-shrink: jbig2enc not found. Install it (see README) or pass --jbig2 PATH.")
    for tool in ("pdfinfo", "pdftoppm"):
        if not shutil.which(tool):
            sys.exit(f"pdf-shrink: {tool} not found (sudo apt install poppler-utils)")
    try:
        import pikepdf  # noqa: F401
    except ImportError:
        sys.exit("pdf-shrink: needs pikepdf (sudo apt install python3-pikepdf)")

    targets = collect(args.inputs, args.recursive, args.suffix)
    dirs = [os.path.abspath(i) for i in args.inputs]
    if len(dirs) == 1:
        args.root = dirs[0] if os.path.isdir(dirs[0]) else os.path.dirname(dirs[0])
    else:
        args.root = os.path.commonpath([d if os.path.isdir(d) else os.path.dirname(d) for d in dirs])

    modes = sum(bool(x) for x in (args.output, args.output_dir, args.in_place))
    if modes > 1:
        ap.error("-o, --output-dir and --in-place are mutually exclusive")
    if args.output and len(targets) != 1:
        ap.error("-o/--output takes exactly one input file")
    if args.output and os.path.abspath(args.output) in targets:
        ap.error("-o/--output names the input file; use --in-place to overwrite")
    if (args.backup_dir or args.no_backup) and not args.in_place:
        ap.error("--backup-dir / --no-backup only apply with --in-place")
    if args.output_dir:
        args.output_dir = os.path.abspath(args.output_dir)
    if args.resume and not args.report:
        ap.error("--resume needs --report")

    if args.in_place and not args.dry_run:
        if args.backup_dir:
            args.backup_dir = os.path.abspath(os.path.realpath(os.path.expanduser(args.backup_dir)))
            bad = [d for d in SYNC_DIRS
                   if os.path.isdir(os.path.expanduser(d))
                   and is_within(args.backup_dir, os.path.realpath(os.path.expanduser(d)))]
            if bad:
                sys.exit(f"refusing: --backup-dir {args.backup_dir} is inside a synced folder "
                         f"({bad[0]}). Pick a path on local disk.")
            if is_within(args.backup_dir, args.root):
                sys.exit(f"refusing: --backup-dir {args.backup_dir} is inside the tree being processed.")
            os.makedirs(args.backup_dir, exist_ok=True)
            if not os.access(args.backup_dir, os.W_OK):
                sys.exit(f"refusing: --backup-dir {args.backup_dir} is not writable.")
        elif not args.no_backup:
            sys.exit("refusing: --in-place needs --backup-dir DIR (recommended) or --no-backup.")

    corpus_files = len(targets)
    corpus_bytes = sum(os.path.getsize(p) for p in targets)
    sampled = False
    if args.limit and args.limit < len(targets):
        if args.first:
            targets = targets[:args.limit]
        else:
            targets = sorted(random.Random(args.seed).sample(targets, args.limit))
        sampled = True

    counts = {}
    prior_in = prior_out = 0
    rep = w = None
    if args.report:
        if args.resume and os.path.exists(args.report):
            done = set()
            with open(args.report, newline="") as fh:
                for r in csv.DictReader(fh):
                    done.add(r["path"])
                    counts[r["status"]] = counts.get(r["status"], 0) + 1
                    prior_in += int(r["bytes_in"])
                    prior_out += int(r["bytes_out"])
            print(f"resuming: {len(done)} files already recorded")
            targets = [p for p in targets if os.path.relpath(p, args.root) not in done]
            rep = open(args.report, "a", newline="")
            w = csv.writer(rep)
        else:
            rep = open(args.report, "w", newline="")
            w = csv.writer(rep)
            w.writerow(["path", "bytes_in", "bytes_out", "status", "note"])

    total_in = sum(os.path.getsize(p) for p in targets)
    if args.in_place and args.backup_dir and not args.dry_run:
        free = shutil.disk_usage(args.backup_dir).free
        need = int(total_in * 1.05)
        if free < need:
            sys.exit(f"refusing: not enough free space for backups "
                     f"({need/2**30:.1f} GB needed, {free/2**30:.1f} GB free).")
    if sampled:
        print(f"corpus : {corpus_files} PDFs, {corpus_bytes/2**20:.1f} MB")
        print(f"sample : {len(targets)} PDFs, {total_in/2**20:.1f} MB")
    else:
        print(f"{len(targets)} PDF(s), {total_in/2**20:.1f} MB")

    saved = prior_in - prior_out
    t0 = time.time()
    n_total = len(targets)

    def record(res):
        p, a, b, status, note = res
        counts[status] = counts.get(status, 0) + 1
        if w:
            w.writerow([os.path.relpath(p, args.root), a, b, status, note])
            rep.flush()
        if n_total <= 20 or status == "FAIL":
            print(f"  {status:12s} {os.path.relpath(p, args.root)}  "
                  f"{a:,} -> {b:,} bytes  {note}")
        return a - b

    def progress(i):
        if n_total > 20 and (i % 25 == 0 or i == n_total):
            el = time.time() - t0
            rate = i / el if el else 0
            eta = (n_total - i) / rate if rate else 0
            print(f"[{i}/{n_total}] saved {saved/2**20:.1f} MB  "
                  f"{rate*60:.0f} files/min  ETA {eta/3600:.1f}h", flush=True)

    if args.jobs > 1 and n_total:
        from concurrent.futures import ProcessPoolExecutor, as_completed
        with ProcessPoolExecutor(max_workers=args.jobs) as ex:
            futs = {ex.submit(process, p, args): p for p in targets}
            for i, fut in enumerate(as_completed(futs), 1):
                try:
                    saved += record(fut.result())
                except Exception as e:
                    p = futs[fut]
                    sz = os.path.getsize(p) if os.path.exists(p) else 0
                    saved += record((p, sz, sz, "SKIP", f"worker error: {type(e).__name__}"))
                progress(i)
    else:
        for i, p in enumerate(targets, 1):
            saved += record(process(p, args))
            progress(i)
    if rep:
        rep.close()

    pct = 100 * saved / total_in if total_in else 0
    print(f"\n{'DRY RUN, nothing written. ' if args.dry_run else ''}"
          f"saved {saved/2**20:.1f} MB of {total_in/2**20:.1f} MB ({pct:.1f}%)")
    for k, v in sorted(counts.items()):
        print(f"  {k:14s} {v}")
    if sampled and total_in:
        proj = corpus_bytes * (saved / total_in)
        print(f"projected over all {corpus_files} files: ~{proj/2**20:.0f} MB ({pct:.1f}%)")
    return 1 if counts.get("FAIL") else 0


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