#!/usr/bin/env python3
"""Bina 0.4 Rizeh: Persian page OCR (Surya layout + 95k Indra CTC student + printed/table/numeric experts and assembly).

usage: python rizeh.py page.jpg [more.png doc.pdf ...] [--out results] [--device cuda|cpu]
"""
import argparse,json,os,sys,time,importlib.util
from pathlib import Path
ROOT=Path(__file__).resolve().parent;ST=ROOT/'stallion';MARK=ROOT/'.bina-root'
LAYOUT_REPO,LAYOUT_REV='datalab-to/surya_layout2','0aee81d5fd9275c0582e545bf3a56944b1e75679'
TEXT=('.py','.json','.yml','.yaml','.txt','.jsonl')

def setup():
 """Point every bundled path at this checkout (idempotent; reruns after a move) and fetch the pinned layout model."""
 here=str(ROOT);old=MARK.read_text().strip()if MARK.exists()else'__BINA_ROOT__'
 if old!=here:
  n=0
  for p in ST.rglob('*'):
   if p.is_file()and p.suffix in TEXT:
    s=p.read_text(encoding='utf-8')
    if old in s:p.write_text(s.replace(old,here),encoding='utf-8');n+=1
  MARK.write_text(here);print(f'[rizeh] bound {n} files to {here}',file=sys.stderr)
 snap=ST/'indra-scale/hf/hub/models--datalab-to--surya_layout2/snapshots'/LAYOUT_REV
 if not(snap/'rfdetr_layout.pth').exists():
  from huggingface_hub import snapshot_download
  print('[rizeh] downloading layout model',LAYOUT_REPO,LAYOUT_REV,file=sys.stderr)
  snapshot_download(LAYOUT_REPO,revision=LAYOUT_REV,cache_dir=str(ST/'indra-scale/hf/hub'))
 return snap

def pages(path):
 p=Path(path)
 if p.suffix.lower()=='.pdf':
  import pypdfium2 as pdfium
  doc=pdfium.PdfDocument(str(p))
  for i in range(len(doc)):yield f'{p.stem}-p{i+1:03d}',doc[i].render(scale=200/72).to_pil().convert('RGB')
 else:
  from PIL import Image
  yield p.stem,Image.open(p).convert('RGB')

class Rizeh:
 def __init__(self,device='cuda'):
  snap=setup()
  os.environ.setdefault('GATE_DEVICE',device);os.environ.setdefault('FAST_DETECTOR_DEVICE','cpu');os.environ.setdefault('DETECTOR_IMAGE_CHUNK_HEIGHT','1400');os.environ['FAST_ORDER_MODEL_CHECKPOINT']=str(snap/'order')
  P=ST/'indra-scale';sys.path[:0]=[str(P),str(ST/'bina04-pagepath-v3-round3-clean-20261031'),str(P/'scoringTT')]
  spec=importlib.util.spec_from_file_location('bina04_rizeh_runtime',ST/'bina04-ink-loss-guard-95k-20261104/production.py');m=importlib.util.module_from_spec(spec);spec.loader.exec_module(m)
  self.engine=m.Runtime(device)
  from surya.fast_layout.server import LayoutEngine
  self.layout=LayoutEngine(checkpoint=str(snap))
 def __call__(self,image,name='page'):
  raw=self.layout.run_batch([image],[dict(threshold=None,use_order=True)])[0].model_dump(mode='json')
  out=self.engine(image,dict(raw_layout=raw),'page',name)
  lines=[dict(text=l.get('text',''),layer=l.get('layer'),box=[round(float(v),1)for v in l['box']],confidence=l.get('confidence'))for l in out.get('lines',[])if l.get('box')]
  return dict(text=out['text'],lines=lines)
 def close(self):self.engine.close()

def main():
 ap=argparse.ArgumentParser(description=__doc__,formatter_class=argparse.RawDescriptionHelpFormatter);ap.add_argument('inputs',nargs='+');ap.add_argument('--out',default='rizeh-out');ap.add_argument('--device',default='cuda',choices=['cuda','cpu']);a=ap.parse_args()
 O=Path(a.out);O.mkdir(parents=True,exist_ok=True);r=Rizeh(a.device)
 for src in a.inputs:
  for name,im in pages(src):
   t=time.perf_counter();res=r(im,name);res.update(source=str(src),seconds=round(time.perf_counter()-t,2),model='Bina-0.4-Rizeh')
   (O/f'{name}.json').write_text(json.dumps(res,ensure_ascii=False,indent=1),encoding='utf-8');(O/f'{name}.txt').write_text(res['text'],encoding='utf-8')
   print(f'{name}: {len(res["lines"])} lines, {res["seconds"]}s -> {O/name}.txt',flush=True)
 r.close()
if __name__=='__main__':main()