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_sample.py INPUT_IMAGE OUTPUT_JSON")
        return 2

    source = Path(sys.argv[1]).resolve()
    destination = Path(sys.argv[2]).resolve()
    destination.parent.mkdir(parents=True, exist_ok=True)

    engine = PaddleOCR(
        lang="fa",
        device="cpu",
        enable_mkldnn=False,
        text_detection_model_name=os.environ.get("GOVAHINO_OCR_DET_MODEL", "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,
    )
    predictions = list(engine.predict(input=str(source)))
    extracted = []
    for prediction in predictions:
        payload = prediction.json
        if callable(payload):
            payload = payload()
        extracted.append(payload)

    destination.write_text(
        json.dumps(extracted, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(f"saved {destination}")
    return 0


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