#!/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
import traceback
from datetime import datetime

# ── 崩溃日志文件：写入源码目录根（easyocr_crash.log），统一管理、便于父进程读取 ──
def _get_virtual_memory_info() -> dict:
    """获取虚拟内存（页面文件）状态，用于诊断 WinError 1455。"""
    try:
        import ctypes
        class MEMORYSTATUSEX(ctypes.Structure):
            _fields_ = [
                ("dwLength", ctypes.c_ulong),
                ("dwMemoryLoad", ctypes.c_ulong),
                ("ullTotalPhys", ctypes.c_ulonglong),
                ("ullAvailPhys", ctypes.c_ulonglong),
                ("ullTotalPageFile", ctypes.c_ulonglong),
                ("ullAvailPageFile", ctypes.c_ulonglong),
                ("ullTotalVirtual", ctypes.c_ulonglong),
                ("ullAvailVirtual", ctypes.c_ulonglong),
                ("ullAvailExtendedVirtual", ctypes.c_ulonglong),
            ]
        mem = MEMORYSTATUSEX()
        mem.dwLength = ctypes.sizeof(MEMORYSTATUSEX)
        ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(mem))
        return {
            "total_phys_gb": round(mem.ullTotalPhys / (1024**3), 1),
            "avail_phys_gb": round(mem.ullAvailPhys / (1024**3), 1),
            "total_pagefile_gb": round(mem.ullTotalPageFile / (1024**3), 1),
            "avail_pagefile_gb": round(mem.ullAvailPageFile / (1024**3), 1),
            "commit_pct": round((1 - mem.ullAvailPageFile / max(mem.ullTotalPageFile, 1)) * 100, 1),
        }
    except Exception:
        return {}


def _get_crash_log_path() -> str:
    try:
        base = os.path.dirname(os.path.abspath(__file__))
    except Exception:
        base = "."
    return os.path.join(base, "easyocr_crash.log")

CRASH_LOG = _get_crash_log_path()

# ── 状态文件：记录 worker 当前在处理什么/卡多久，父进程 kill 前读取，避免"一声不吭" ──
def _get_state_path() -> str:
    try:
        base = os.path.dirname(os.path.abspath(__file__))
    except Exception:
        base = "."
    return os.path.join(base, "easyocr_worker_state.json")


STATE_PATH = _get_state_path()

# worker 当前运行模式（GPU/CPU），由 _init_reader 设置
_WORKER_MODE = "未知"


def _write_state(status: str, last_img_len: int = 0, ts: float = None, error: str = ""):
    """把当前 worker 状态落盘为 JSON，便于父进程 kill 前诊断死前状态。

    status: idle / busy / crashed / terminated
    """
    try:
        _ts = ts if ts is not None else datetime.now().timestamp()
        _rec = {
            "pid": os.getpid(),
            "mode": _WORKER_MODE,
            "status": status,
            "last_img_len": last_img_len,
            "ts": _ts,
            "iso": datetime.fromtimestamp(_ts).isoformat() if _ts else "",
            "error": error,
        }
        with open(STATE_PATH, "w", encoding="utf-8") as f:
            json.dump(_rec, f, ensure_ascii=False)
            f.flush()
    except Exception:
        pass


def _write_crash_log(msg: str):
    try:
        with open(CRASH_LOG, "a", encoding="utf-8") as f:
            f.write(msg + "\n")
            f.flush()
    except Exception:
        pass

# ── 顶层异常捕获：确保崩溃原因写入 stderr + 落盘（否则父进程看到空 stderr） ──
def _flush_exit(msg, code=1):
    full = f"[easyocr_worker] FATAL: {msg}"
    sys.stderr.write(full + "\n")
    sys.stderr.flush()
    _write_crash_log(f"{datetime.now()} {full}")
    _write_state("crashed", error=full[:500])
    sys.exit(code)

# ── 全局未捕获异常钩子：任何漏网的 Python 异常都落盘 ──
def _global_excepthook(exc_type, exc_val, exc_tb):
    msg = "".join(traceback.format_exception(exc_type, exc_val, exc_tb))
    sys.stderr.write(f"[easyocr_worker] 未捕获异常:\n{msg}")
    sys.stderr.flush()
    _write_crash_log(f"{datetime.now()} [easyocr_worker] 未捕获异常:\n{msg}")
    _write_state("crashed", error=msg[:500])

sys.excepthook = _global_excepthook

# ── faulthandler：捕获原生层崩溃（segfault / CUDA 初始化竞态），Python try/except 抓不到 ──
#    必须在重型 import 之前启用，连 import 阶段的原生崩溃也能 dump 出来。
try:
    import faulthandler
    _fh = open(CRASH_LOG, "a", encoding="utf-8")
    faulthandler.enable(file=_fh, all_threads=True)
except Exception:
    pass

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

    # ── CUDA 驱动预检：在 import torch 之前探测 CUDA 驱动是否可用 ──
    #    防止 CUDA 驱动损坏导致 import torch 时原生崩溃（crash log 为空的静默崩溃）
    if os.environ.get("EASYOCR_FORCE_CPU") != "1":
        try:
            import ctypes
            ctypes.windll.nvcuda  # 尝试加载 CUDA 驱动 DLL
            sys.stderr.write("[easyocr_worker] CUDA 驱动 DLL 加载成功\n")
            sys.stderr.flush()
        except Exception as _cuda_drv_err:
            os.environ["EASYOCR_FORCE_CPU"] = "1"
            os.environ["CUDA_VISIBLE_DEVICES"] = ""
            sys.stderr.write(f"[easyocr_worker] ⚠️ CUDA 驱动加载失败({_cuda_drv_err})，强制 CPU 模式\n")
            sys.stderr.flush()
            _write_crash_log(f"{datetime.now()} [easyocr_worker] CUDA 驱动预检失败: {_cuda_drv_err}")

    # ── 预检：物理内存+虚拟内存是否足够加载 CUDA 库 ──
    _vm = _get_virtual_memory_info()
    # ⭐ 统一使用 GlobalMemoryStatusEx（与 monitor.py / Task Manager 一致）
    #    psutil.virtual_memory().available 在 Windows 上可比 Task Manager 少报 ~1.7GB，
    #    导致 4.2GB 可用时误判为 2.5GB 触发降级。
    if _vm:
        try:
            _avail_phys_gb = _vm.get("avail_phys_gb", 0)
            _total_phys_gb = _vm.get("total_phys_gb", 0)
            _used_phys_gb = _total_phys_gb - _avail_phys_gb
            _pagefile_pct = _vm.get("commit_pct", 0)
            _pagefile_used_gb = _vm.get("total_pagefile_gb", 0) - _vm.get("avail_pagefile_gb", 0)
            _commit_gb = _used_phys_gb + _pagefile_used_gb
            _commit_total_gb = _total_phys_gb + _vm.get("total_pagefile_gb", 0)
            _commit_pct = (_commit_gb / max(_commit_total_gb, 1)) * 100
            # ⭐ 放宽阈值：提交量>95%才强制CPU，物理内存<2GB才强制CPU
            #   避免内存充足时误判（用户报告提交量81.8%被误判为不足）
            if _commit_pct > 95.0:
                _msg = (
                    f"系统内存压力过高: 物理剩余{_avail_phys_gb:.1f}GB，"
                    f"提交量{_commit_pct:.1f}% ({_commit_gb:.1f}GB/{_commit_total_gb:.1f}GB)"
                )
                _write_crash_log(f"{datetime.now()} [easyocr_worker] 系统内存压力预检失败: {_msg}")
                os.environ["EASYOCR_FORCE_CPU"] = "1"
                os.environ["CUDA_VISIBLE_DEVICES"] = ""  # 禁用CUDA，防止加载CUDA库导致崩溃
                sys.stderr.write(f"[easyocr_worker] ⚠️ {_msg}，降级CPU模式\n")
                sys.stderr.flush()
            elif _avail_phys_gb < 2.0:
                _msg = (
                    f"物理内存不足: 剩余{_avail_phys_gb:.1f}GB"
                    f"(总计{_total_phys_gb:.1f}GB)"
                    f"，CUDA加载需要≥2GB物理内存"
                )
                _write_crash_log(f"{datetime.now()} [easyocr_worker] 物理内存预检失败: {_msg}")
                os.environ["EASYOCR_FORCE_CPU"] = "1"
                os.environ["CUDA_VISIBLE_DEVICES"] = ""  # 禁用CUDA，防止加载CUDA库导致崩溃
                sys.stderr.write(f"[easyocr_worker] ⚠️ {_msg}，降级CPU模式\n")
                sys.stderr.flush()
        except Exception:
            pass  # GlobalMemoryStatusEx 失败时不阻断
    # ⭐ 放宽退出阈值：百分比>98% 或 剩余<1GB 时退出（防冻结）
    #   避免内存充足时误退出（用户报告提交量81.8%被误判为不足）
    if _vm and (_vm.get("commit_pct", 0) > 98.0 or _vm.get("avail_pagefile_gb", 99) < 1.0):
        _msg = (
            f"页面文件不足: 已用{_vm.get('commit_pct', '?')}%"
            f"(剩余{_vm.get('avail_pagefile_gb', '?')}GB/{_vm.get('total_pagefile_gb', '?')}GB)"
            f"，CUDA库需要充足的虚拟地址空间"
        )
        _write_crash_log(f"{datetime.now()} [easyocr_worker] 页面文件预检失败: {_msg}")
        _flush_exit(f"页面文件不足: {_msg}", code=2)

    from PIL import Image
    import numpy as np
    import easyocr
except Exception as e:
    _vm = _get_virtual_memory_info()
    _vm_hint = ""
    if _vm:
        _vm_hint = (
            f"\n[virtual memory] 物理内存: {_vm.get('avail_phys_gb',0)}/{_vm.get('total_phys_gb',0)} GB"
            f" | 页面文件: {_vm.get('avail_pagefile_gb',0)}/{_vm.get('total_pagefile_gb',0)} GB"
            f" | 提交使用率: {_vm.get('commit_pct',0)}%"
        )
    # WinError 1455 = 页面文件太小 → 标记为永久失败，不重试
    _is_pagefile_error = "1455" in str(e) or "页面文件" in str(e)
    _exit_code = 2 if _is_pagefile_error else 1  # exit(2) = 虚拟内存不足
    _write_crash_log(
        f"{datetime.now()} [easyocr_worker] import {_exit_code}: "
        f"{'页面文件不足' if _is_pagefile_error else 'import失败'}: {e}{_vm_hint}"
    )
    _flush_exit(f"import 失败{'(页面文件不足)' if _is_pagefile_error else ''}: {e}{_vm_hint}", code=_exit_code)


def _init_reader():
    """初始化 EasyOCR reader（GPU 优先，父进程可强制 CPU）。"""
    global _WORKER_MODE
    # 先检查环境变量是否强制CPU
    _force_cpu = os.environ.get("EASYOCR_FORCE_CPU") == "1"

    # 再次检查物理内存：加载torch+CUDA需要大量内存，内存不足时强制CPU
    if not _force_cpu:
        try:
            import psutil
            _mem = psutil.virtual_memory()
            _avail_phys_gb = _mem.available / (1024**3)
            if _avail_phys_gb < 3.0:
                _force_cpu = True
                os.environ["CUDA_VISIBLE_DEVICES"] = ""  # 禁用CUDA，防止加载CUDA库导致崩溃
                sys.stderr.write(
                    f"[easyocr_worker] ⚠️ 物理内存不足(剩余{_avail_phys_gb:.1f}GB)，"
                    f"强制CPU模式（避免CUDA加载导致冻结）\n"
                )
                sys.stderr.flush()
        except Exception:
            pass  # psutil不可用时不阻断
    
    try:
        import torch
        # 父进程 EASYOCR_FORCE_CPU=1 → 强制 CPU（显存反复 OOM 时降级，防冻结）
        use_gpu = (not _force_cpu) and torch.cuda.is_available()
        if use_gpu:
            # CUDA 健康检查：尝试创建一个小张量，避免驱动/显存问题导致后续崩溃
            try:
                _test_tensor = torch.zeros(1, device="cuda")
                del _test_tensor
                torch.cuda.empty_cache()
                sys.stderr.write(f"[easyocr_worker] CUDA 健康检查通过 (设备: {torch.cuda.get_device_name(0)})\n")
                sys.stderr.flush()
            except Exception as cuda_err:
                sys.stderr.write(f"[easyocr_worker] CUDA 健康检查失败({cuda_err})，降级 CPU\n")
                sys.stderr.flush()
                use_gpu = False
                os.environ["CUDA_VISIBLE_DEVICES"] = ""  # 禁用CUDA，防止后续崩溃
    except Exception:
        use_gpu = False

    try:
        reader = easyocr.Reader(["ch_sim", "en"], gpu=use_gpu)
        mode = "GPU" if use_gpu else "CPU"
        _WORKER_MODE = mode
        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)
        _WORKER_MODE = "CPU"
        sys.stderr.write("[easyocr_worker] Reader 就绪（CPU 模式）\n")
        sys.stderr.flush()
        return reader


def _process_image(img_b64: str, reader, timeout: float = 30.0) -> str:
    """解码 base64 图片 → EasyOCR 识别 → 返回空格拼接文本。
    timeout: 单张图片OCR超时秒数，防止GPU/CPU推理卡死。
    """
    if reader is None:
        return ""
    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)

    # ⭐ 超时保护：用 concurrent.futures 包装 readtext，防止GPU推理卡死
    from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(reader.readtext, img_np)
        try:
            results = future.result(timeout=timeout)
        except FuturesTimeoutError:
            sys.stderr.write(f"[easyocr_worker] readtext 超时({timeout}s)，放弃本次OCR\n")
            sys.stderr.flush()
            _write_crash_log(f"{datetime.now()} [easyocr_worker] readtext 超时({timeout}s)")
            return ""
        except Exception as e:
            sys.stderr.write(f"[easyocr_worker] readtext 异常: {e}\n")
            sys.stderr.flush()
            raise

    return "\n".join([r[1] for r in results]).strip()


def _main():
    _write_crash_log(f"{datetime.now()} [easyocr_worker] 启动 (pid={os.getpid()})")
    reader = None
    try:
        reader = _init_reader()
    except Exception as e:
        sys.stderr.write(f"[easyocr_worker] Reader 初始化失败: {e}\n")
        sys.stderr.flush()
        _write_crash_log(f"{datetime.now()} [easyocr_worker] Reader 初始化失败: {e}")
    # ⭐ 就绪握手：即使 Reader 失败也要发 ready，否则父进程永远等握手
    sys.stdout.write(json.dumps({"ready": True}, ensure_ascii=False) + "\n")
    sys.stdout.flush()
    sys.stderr.write("[easyocr_worker] 就绪，等待 stdin 输入\n")
    sys.stderr.flush()
    _write_state("idle")

    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 字段")
            _write_state("busy", last_img_len=len(img_b64))
            sys.stderr.write(f"[easyocr_worker] 收到请求，处理中... (img长度={len(img_b64)})\n")
            sys.stderr.flush()
            text = _process_image(img_b64, reader)
            sys.stderr.write(f"[easyocr_worker] 处理完成，结果长度={len(text)}\n")
            sys.stderr.flush()
            out = json.dumps({"ok": True, "text": text}, ensure_ascii=False)
            sys.stdout.write(out + "\n")
            sys.stdout.flush()
            _write_state("idle")
        except Exception as e:
            estr = str(e)
            # ⭐ CUDA OOM / kernel error：GPU context 已损坏，继续跑只会反复卡死占显存
            #   → 立即退出进程释放显存，父进程检测到 stdout 关闭后自动重启/降级 CPU
            if "cuda" in estr.lower() or "out of memory" in estr.lower():
                _write_crash_log(f"{datetime.now()} [easyocr_worker] CUDA 错误，主动退出释放显存: {estr[:200]}")
                _flush_exit(f"CUDA 错误，主动退出释放显存: {estr[:200]}")
            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()


def _cleanup_gpu():
    """退出时清理 CUDA/GPU 内存，防止子进程被 kill 后 GPU 内存泄漏到系统。"""
    try:
        import gc
        gc.collect()
    except Exception:
        pass
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    except Exception:
        pass


def _signal_handler(signum, frame):
    """捕获 SIGTERM/SIGINT，先记录被终止状态再清理 GPU 退出（proc.terminate() 触发此处理）。"""
    _write_state("terminated", error=f"收到信号 {signum}，父进程请求终止，正在清理退出")
    _cleanup_gpu()
    sys.exit(0)


# 注册清理钩子：atexit（正常退出）+ signal（SIGTERM/SIGINT）
import atexit
atexit.register(_cleanup_gpu)
try:
    import signal
    signal.signal(signal.SIGTERM, _signal_handler)
    signal.signal(signal.SIGINT, _signal_handler)
except Exception:
    pass


if __name__ == "__main__":
    try:
        _main()
    except Exception as e:
        _flush_exit(f"_main 崩溃: {e}\n{traceback.format_exc()}")
    finally:
        _cleanup_gpu()