import json
import os
import sys
from pathlib import Path

os.environ.setdefault(
    "PADDLE_PDX_CACHE_HOME",
    str(Path(__file__).resolve().parents[1] / ".cache" / "paddlex"),
)

from paddleocr import PaddleOCR


def main() -> int:
    if len(sys.argv) != 3:
        print("usage: ocr_batch.py INPUT_DIR OUTPUT_DIR")
        return 2

    input_dir = Path(sys.argv[1])
    output_dir = Path(sys.argv[2])
    output_dir.mkdir(parents=True, exist_ok=True)
    images = sorted(input_dir.glob("q-*.png"))

    engine = PaddleOCR(
        device="cpu",
        enable_mkldnn=False,
        text_detection_model_name="PP-OCRv5_mobile_det",
        text_recognition_model_name="arabic_PP-OCRv5_mobile_rec",
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
    )

    for index, image in enumerate(images, start=1):
        predictions = list(engine.predict(input=str(image)))
        extracted = []
        for prediction in predictions:
            payload = prediction.json
            if callable(payload):
                payload = payload()
            extracted.append(payload)
        destination = output_dir / f"{image.stem}.json"
        destination.write_text(json.dumps(extracted, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"{index}/{len(images)} {image.name}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())