"""CPU entry point for page-by-page PaddleOCR-VL PDF recognition."""

from __future__ import annotations

import argparse
import os
import platform
import sys
import time
from pathlib import Path

from pdf_ocr_core import preflight_pdf, process_pdf

PROJECT_ROOT = Path(__file__).resolve().parent
DEFAULT_OUTPUT = PROJECT_ROOT / "outputs"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="PaddleOCR-VL-1.6 CPU PDF OCR（逐页、可恢复）",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("pdf", type=Path, help="输入 PDF 文件")
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT, help="输出根目录")
    parser.add_argument("--pages", help="一页或多个页码范围，例如 1-5,8,10-")
    parser.add_argument("--dpi", type=int, default=144, help="PDF 页面渲染 DPI")
    parser.add_argument("--password", help="加密 PDF 密码")
    parser.add_argument("--threads", type=int, default=None, help="Paddle CPU 线程数")
    parser.add_argument("--resume", action="store_true", help="跳过已完成页，继续现有任务")
    parser.add_argument("--overwrite", action="store_true", help="删除已有输出并重新处理")
    parser.add_argument("--keep-rendered", action="store_true", help="保留逐页渲染 PNG")
    parser.add_argument("--fail-fast", action="store_true", help="任一页失败后立即停止")
    parser.add_argument("--max-new-tokens", type=int, default=None, help="限制每个文本块最大生成 token")
    parser.add_argument("--min-pixels", type=int, default=None, help="VLM 最小输入像素参数")
    parser.add_argument("--max-pixels", type=int, default=None, help="VLM 最大输入像素参数")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    total_cores = os.cpu_count() or 4
    safe_default_threads = max(1, total_cores - 2)
    threads = args.threads or int(os.environ.get("PADDLE_THREADS", safe_default_threads))
    if threads < 1:
        print("[ERROR] --threads 必须大于等于 1", file=sys.stderr)
        return 2

    try:
        preflight = preflight_pdf(
            pdf_path=args.pdf,
            output_root=args.output,
            pages=args.pages,
            dpi=args.dpi,
            password=args.password,
            resume=args.resume,
            overwrite=args.overwrite,
        )
    except Exception as exc:
        print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1

    print(
        f"[PDF] {preflight['page_count']} pages, selected: "
        f"{len(preflight['selected_pages'])}, output: {preflight['document_dir']}"
    )

    from paddle import core
    from paddleocr import PaddleOCRVL

    core.set_num_threads(threads)
    print(f"[CPU] Threads: {threads} / {total_cores} (reserved: {max(0, total_cores - threads)})")
    print("Loading PaddleOCR-VL-1.6...")
    init_started = time.perf_counter()
    pipeline = PaddleOCRVL(pipeline_version="v1.6", device="cpu")
    init_seconds = time.perf_counter() - init_started
    print(f"Model init: {init_seconds:.2f}s")

    predict_kwargs = {
        key: value
        for key, value in {
            "max_new_tokens": args.max_new_tokens,
            "min_pixels": args.min_pixels,
            "max_pixels": args.max_pixels,
        }.items()
        if value is not None
    }
    metadata = {
        "device": "cpu",
        "cpu_threads": threads,
        "python_version": platform.python_version(),
        "platform": platform.platform(),
        "model_init_seconds": round(init_seconds, 3),
        "pipeline_version": "v1.6",
    }

    try:
        summary = process_pdf(
            pipeline=pipeline,
            pdf_path=args.pdf,
            output_root=args.output,
            pages=args.pages,
            dpi=args.dpi,
            password=args.password,
            resume=args.resume,
            overwrite=args.overwrite,
            keep_rendered=args.keep_rendered,
            fail_fast=args.fail_fast,
            run_metadata=metadata,
            predict_kwargs=predict_kwargs,
        )
    except KeyboardInterrupt:
        print("\n[INTERRUPTED] 已保存当前进度；使用 --resume 继续。", file=sys.stderr)
        return 130
    except Exception as exc:
        print(f"[ERROR] {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1

    print("\n[PDF OCR Summary]")
    print(f"Status:    {summary['status']}")
    print(f"Completed: {summary['completed_pages']} / {summary['selected_pages']}")
    print(f"Failed:    {summary['failed_pages']}")
    print(f"Output:    {summary['document_dir']}")
    return 0 if not summary["failed_pages"] else 3


if __name__ == "__main__":
    raise SystemExit(main())