#!/usr/bin/env python3
"""EasyOCR 子进程 worker。

通过 stdin 读 base64 图片、stdout 写 JSON 结果。GPU crash 只杀此进程，不冻主进程。
协议：
  stdin:  {"img":"<base64_png_bytes>"}
  stdout: {"ok":true,"text":"识别结果"}
  stdout: {"ok":false,"error":"错误描述"}

初始化慢（首次 3-10s），初始化完成后常驻处理。
"""

import json
import sys
import base64
import io
import os

# ── 线程/进程限制：避免 MKL/OpenMP 在子进程多线程冲突 ──
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")

from PIL import Image
import numpy as np
import easyocr


def _init_reader():
    """初始化 EasyOCR reader（GPU 优先）。"""
    try:
        import torch
        use_gpu = torch.cuda.is_available()
    except Exception:
        use_gpu = False

    try:
        reader = easyocr.Reader(["ch_sim", "en"], gpu=use_gpu)
        mode = "GPU" if use_gpu else "CPU"
        sys.stderr.write(f"[easyocr_worker] Reader 就绪（{mode}模式）\n")
        sys.stderr.flush()
        return reader
    except Exception as e:
        # GPU 不可用 → CPU 兜底
        sys.stderr.write(f"[easyocr_worker] GPU 初始化失败({e})，尝试 CPU\n")
        sys.stderr.flush()
        reader = easyocr.Reader(["ch_sim", "en"], gpu=False)
        sys.stderr.write("[easyocr_worker] Reader 就绪（CPU 模式）\n")
        sys.stderr.flush()
        return reader


def _process_image(img_b64: str, reader) -> str:
    """解码 base64 图片 → EasyOCR 识别 → 返回空格拼接文本。"""
    raw = base64.b64decode(img_b64)
    img = Image.open(io.BytesIO(raw))
    # 保持与 monitor.py 一致：RGBA → numpy
    if img.mode != "RGBA":
        img = img.convert("RGBA")
    img_np = np.array(img)
    results = reader.readtext(img_np)
    return "\n".join([r[1] for r in results]).strip()


def _main():
    reader = _init_reader()
    # ⭐ 就绪握手：初始化完成后先输出 {"ready": true}，让父进程知道可以开始请求
    # 否则父进程的 10s 请求超时会包含 3-10s 的初始化时间 → 永远超时
    sys.stdout.write(json.dumps({"ready": True}, ensure_ascii=False) + "\n")
    sys.stdout.flush()
    sys.stderr.write("[easyocr_worker] 就绪，等待 stdin 输入\n")
    sys.stderr.flush()

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            payload = json.loads(line)
            img_b64 = payload.get("img", "")
            if not img_b64:
                raise ValueError("缺少 img 字段")
            text = _process_image(img_b64, reader)
            out = json.dumps({"ok": True, "text": text}, ensure_ascii=False)
            sys.stdout.write(out + "\n")
            sys.stdout.flush()
        except Exception as e:
            out = json.dumps({"ok": False, "error": str(e)}, ensure_ascii=False)
            sys.stdout.write(out + "\n")
            sys.stdout.flush()
            sys.stderr.write(f"[easyocr_worker] 处理失败: {e}\n")
            sys.stderr.flush()


if __name__ == "__main__":
    _main()