#!/usr/bin/env python3
"""
Doble OCR tesseract de los escaneos autoritativos (la fuente de verdad).

Cada pagina se renderiza a 300 dpi y se OCR-ea DOS veces con configuraciones
distintas de tesseract:
  - psm 6  (bloque uniforme de texto)
  - psm 4  (columna de texto de tamanos variables)
ambas con --oem 1 (LSTM).

Donde las dos lecturas coinciden -> consenso fiable.
Donde discrepan -> 'dudoso' (se escalara a Surya + Llama mas adelante).

Reanudable: si los dos .txt de una pagina existen, se salta.
Paralelo: ProcessPoolExecutor.

Uso:
  python3 ocr_pages.py [--workers N] [--book BOOKID] [--dpi 300]
Salida:
  ocr/<book>/<pdfstem>/pNNNN.psm6.txt
  ocr/<book>/<pdfstem>/pNNNN.psm4.txt
"""
import json, os, re, subprocess, sys, tempfile, argparse, time
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor, as_completed

HERE = Path(__file__).resolve().parent
CFG = json.loads((HERE / "books.json").read_text())
SCAN_ROOT = Path(CFG["_scan_root"])
OCR_DIR = HERE / "ocr"

def page_count(pdf: Path) -> int:
    r = subprocess.run(["pdfinfo", str(pdf)], capture_output=True, text=True)
    m = re.search(r'^Pages:\s+(\d+)', r.stdout, re.M)
    return int(m.group(1)) if m else 0

def ocr_one(args):
    pdf_str, page, outdir_str, dpi = args
    pdf = Path(pdf_str)
    outdir = Path(outdir_str)
    o6 = outdir / f"p{page:04d}.psm6.txt"
    o4 = outdir / f"p{page:04d}.psm4.txt"
    if o6.exists() and o4.exists():
        return (pdf.name, page, "skip")
    with tempfile.TemporaryDirectory() as td:
        prefix = os.path.join(td, "pg")
        subprocess.run(["pdftoppm", "-f", str(page), "-l", str(page),
                        "-r", str(dpi), "-png", pdf_str, prefix],
                       capture_output=True, timeout=120)
        imgs = sorted(Path(td).glob("pg*.png"))
        if not imgs:
            o6.write_text(""); o4.write_text("")
            return (pdf.name, page, "noimg")
        img = str(imgs[0])
        for psm, out in ((6, o6), (4, o4)):
            r = subprocess.run(
                ["tesseract", img, "-", "--oem", "1", "--psm", str(psm), "-l", "eng"],
                capture_output=True, text=True, timeout=120)
            out.write_text(r.stdout)
    return (pdf.name, page, "ok")

def build_jobs(only_book, dpi):
    jobs = []
    books = CFG["books"]
    # ordenar por total de paginas asc para que los libros cortos terminen antes
    order = []
    for book, pdfs in books.items():
        if only_book and book != only_book:
            continue
        total = 0
        for rel in pdfs:
            p = SCAN_ROOT / rel
            if p.exists():
                total += page_count(p)
        order.append((total, book, pdfs))
    order.sort()
    for total, book, pdfs in order:
        for rel in pdfs:
            pdf = SCAN_ROOT / rel
            if not pdf.exists():
                print(f"  WARN missing {pdf}", file=sys.stderr); continue
            stem = pdf.stem
            outdir = OCR_DIR / book / stem
            outdir.mkdir(parents=True, exist_ok=True)
            n = page_count(pdf)
            for pg in range(1, n + 1):
                jobs.append((str(pdf), pg, str(outdir), dpi))
    return jobs, order

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=max(2, (os.cpu_count() or 4) - 2))
    ap.add_argument("--book", default=None)
    ap.add_argument("--dpi", type=int, default=300)
    a = ap.parse_args()

    jobs, order = build_jobs(a.book, a.dpi)
    print(f"Libros: {len(order)} | paginas-trabajo: {len(jobs)} | workers: {a.workers}",
          file=sys.stderr, flush=True)
    for total, book, _ in order:
        print(f"   {total:5d}p  {book}", file=sys.stderr)

    done = skip = 0
    t0 = time.time()
    with ProcessPoolExecutor(max_workers=a.workers) as ex:
        futs = [ex.submit(ocr_one, j) for j in jobs]
        for i, f in enumerate(as_completed(futs), 1):
            name, pg, st = f.result()
            if st == "skip": skip += 1
            else: done += 1
            if i % 100 == 0:
                el = time.time() - t0
                rate = done / el if el else 0
                print(f"  {i}/{len(jobs)}  ok={done} skip={skip}  "
                      f"{rate:.1f}pg/s  ult={name} p{pg}", file=sys.stderr, flush=True)
    print(f"FIN: {done} OCR nuevas, {skip} saltadas, {time.time()-t0:.0f}s",
          file=sys.stderr)

if __name__ == "__main__":
    main()