import os
import json
from pathlib import Path

# PaddleOCR を import する前に設定する必要があります。
os.environ["FLAGS_use_mkldnn"] = "0"
os.environ["FLAGS_use_onednn"] = "0"
os.environ["PADDLE_PDX_ENABLE_MKLDNN_BYDEFAULT"] = "0"
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"

from paddleocr import PaddleOCR


def to_py(obj):
    """numpy array / PaddleOCR の結果を、JSON 保存しやすい通常の Python オブジェクトに変換します。"""
    if hasattr(obj, "tolist"):
        return obj.tolist()
    return obj


def main():
    img_path = "./samples/jp_001.jpg"
    output_dir = Path("./output")
    output_dir.mkdir(parents=True, exist_ok=True)

    ocr = PaddleOCR(
        text_detection_model_name="PP-OCRv6_small_det",
        text_recognition_model_name="PP-OCRv6_small_rec",

        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,

        text_det_limit_side_len=736,
        text_det_limit_type="max",
        text_recognition_batch_size=1,

        device="cpu",
    )

    results = ocr.predict(img_path)

    all_items = []

    for page_idx, res in enumerate(results):
        # PaddleOCR v3.x の OCRResult オブジェクトは、結果オブジェクトそのものです。
        # そのため、res["res"] のように下位キーを参照する必要はありません。
        data = res

        texts = to_py(data["rec_texts"])
        scores = to_py(data["rec_scores"])
        boxes = to_py(data["rec_boxes"])

        print(f"\nImage: {img_path}")
        print(f"Detected text lines: {len(texts)}")
        print("-" * 80)

        for i, text in enumerate(texts):
            score = float(scores[i])
            box = boxes[i]

            print(f"{i + 1:02d}. {text}\t score={score:.4f}")

            all_items.append({
                "index": i + 1,
                "text": text,
                "score": score,
                "box": box,
            })

        # PaddleOCR 標準の可視化画像と JSON も保存します。
        res.save_to_img(str(output_dir))
        res.save_to_json(str(output_dir))

    with open(output_dir / "texts_only.json", "w", encoding="utf-8") as f:
        json.dump(all_items, f, ensure_ascii=False, indent=2)

    print("\nSaved:")
    print(f"  {output_dir}/texts_only.json")
    print(f"  {output_dir}/PaddleOCR visualization/json")


if __name__ == "__main__":
    main()
