#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
使用 OnnxOCR 对图片做 OCR，结果输出为 JSON 到 stdout。
用法1: python ocr-onnx.py --image <图片路径> [--project-root <项目根目录>]
       输出: {"success": true, "text": "识别结果"} 或 {"success": false, "error": "..."}
用法2: python ocr-onnx.py --image <图片路径> --find-text "要查找的文字" [--project-root <项目根目录>]
       在图中查找该文字，返回中心点: {"success": true, "x": 123, "y": 456} 或 {"success": false, "error": "..."}
"""
import sys
import os
import json
import argparse

def box_center(box):
    """box 为 4 个点 [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]，返回中心 (cx, cy)"""
    if not box or len(box) < 4:
        return None
    xs = [p[0] for p in box]
    ys = [p[1] for p in box]
    return (sum(xs) / len(xs), sum(ys) / len(ys))


def normalize_for_match(s):
    """规范化后用于匹配：去空格、全角括号/数字转半角，便于 OCR 变体（如 下一步(1)、下一步（1））都能匹配 下一步"""
    if not s:
        return ''
    s = (s or '').strip().replace(' ', '').replace('\u3000', '')
    # 全角括号、数字 → 半角
    t = []
    for c in s:
        if c in ('（', '［', '｛'): t.append('(')
        elif c in ('）', '］', '｝'): t.append(')')
        elif '\uff10' <= c <= '\uff19': t.append(chr(ord(c) - 0xfee0))
        else: t.append(c)
    return ''.join(t)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--image', required=True, help='图片路径（绝对或相对）')
    ap.add_argument('--find-text', default=None, help='要查找的文字；若指定则返回该文字在图中的中心点 x,y')
    ap.add_argument('--project-root', default=None, help='项目根目录，用于解析相对路径及加载 onnxocr')
    args = ap.parse_args()

    project_root = args.project_root
    if not project_root:
        project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
    project_root = os.path.normpath(project_root)

    # 确保可导入 onnxocr（包在 python/onnxocr/onnxocr/，故将 python/onnxocr 加入 path）
    onnxocr_root = os.path.join(project_root, 'python', 'onnxocr')
    if onnxocr_root not in sys.path:
        sys.path.insert(0, onnxocr_root)

    image_path = args.image
    if not os.path.isabs(image_path):
        image_path = os.path.normpath(os.path.join(project_root, image_path))

    if not os.path.isfile(image_path):
        out = {'success': False, 'error': f'图片不存在: {image_path}'}
        print(json.dumps(out, ensure_ascii=False))
        return

    try:
        import cv2
    except ImportError:
        out = {'success': False, 'error': 'OpenCV 未安装，请安装: pip install opencv-python'}
        print(json.dumps(out, ensure_ascii=False))
        sys.exit(1)

    try:
        from onnxocr.onnx_paddleocr import ONNXPaddleOcr
    except ImportError as e:
        out = {'success': False, 'error': f'OnnxOCR 导入失败: {e}。请确保 python/onnxocr 已存在且可导入。'}
        print(json.dumps(out, ensure_ascii=False))
        sys.exit(1)

    try:
        img = cv2.imread(image_path)
        if img is None:
            out = {'success': False, 'error': f'无法读取图片: {image_path}'}
            print(json.dumps(out, ensure_ascii=False))
            sys.exit(1)

        model = ONNXPaddleOcr(
            use_angle_cls=True,
            use_gpu=False,
            det_db_thresh=0.2,
            det_db_box_thresh=0.45,
            drop_score=0.3,
        )
        result = model.ocr(img)

        find_text = (args.find_text or '').strip()
        find_norm = normalize_for_match(find_text)

        if find_text:
            # 模式2：查找文字，返回中心点（支持包含匹配与规范化匹配，如 "下一步" 可匹配 "下一步(1)"、"下一步（1）"）
            if not result or not result[0]:
                out = {'success': False, 'error': f'图中未识别到文字，或未找到: "{find_text}"'}
                print(json.dumps(out, ensure_ascii=False))
                sys.exit(1)
            for line in result[0]:
                box, rec = line[0], line[1] if len(line) > 1 else None
                if not rec:
                    continue
                text = (rec[0] or '').strip()
                text_norm = normalize_for_match(text)
                if (find_text in text or text == find_text or
                        (find_norm and (find_norm in text_norm or text_norm == find_norm))):
                    center = box_center(box)
                    if center is not None:
                        cx, cy = int(round(center[0])), int(round(center[1]))
                        out = {'success': True, 'x': cx, 'y': cy}
                        print(json.dumps(out, ensure_ascii=False))
                        return
            out = {'success': False, 'error': f'图中未找到文字: "{find_text}"'}
            print(json.dumps(out, ensure_ascii=False))
            sys.exit(1)
        else:
            # 模式1：整图 OCR，返回全文
            if not result or not result[0]:
                text = ''
            else:
                lines = [line[1][0] for line in result[0] if line[1]]
                text = '\n'.join(lines) if lines else ''
            out = {'success': True, 'text': text}
            print(json.dumps(out, ensure_ascii=False))
    except Exception as e:
        out = {'success': False, 'error': str(e)}
        print(json.dumps(out, ensure_ascii=False))
        sys.exit(1)

if __name__ == '__main__':
    main()