#!/usr/bin/env python3
"""
Surya OCR (motor neuronal, distinto a tesseract) SOLO de las paginas que tienen
candidatos a escalar. Surya no comete los mismos errores sistematicos que
tesseract (h->il, e->c...), asi que sirve de arbitro fiable.

Flujo:
  1. lee out/*.candidates.jsonl, toma paginas de los candidatos REVISAR/REPO_TYPO/AMBIGUO
  2. renderiza esas paginas a PNG (300dpi) en un staging por libro
  3. corre surya_ocr sobre la carpeta (una invocacion -> carga el modelo una vez)
  4. guarda surya/<book>/<stem>/pNNNN.surya.txt

Debe ejecutarse dentro del venv: surya-venv/bin/python surya_pages.py
Uso: python surya_pages.py [--book BOOKID] [--only REVISAR,REPO_TYPO,AMBIGUO]
"""
import json, sys, argparse, subprocess, tempfile, shutil, os, re
from pathlib import Path

HERE = Path(__file__).resolve().parent
CFG = json.loads((HERE / "books.json").read_text())
SCAN_ROOT = Path(CFG["_scan_root"])
OUT = HERE / "out"
SURYA_DIR = HERE / "surya"; SURYA_DIR.mkdir(exist_ok=True)

# stem (nombre sin extension del PDF) -> ruta PDF
STEM2PDF = {}
for book, rels in CFG["books"].items():
    for rel in rels:
        p = SCAN_ROOT / rel
        STEM2PDF.setdefault(book, {})[p.stem] = p

def pages_needed(book, classes):
    """devuelve {stem: set(page_int)} de candidatos a escalar."""
    cf = OUT / f"{book}.candidates.jsonl"
    if not cf.exists():
        return {}
    need = {}
    for l in cf.open():
        c = json.loads(l)
        if c["verdict"] not in classes:
            continue
        pr = c.get("page", "")
        if "#" not in pr:
            continue
        stem, pg = pr.split("#", 1)
        m = re.search(r'(\d+)', pg)
        if not m:
            continue
        need.setdefault(stem, set()).add(int(m.group(1)))
    return need

# --- Surya: import API una vez ---
_PREDICTORS = {}
def get_predictors():
    if _PREDICTORS:
        return _PREDICTORS
    from surya.recognition import RecognitionPredictor
    from surya.detection import DetectionPredictor
    _PREDICTORS["det"] = DetectionPredictor()
    _PREDICTORS["rec"] = RecognitionPredictor()
    return _PREDICTORS

def surya_text(img_path):
    import re as _re
    from PIL import Image
    P = get_predictors()
    img = Image.open(img_path).convert("RGB")
    # surya 0.20: PageOCRResult.blocks -> List[BlockOCRResult]; each block.html has text
    preds = P["rec"]([img], full_page=True)
    page = preds[0]
    texts = []
    for blk in page.blocks:
        if blk.skipped or blk.error:
            continue
        # strip HTML tags from block.html
        txt = _re.sub(r'<[^>]+>', ' ', blk.html or '').strip()
        if txt:
            texts.append(txt)
    return "\n".join(texts)

def render_page(pdf, page, outdir):
    prefix = str(outdir / f"r{page:04d}")
    subprocess.run(["pdftoppm", "-f", str(page), "-l", str(page),
                    "-r", "300", "-png", str(pdf), prefix],
                   capture_output=True, timeout=120)
    imgs = sorted(outdir.glob(f"r{page:04d}*.png"))
    return imgs[0] if imgs else None

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--book")
    ap.add_argument("--only", default="REVISAR,REPO_TYPO,AMBIGUO")
    a = ap.parse_args()
    classes = set(a.only.split(","))
    books = [a.book] if a.book else list(CFG["books"].keys())

    total = 0
    for book in books:
        need = pages_needed(book, classes)
        if not need:
            continue
        for stem, pgs in need.items():
            pdf = STEM2PDF.get(book, {}).get(stem)
            if not pdf or not pdf.exists():
                print(f"  WARN sin PDF para {book}/{stem}", file=sys.stderr); continue
            destdir = SURYA_DIR / book / stem
            destdir.mkdir(parents=True, exist_ok=True)
            with tempfile.TemporaryDirectory() as td:
                tdp = Path(td)
                for pg in sorted(pgs):
                    outf = destdir / f"p{pg:04d}.surya.txt"
                    if outf.exists():
                        continue
                    img = render_page(pdf, pg, tdp)
                    if not img:
                        outf.write_text(""); continue
                    try:
                        outf.write_text(surya_text(img))
                        total += 1
                    except Exception as e:
                        print(f"  ERR surya {book}/{stem} p{pg}: {e}", file=sys.stderr)
                    img.unlink(missing_ok=True)
            print(f"  {book}/{stem}: {len(pgs)} paginas surya", file=sys.stderr, flush=True)
    print(f"FIN surya: {total} paginas nuevas", file=sys.stderr)

if __name__ == "__main__":
    main()