inkreach-prefect/flows/label_ocr.py

"""面单 OCR enrichment:对所有系统(RIIN+SDS)有待处理面单的订单,用 v2 batch API 批量识别。

对应 label-ocr。每小时跑一次,跟上增量同步节奏。
按国家分组,每组 200 单/批调 /api/v2/batch(服务端并发下载+重试)。
"""

from __future__ import annotations

import os

import httpx
from prefect import flow, task
from prefect.runtime import deployment, flow_run

from api.ocr import OcrClient, country_to_ocr_code
from flows.shared import SyncDB, build_stats_rows, emit_sync_summary

BATCH_SIZE = 200

_OCR_FIELD_MAP = [
    ("layout", "ocrLayout"), ("courier", "ocrCourier"), ("tracking", "ocrTracking"),
    ("marker", "ocrMarker"), ("s_name", "ocrSName"), ("s_street", "ocrSStreet"),
    ("s_city", "ocrSCity"), ("s_state", "ocrSState"), ("s_zip", "ocrSZip"),
    ("r_name", "ocrRName"), ("r_street", "ocrRStreet"), ("r_city", "ocrRCity"),
    ("r_state", "ocrRState"), ("r_zip", "ocrRZip"),
]


def _result_to_fields(r) -> dict[str, str | None]:
    if not r.ok:
        return {"ocrLayout": "ERROR"}
    return {col: (getattr(r, attr) or "") for attr, col in _OCR_FIELD_MAP}


def _has_running_peer() -> bool:
    """检查是否有同类 flow run 正在运行(排除自身)。"""
    api_url = os.environ.get("PREFECT_API_URL", "http://localhost:4200/api")
    cur_id = str(flow_run.id) if flow_run.id else ""
    dep_id = str(deployment.id) if deployment.id else ""
    if not dep_id:
        return False
    try:
        resp = httpx.post(
            f"{api_url}/flow_runs/filter",
            json={
                "flow_runs": {"deployment_id": {"any_": [dep_id]}},
                "state": {"type": {"any_": ["RUNNING"]}},
            },
            params={"limit": 10},
            timeout=10,
        )
        resp.raise_for_status()
        return any(r["id"] != cur_id for r in resp.json())
    except Exception:
        return False


@task
def process_all(pending: list) -> dict[str, int]:
    db = SyncDB()
    ocr = OcrClient()
    stats = {"success": 0, "error": 0, "unsupported": 0, "batches": 0}
    try:
        # 1. 按国家分组;不支持的国家直接标 UNSUPPORTED
        by_country: dict[str, list[tuple[str, object, str]]] = {}
        unsupp: list[tuple[str, object]] = []
        for order_no, created_at, pdf_url, dest_country in pending:
            code = country_to_ocr_code(dest_country)
            if code:
                by_country.setdefault(code, []).append((order_no, created_at, pdf_url))
            else:
                unsupp.append((order_no, created_at))

        if unsupp:
            db.update_ocr_batch([(no, ca, {"ocrLayout": "UNSUPPORTED"}) for no, ca in unsupp])
            stats["unsupported"] = len(unsupp)
            print(f"[OCR] 不支持国家: {len(unsupp)} 单 → UNSUPPORTED")

        # 2. 每个国家分批调 v2 batch API
        for country_code, orders in by_country.items():
            total = len(orders)
            for i in range(0, total, BATCH_SIZE):
                batch = orders[i : i + BATCH_SIZE]
                created_map = {no: ca for no, ca, _ in batch}
                urls = [{"url": url, "filename": no} for no, _, url in batch]
                batch_no = i // BATCH_SIZE + 1
                total_batches = -(-total // BATCH_SIZE)
                try:
                    results = ocr.recognize_batch(urls, country_code)
                except Exception as err:  # noqa: BLE001
                    print(f"[OCR] {country_code} batch {batch_no}/{total_batches} 请求失败: {str(err)[:100]}")
                    db.update_ocr_batch([(no, created_map[no], {"ocrLayout": "ERROR"}) for no in created_map])
                    stats["error"] += len(batch)
                    continue

                updates = []
                for r in results:
                    order_no = r.file
                    created_at = created_map.get(order_no)
                    if created_at is None:
                        continue
                    fields = _result_to_fields(r)
                    updates.append((order_no, created_at, fields))
                    stats["success" if r.ok else "error"] += 1
                db.update_ocr_batch(updates)
                stats["batches"] += 1
                ok = sum(1 for r in results if r.ok)
                print(f"[OCR] {country_code} batch {batch_no}/{total_batches}: "
                      f"{ok}/{len(results)} 成功")
        return stats
    finally:
        ocr.close()


@flow(name="label-ocr", log_prints=True)
def label_ocr(
    destination_countries: list[str] | None = None,
    courier_source: str | None = None,
    days: int = 7,
) -> dict:
    """面单 OCR enrichment。默认处理英国(GB)+美国(US),近 7 天下单的订单。

    手动跑其他国家: label_ocr(destination_countries=['DE'])
    """
    if destination_countries is None:
        destination_countries = ['GB', 'US']
    if _has_running_peer():
        print("已有 label-ocr 在运行,跳过本次")
        return {"skipped": True}
    db = SyncDB()
    before = db.count_ocr_status()
    print(f"OCR 状态(开始): {before}")

    pending = db.get_ocr_pending(
        destination_countries=destination_countries,
        courier_source=courier_source,
        days=days,
    )
    scope = f" [{'/'.join(destination_countries)}{f' {courier_source}' if courier_source else ''}]"
    print(f"待 OCR: {len(pending)} 单{scope} (RIIN+SDS, 近{days}天)")
    if not pending:
        emit_sync_summary(f"面单 OCR{scope}", build_stats_rows("ALL", {}, [
            ("待处理", 0), ("累计已完成", before.get("done", 0)),
        ]))
        return {"processed": 0, **before}

    stats = process_all(pending)
    after = db.count_ocr_status()
    total_processed = stats["success"] + stats["error"] + stats["unsupported"]
    print(f"完成{scope}: 处理 {total_processed} 单 / {stats['batches']} 批 | "
          f"成功 {stats['success']} 失败 {stats['error']} 不支持 {stats['unsupported']}")
    print(f"累计: {after}")
    emit_sync_summary(f"面单 OCR{scope}", build_stats_rows("ALL", {}, [
        ("本批处理", total_processed), ("成功", stats["success"]),
        ("失败", stats["error"]), ("不支持", stats["unsupported"]),
        ("累计已完成", after.get("done", 0)), ("累计待处理", after.get("pending", 0)),
    ]))
    return {"processed": total_processed, **stats, **after}


if __name__ == "__main__":
    label_ocr.serve(name="label-ocr", cron="0 * * * *")