"""
批量 OCR 识别 — 多进程并行加速

原理：每个进程独立加载一份 pipeline 实例，同时处理不同图片。
适用场景：一次处理多张图片（如文件夹批量 OCR）。

用法:
    python batch_ocr.py <图片目录> [--workers 4] [--threads 10]
"""
import time
import os
import sys
import argparse
from multiprocessing import Pool, cpu_count
from pathlib import Path


def ocr_single(image_path: str, threads: int) -> dict:
    """单张图片 OCR（在子进程中执行）"""
    from paddle import core
    core.set_num_threads(threads)

    from paddleocr import PaddleOCRVL

    pipeline = PaddleOCRVL(pipeline_version="v1.6")

    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():
    parser = argparse.ArgumentParser(description="Batch OCR with multiprocessing")
    parser.add_argument("dir", type=str, help="图片目录")
    parser.add_argument("--workers", type=int, default=4, help="并行进程数 (默认 4)")
    parser.add_argument("--threads", type=int, default=None, help="每进程线程数 (默认: 总核心/workers)")
    args = parser.parse_args()

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

    images = list(image_dir.glob("*.png")) + list(image_dir.glob("*.jpg")) + list(image_dir.glob("*.jpeg"))
    if not images:
        print(f"[ERROR] 目录中没有图片: {args.dir}")
        sys.exit(1)

    workers = min(args.workers, len(images))
    threads = args.threads or max(1, cpu_count() // workers)

    print(f"图片数量: {len(images)}")
    print(f"并行进程: {workers}")
    print(f"每进程线程: {threads}")
    print(f"总核心利用: {workers * threads} / {cpu_count()}")
    print("=" * 60)

    t0 = time.perf_counter()

    # 每个子进程独立加载模型 + 推理
    with Pool(workers) as pool:
        tasks = [(str(img), threads) for img in images]
        results = pool.starmap(ocr_single, tasks)

    total_elapsed = time.perf_counter() - t0

    # 输出结果
    print("\n" + "=" * 60)
    for r in results:
        print(f"\n[文件] {r['path']} ({r['elapsed']:.1f}s)")
        for block in r["blocks"]:
            print(f"  [{block['label']}] {block['content'][:60]}{'...' if len(block['content']) > 60 else ''}")

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


if __name__ == "__main__":
    main()