"""
批量 OCR 识别 — 多进程并行加速（系统友好版）

修复要点:
  1. 进程错峰启动（随机延迟），避免同时加载 N 个模型导致内存/CUP 打满
  2. 降低子进程优先级，保证系统 UI 正常响应
  3. 预留 1-2 个核心给 OS，避免 CPU 完全饱和
  4. 用 imap_unordered 逐任务分发，而非一次性灌满

用法:
    python batch_ocr.py <图片目录> [--workers 4] [--threads 5]

安全建议:
    - 32GB RAM 建议 --workers <= 4
    - 16GB RAM 建议 --workers <= 2
    - 不确定时先用 --workers 1 测试
"""
import time
import os
import sys
import random
import argparse
from multiprocessing import Pool, cpu_count
from pathlib import Path

# ── Worker 初始化（在子进程中执行） ──

def _init_worker(threads: int, stagger_max: float):
    """
    每个 Worker 启动时：随机延迟 → 设线程数 → 降优先级 → 加载模型。

    随机延迟是关键：避免 N 个进程同时读磁盘/分配内存，
    将 4×2GB=8GB 的内存峰值分散到 0~15s 的时间窗口中。
    """
    delay = random.uniform(0, stagger_max)
    time.sleep(delay)

    # 算子级线程数
    from paddle import core
    core.set_num_threads(threads)

    # 降低进程优先级（不影响计算吞吐，但让 OS 调度更公平）
    try:
        import psutil
        p = psutil.Process()
        if sys.platform == "win32":
            p.nice(psutil.BELOW_NORMAL_PRIORITY_CLASS)
        else:
            p.nice(10)
    except ImportError:
        pass
    except Exception:
        pass

    # 加载 pipeline（~2GB，耗时 ~40s）
    from paddleocr import PaddleOCRVL
    global _pipeline
    _pipeline = PaddleOCRVL(pipeline_version="v1.6")


def _ocr_task(image_path: str) -> dict:
    """单张图片 OCR（使用全局 pipeline）"""
    global _pipeline
    t0 = time.perf_counter()
    result = _pipeline.predict(image_path)
    elapsed = time.perf_counter() - t0

    blocks = []
    for block in result[0]["parsing_res_list"]:
        if block.content.strip():
            blocks.append({
                "label": block.label,
                "bbox": block.bbox,
                "content": block.content,
            })

    return {
        "path": str(image_path),
        "elapsed": round(elapsed, 2),
        "blocks": blocks,
    }


# ── 主流程 ──

def main():
    total_cores = cpu_count()

    parser = argparse.ArgumentParser(
        description="批量 OCR — 多进程并行（系统友好版）",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python batch_ocr.py images/                          # 默认 2 进程
  python batch_ocr.py images/ --workers 4              # 4 进程（需 32GB RAM）
  python batch_ocr.py images/ --workers 2 --threads 8  # 指定每进程线程数
        """,
    )
    parser.add_argument("dir", type=str, help="图片目录")
    parser.add_argument(
        "--workers", type=int, default=2,
        help="并行进程数 (默认 2，安全值；最大建议不超过 RAM_GB/2)",
    )
    parser.add_argument(
        "--threads", type=int, default=None,
        help=f"每进程线程数 (默认: (总核心-1)/workers，保证 OS 有 1 核可用)",
    )
    parser.add_argument(
        "--stagger", type=float, default=15.0,
        help="进程启动错峰窗口秒数 (默认 15s，值越大内存峰值越低)",
    )
    args = parser.parse_args()

    # ── 扫描图片 ──
    image_dir = Path(args.dir)
    if not image_dir.is_dir():
        print(f"[ERROR] 目录不存在: {args.dir}")
        sys.exit(1)

    extensions = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tiff", "*.tif", "*.webp")
    images = []
    for ext in extensions:
        images.extend(image_dir.glob(ext))
    images = sorted(images)

    if not images:
        print(f"[ERROR] 目录中没有图片: {args.dir}")
        sys.exit(1)

    # ── 资源规划 ──
    workers = min(args.workers, len(images))
    reserved_for_os = 1  # 至少给 OS 留 1 个逻辑核心
    if args.threads:
        threads = args.threads
    else:
        threads = max(1, (total_cores - reserved_for_os) // workers)

    total_cpu_used = workers * threads
    stagger = args.stagger

    # 内存估算
    model_mem_per_worker = 2.0  # GB, 模型 ~1.8GB + 运行时开销
    estimated_mem = workers * model_mem_per_worker + 2  # +2GB for OS

    try:
        import psutil
        avail_gb = psutil.virtual_memory().available / (1024**3)
        mem_ok = avail_gb > estimated_mem
    except ImportError:
        avail_gb = None
        mem_ok = True  # 无法检测，假定 OK

    # ── 打印配置 ──
    print("=" * 60)
    print(f"  图片数量:       {len(images)}")
    print(f"  并行进程:       {workers}")
    print(f"  每进程线程:     {threads}")
    print(f"  CPU 占用:       {total_cpu_used} / {total_cores} 核 (保留 {total_cores - total_cpu_used} 给 OS)")
    print(f"  错峰窗口:       {stagger}s")
    print(f"  预估内存:       ~{estimated_mem:.0f}GB (可用: {avail_gb:.0f}GB)" if avail_gb else f"  预估内存:       ~{estimated_mem:.0f}GB")
    if not mem_ok:
        print(f"  [WARNING] 可用内存不足！建议降低 --workers 到 {max(1, int((avail_gb - 2) / model_mem_per_worker))}")
    print("=" * 60)

    if not mem_ok:
        resp = input("内存不足，是否继续？[y/N] ").strip().lower()
        if resp != "y":
            print("已取消。")
            sys.exit(0)

    # ── 执行 ──
    t0 = time.perf_counter()

    with Pool(
        processes=workers,
        initializer=_init_worker,
        initargs=(threads, stagger),
    ) as pool:
        # imap_unordered: 逐任务分发，先完成的先返回
        # chunk 大 → 吞吐高但内存峰值高；chunk=1 → 最平滑
        image_paths = [str(img) for img in images]
        results = list(pool.imap_unordered(_ocr_task, image_paths, chunksize=1))

    total_elapsed = time.perf_counter() - t0

    # ── 输出 ──
    print("\n" + "=" * 60)
    for r in sorted(results, key=lambda x: x["path"]):
        print(f"\n[文件] {r['path']} ({r['elapsed']:.1f}s)")
        for block in r["blocks"]:
            preview = block["content"].replace("\n", "\\n")
            if len(preview) > 80:
                preview = preview[:80] + "..."
            print(f"  [{block['label']}] {preview}")

    total_per_image = sum(r["elapsed"] for r in results)
    print("\n" + "=" * 60)
    print(f"  总图片:   {len(images)}")
    print(f"  总耗时:   {total_elapsed:.1f}s ({total_elapsed/60:.1f}min)")
    print(f"  平均每图: {total_elapsed / len(images):.1f}s")
    print(f"  串行预计: {total_per_image:.1f}s")
    if total_elapsed > 0:
        print(f"  加速比:   {total_per_image / total_elapsed:.2f}x")
    print("=" * 60)


if __name__ == "__main__":
    main()