# -*- coding: utf-8 -*-
"""
向 OCR 服务发一次请求做 OCR，验证端点可用、看识别结果。

用法:
  python test_ocr.py <图片路径>
  python test_ocr.py <图片路径> --url http://其它IP:18081
  python test_ocr.py <图片路径> --raw          # 额外打印原始 JSON 响应

不给图片时，自动在 ./composite、./split 下找一张 .jpg 试。
"""
import argparse
import base64
import glob
import json
import os
import random
import time

import requests

DEFAULT_URL = "http://asr.ocr.general.smart.l20.production.polaris:18081"
OCR_PATH = "/youtu/ocrapi/hpgeneralocr_v3"


def find_sample_image(given):
    if given:
        return given if os.path.exists(given) else None
    for root in ("./composite", "./split"):
        hits = glob.glob(os.path.join(root, "**", "*.jpg"), recursive=True)
        if hits:
            return hits[0]
    return None


def img2base64(img_path):
    with open(img_path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def ocr_request(base_url, image_b64, timeout=30):
    """向端点发一次 OCR 请求，返回 (原始响应文本, 耗时秒)。"""
    url = base_url.rstrip("/") + OCR_PATH
    app_id = f"{int(time.time() * 1000)}{random.randint(100000, 999999)}"
    payload = {
        "app_id": app_id,
        "image": image_b64,
        "ocr_template": "OCR",
        "options": {"pdf_page_index": 0, "detect_text": True},
    }
    t0 = time.time()
    resp = requests.post(url, headers={"Content-Type": "application/json"},
                         json=payload, timeout=timeout)
    dt = time.time() - t0
    resp.raise_for_status()
    return resp.text, dt


def extract_texts(raw):
    """从 Youtu OCR 响应里抽出每行识别文字 (content, conf)。"""
    try:
        data = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        return None
    if data.get("error_code", 0) != 0:
        return f"OCR 返回错误: error_code={data.get('error_code')}, msg={data.get('error_msg')}"
    lines = []
    for rec in data.get("recognize_list", []):
        for item in rec.get("item_content", {}).get("item_list", []):
            lines.append((item.get("content", ""), item.get("conf", 0)))
    return lines


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image", nargs="?", default=None, help="图片路径")
    ap.add_argument("--url", default=DEFAULT_URL, help="OCR 端点 (默认 Polaris 地址)")
    ap.add_argument("--raw", action="store_true", help="打印原始 JSON 响应")
    args = ap.parse_args()

    img = find_sample_image(args.image)
    if not img:
        print("找不到图片，请把一张 jpg 路径作为参数传入。")
        return

    print(f"端点: {args.url}{OCR_PATH}")
    print(f"图片: {img}")
    print("-" * 60)

    try:
        raw, dt = ocr_request(args.url, img2base64(img))
    except requests.exceptions.RequestException as e:
        print(f"✗ 请求失败: {e}")
        return

    print(f"✓ 请求成功，耗时 {dt * 1000:.0f}ms")

    if args.raw:
        print("\n原始响应:")
        print(raw)
        print("-" * 60)

    result = extract_texts(raw)
    if result is None:
        print("响应不是合法 JSON")
    elif isinstance(result, str):
        print(result)
    elif not result:
        print("（未识别到文字）")
    else:
        print(f"\n识别到 {len(result)} 行文字:")
        for content, conf in result:
            print(f"  [{conf:.2f}] {content}")


if __name__ == "__main__":
    main()