#!/usr/bin/env python3
"""Re-OCR an image-based (scanned) book into an improved searchable PDF.

For each page:
  - if the page is image-based (a large raster + a poor OCR text layer),
    render it to grayscale, run tesseract 5 (psm 6, eng) to produce a fresh
    searchable single-page PDF (image + invisible text layer), and use that;
  - if the page is already digital text (no big image), copy it unchanged.

The result is a same-look PDF whose extractable text (get_text) is the fresh
tesseract-5 OCR — which fixes the garbled transliteration in the old layer.
compare.py then reads the improved text with zero code changes.

Usage: python3 reocr_book.py <input.pdf> <output.pdf> [dpi] [first] [last]
"""
import fitz, subprocess, sys, os, tempfile

DPI = 300

def is_image_page(page) -> bool:
    A = page.rect.width * page.rect.height
    if not A:
        return False
    for b in page.get_text("rawdict")["blocks"]:
        if b.get("type") == 1:  # image block
            bb = b["bbox"]
            ia = (bb[2] - bb[0]) * (bb[3] - bb[1])
            if ia > 0.45 * A:
                return True
    return False

def main():
    inp, outp = sys.argv[1], sys.argv[2]
    dpi = int(sys.argv[3]) if len(sys.argv) > 3 else DPI
    first = int(sys.argv[4]) if len(sys.argv) > 4 else 0
    last = int(sys.argv[5]) if len(sys.argv) > 5 else None

    src = fitz.open(inp)
    out = fitz.open()
    n = src.page_count
    last = n if last is None else min(last, n)
    work = tempfile.mkdtemp(prefix="reocr_")
    reocr = passthru = 0
    for i in range(first, last):
        page = src[i]
        if is_image_page(page):
            pm = page.get_pixmap(dpi=dpi, colorspace=fitz.csGRAY)
            jpg = os.path.join(work, "p.jpg")
            with open(jpg, "wb") as fh:
                fh.write(pm.tobytes("jpeg", jpg_quality=80))
            base = os.path.join(work, "p")
            r = subprocess.run(["tesseract", jpg, base, "--psm", "6", "pdf"],
                               capture_output=True)
            pdfp = base + ".pdf"
            if os.path.exists(pdfp):
                with fitz.open(pdfp) as pg:
                    out.insert_pdf(pg)
                reocr += 1
            else:
                out.insert_pdf(src, from_page=i, to_page=i)
                passthru += 1
        else:
            out.insert_pdf(src, from_page=i, to_page=i)
            passthru += 1
    out.save(outp, deflate=True, garbage=3)
    out.close(); src.close()
    print(f"{os.path.basename(inp)}: pages {first}-{last}  re-OCR={reocr}  passthru={passthru}  -> {outp} ({os.path.getsize(outp)/1e6:.1f} MB)")

if __name__ == "__main__":
    main()