# qs_ocr_run.py - batch RapidOCR over extracted frames (PC, Python 3.8).
# Usage: python qs_ocr_run.py --frames-dir DIR --out OUT.jsonl --workers 3 [--step 1] [--from F] [--to F]
# Input frames: f_%06d.png (or .jpg) with global frame numbers in the filename.
# Output: JSONL, one line per processed frame: {"f": int, "texts": [{"t": str, "s": float, "b": [[x,y]x4]}]}
import argparse, glob, json, os, re, sys, time
from concurrent.futures import ProcessPoolExecutor

ENGINE = None
BOX_THRESH = None

def _init(box_thresh):
    global ENGINE, BOX_THRESH
    BOX_THRESH = box_thresh
    import cv2  # noqa: F401  (ensure cv2 loads in worker)
    from rapidocr_onnxruntime import RapidOCR
    if box_thresh and box_thresh > 0:
        ENGINE = RapidOCR(det_db_box_thresh=box_thresh)
    else:
        ENGINE = RapidOCR()

def _work(path):
    import cv2
    m = re.search(r'f_(\d+)\.', os.path.basename(path))
    fno = int(m.group(1)) if m else -1
    img = cv2.imread(path)
    if img is None:
        return {"f": fno, "texts": [], "err": "imread_fail"}
    try:
        result, _elapse = ENGINE(img)
    except Exception as e:  # keep batch alive
        return {"f": fno, "texts": [], "err": str(e)[:200]}
    texts = []
    if result:
        for box, txt, score in result:
            texts.append({
                "t": txt,
                "s": round(float(score), 3),
                "b": [[int(round(p[0])), int(round(p[1]))] for p in box],
            })
    return {"f": fno, "texts": texts}

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--frames-dir', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--workers', type=int, default=3)
    ap.add_argument('--step', type=int, default=1)
    ap.add_argument('--from', dest='frm', type=int, default=-1)
    ap.add_argument('--to', dest='to', type=int, default=-1)
    ap.add_argument('--box-thresh', type=float, default=0.0)
    args = ap.parse_args()

    paths = sorted(glob.glob(os.path.join(args.frames_dir, 'f_*.png')))
    if not paths:
        paths = sorted(glob.glob(os.path.join(args.frames_dir, 'f_*.jpg')))
    sel = []
    for p in paths:
        m = re.search(r'f_(\d+)\.', os.path.basename(p))
        if not m:
            continue
        fno = int(m.group(1))
        if args.frm >= 0 and fno < args.frm:
            continue
        if args.to >= 0 and fno > args.to:
            continue
        sel.append(p)
    if args.step > 1:
        sel = sel[::args.step]
    print("OCR START files=%d workers=%d step=%d box=%s" % (len(sel), args.workers, args.step, args.box_thresh), flush=True)
    t0 = time.time()
    done = 0
    with open(args.out, 'w', encoding='utf-8') as fh:
        with ProcessPoolExecutor(max_workers=args.workers, initializer=_init, initargs=(args.box_thresh,)) as ex:
            for rec in ex.map(_work, sel, chunksize=8):
                fh.write(json.dumps(rec, ensure_ascii=False) + "\n")
                done += 1
                if done % 200 == 0:
                    fh.flush()
                    print("OCR PROGRESS %d/%d elapsed=%.0fs" % (done, len(sel), time.time() - t0), flush=True)
    print("OCR DONE files=%d elapsed=%.0fs" % (done, time.time() - t0), flush=True)

if __name__ == '__main__':
    main()