#!/usr/bin/env python3
"""
OCR completo con Surya de TODAS las páginas de todos los PDFs.
Salta páginas que ya tienen .surya.txt. Guarda progreso página a página.

Uso:
  surya-venv/bin/python surya_full.py                    # todos los libros
  surya-venv/bin/python surya_full.py --book srimad-bhagavatam
  surya-venv/bin/python surya_full.py --book sri-caitanya-caritamrta
"""
import argparse, json, os, re, subprocess, sys, tempfile
from pathlib import Path

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

# Forzar MPS en Apple Silicon
os.environ.setdefault("RECOGNITION_BATCH_SIZE", "64")   # más grande = más rápido en MPS
os.environ.setdefault("DETECTOR_BATCH_SIZE", "8")

# CRÍTICO para Caitanya-caritāmṛta (y todo libro con verso en bengalí):
# el modo full-page entra en BUCLE DE REPETICIÓN con el bengalí y genera hasta
# su tope de tokens (~10 min/página → timeout, página vacía). Se usa modo POR
# BLOQUES (layout → OCR por bloque): cada bloque tiene presupuesto de tokens
# acotado a su tamaño, así el bengalí no se dispara y el inglés/IAST sale bien
# (~50 s/página, estable). El layout guiado (JSON-schema por guided decoding)
# está pensado para vLLM y ROMPE con el backend llama.cpp ("400: failed to
# parse grammar") → hay que desactivarlo para que el layout funcione.
os.environ.setdefault("SURYA_GUIDED_LAYOUT", "false")

# stem → Path del PDF
STEM2PDF: dict[str, dict[str, Path]] = {}
for _book, _rels in CFG["books"].items():
    for _rel in _rels:
        _p = SCAN_ROOT / _rel
        STEM2PDF.setdefault(_book, {})[_p.stem] = _p

# ── Surya (carga modelo una sola vez) ────────────────────────────────────────
_PREDS: dict = {}

def get_preds():
    if _PREDS:
        return _PREDS
    import torch
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"[surya] usando device={device}", flush=True)
    from surya.recognition import RecognitionPredictor
    from surya.layout import LayoutPredictor
    _PREDS["lay"] = LayoutPredictor()
    _PREDS["rec"] = RecognitionPredictor()
    return _PREDS

def surya_text(img_path: Path) -> str:
    from PIL import Image
    P = get_preds()
    img = Image.open(img_path).convert("RGB")
    # Modo por bloques: layout detecta los bloques y cada uno se OCR-ea con un
    # presupuesto de tokens acotado a su tamaño (los ~12 bloques de la página se
    # procesan en paralelo por los slots del servidor). Evita el bucle de
    # repetición del full-page con el bengalí. Ver nota SURYA_GUIDED_LAYOUT.
    layout = P["lay"]([img])
    preds  = P["rec"]([img], layout_results=layout, full_page=False)
    page   = preds[0]
    texts = []
    for blk in page.blocks:
        if getattr(blk, 'skipped', False) or getattr(blk, 'error', False):
            continue
        txt = re.sub(r'<[^>]+>', ' ', blk.html or '').strip()
        if txt:
            texts.append(txt)
    return "\n".join(texts)

def render_page(pdf: Path, page: int, outdir: Path) -> Path | None:
    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 pdf_page_count(pdf: Path) -> int:
    try:
        out = subprocess.check_output(["pdfinfo", str(pdf)], text=True, stderr=subprocess.DEVNULL)
        for line in out.splitlines():
            if line.startswith("Pages:"):
                return int(line.split()[-1])
    except Exception:
        pass
    return 0

# ── Main ─────────────────────────────────────────────────────────────────────
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--book", help="procesar solo este libro")
    a  = ap.parse_args()

    books = [a.book] if a.book else list(CFG["books"].keys())

    grand_new   = 0
    grand_skip  = 0
    grand_total = 0

    for book in books:
        pdfs = STEM2PDF.get(book, {})
        if not pdfs:
            continue

        for stem, pdf in sorted(pdfs.items()):
            if not pdf.exists():
                print(f"  WARN PDF no existe: {pdf}", file=sys.stderr)
                continue

            n_pages = pdf_page_count(pdf)
            if n_pages == 0:
                print(f"  WARN no puedo contar páginas: {pdf.name}", file=sys.stderr)
                continue

            dest = SURYA_DIR / book / stem
            dest.mkdir(parents=True, exist_ok=True)

            new   = 0
            skip  = 0
            errs  = 0
            print(f"\n{book}/{stem}  ({n_pages} pgs)", flush=True)

            with tempfile.TemporaryDirectory() as td:
                tdp = Path(td)
                for pg in range(1, n_pages + 1):
                    outf = dest / f"p{pg:04d}.surya.txt"
                    if outf.exists():
                        skip += 1
                        continue
                    img = render_page(pdf, pg, tdp)
                    if not img:
                        outf.write_text("")
                        errs += 1
                        continue
                    try:
                        text = surya_text(img)
                        outf.write_text(text, encoding="utf-8")
                        new += 1
                    except Exception as e:
                        print(f"  ERR p{pg}: {e}", file=sys.stderr)
                        outf.write_text("")
                        errs += 1
                    finally:
                        img.unlink(missing_ok=True)

                    # Progreso en línea
                    if (pg % 50) == 0 or pg == n_pages:
                        print(f"  {pg}/{n_pages}  new={new} skip={skip} err={errs}", flush=True)

            grand_new   += new
            grand_skip  += skip
            grand_total += n_pages
            print(f"  ✓ {stem}: {new} nuevas, {skip} ya hechas, {errs} errores", flush=True)

    print(f"\n=== FIN === new={grand_new}  skip={grand_skip}  total={grand_total}")

if __name__ == "__main__":
    main()