"""PaddleOCR-VL-1.6 GPU 单图推理与 Benchmark。"""

import argparse
import json
import platform
import statistics
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any

GPU_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = GPU_DIR.parent
DEFAULT_IMAGE = PROJECT_ROOT / "images" / "手写01.png"
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "benchmarks" / "gpu"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="PaddleOCR-VL-1.6 GPU Benchmark")
    parser.add_argument("--image", type=Path, default=DEFAULT_IMAGE, help="待识别图片")
    parser.add_argument("--device-id", type=int, default=0, help="CUDA GPU 编号")
    parser.add_argument("--warmup", type=int, default=1, help="预热轮数")
    parser.add_argument("--rounds", type=int, default=3, help="正式测试轮数")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help="Benchmark JSON 输出目录",
    )
    parser.add_argument("--no-result", action="store_true", help="不在控制台输出 OCR 文本")
    return parser.parse_args()


def validate_args(args: argparse.Namespace) -> None:
    args.image = args.image.expanduser().resolve()
    args.output_dir = args.output_dir.expanduser().resolve()
    if not args.image.is_file():
        raise ValueError(f"图片不存在: {args.image}")
    if args.device_id < 0:
        raise ValueError("--device-id 不能小于 0")
    if args.warmup < 0:
        raise ValueError("--warmup 不能小于 0")
    if args.rounds < 1:
        raise ValueError("--rounds 必须大于等于 1")


def configure_cuda(device_id: int):
    try:
        import paddle
    except ImportError as exc:
        raise RuntimeError("未安装 GPU 子项目依赖，请先运行 gpu/setup_env.py。") from exc

    if not paddle.is_compiled_with_cuda():
        raise RuntimeError(
            "当前 PaddlePaddle 未编译 CUDA 支持。请确认安装的是 paddlepaddle-gpu，"
            "且正在使用 gpu/.venv。"
        )

    try:
        device_count = paddle.device.cuda.device_count()
    except Exception as exc:
        raise RuntimeError(f"无法查询 CUDA 设备: {exc}") from exc

    if device_count < 1:
        raise RuntimeError("未检测到 NVIDIA CUDA GPU；本脚本不会自动回退到 CPU。")
    if device_id >= device_count:
        raise RuntimeError(f"GPU {device_id} 不存在，当前仅检测到 {device_count} 个 CUDA 设备。")

    device = f"gpu:{device_id}"
    try:
        paddle.set_device(device)
        paddle.device.cuda.synchronize(device_id)
    except Exception as exc:
        raise RuntimeError(f"无法启用 {device}: {exc}") from exc

    try:
        device_name = paddle.device.cuda.get_device_name(device_id)
    except Exception:
        device_name = "unknown"

    return paddle, device, device_name


def synchronize(paddle: Any, device_id: int) -> None:
    paddle.device.cuda.synchronize(device_id)


def read_gpu_memory(paddle: Any, device_id: int) -> dict[str, float | None]:
    stats: dict[str, float | None] = {
        "allocated_mb": None,
        "reserved_mb": None,
        "max_allocated_mb": None,
        "max_reserved_mb": None,
    }
    functions = {
        "allocated_mb": "memory_allocated",
        "reserved_mb": "memory_reserved",
        "max_allocated_mb": "max_memory_allocated",
        "max_reserved_mb": "max_memory_reserved",
    }
    for key, function_name in functions.items():
        function = getattr(paddle.device.cuda, function_name, None)
        if function is None:
            continue
        try:
            stats[key] = round(float(function(device_id)) / (1024**2), 2)
        except Exception:
            pass
    return stats


def result_summary(result: list[Any]) -> dict[str, Any]:
    first = result[0]
    blocks = first["parsing_res_list"]
    return {
        "width": first["width"],
        "height": first["height"],
        "layout_boxes": len(first["layout_det_res"]["boxes"]),
        "parsed_blocks": len(blocks),
        "non_empty_blocks": sum(bool(block.content.strip()) for block in blocks),
    }


def print_ocr_result(result: list[Any]) -> None:
    print("\n[OCR Result]")
    for item in result:
        for block in item["parsing_res_list"]:
            if block.content.strip():
                print(f"[{block.label}] {block.bbox}")
                print(block.content)
                print()


def main() -> int:
    args = parse_args()
    try:
        validate_args(args)
        paddle, device, device_name = configure_cuda(args.device_id)
    except (ValueError, RuntimeError) as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        return 1

    from paddleocr import PaddleOCRVL

    print("=" * 70)
    print(f"Device:          {device} ({device_name})")
    print(f"PaddlePaddle:    {paddle.__version__}")
    print(f"Input image:     {args.image}")
    print(f"Warmup/Rounds:   {args.warmup}/{args.rounds}")
    print("=" * 70)

    synchronize(paddle, args.device_id)
    init_started = time.perf_counter()
    pipeline = PaddleOCRVL(pipeline_version="v1.6", device=device)
    synchronize(paddle, args.device_id)
    init_seconds = time.perf_counter() - init_started
    print(f"Model init:      {init_seconds:.3f}s")

    result = None
    for index in range(args.warmup):
        print(f"Warmup {index + 1}/{args.warmup}...", flush=True)
        result = pipeline.predict(str(args.image))
        synchronize(paddle, args.device_id)

    inference_times: list[float] = []
    for index in range(args.rounds):
        synchronize(paddle, args.device_id)
        started = time.perf_counter()
        result = pipeline.predict(str(args.image))
        synchronize(paddle, args.device_id)
        elapsed = time.perf_counter() - started
        inference_times.append(elapsed)
        print(f"Inference {index + 1}/{args.rounds}: {elapsed:.3f}s", flush=True)

    if result is None:
        print("[ERROR] 未产生推理结果。", file=sys.stderr)
        return 2

    summary = result_summary(result)
    benchmark = {
        "status": "completed",
        "timestamp": datetime.now().astimezone().isoformat(),
        "platform": platform.platform(),
        "python_version": platform.python_version(),
        "paddle_version": paddle.__version__,
        "pipeline_version": "v1.6",
        "device": device,
        "device_name": device_name,
        "image_path": str(args.image),
        "image": summary,
        "warmup_rounds": args.warmup,
        "benchmark_rounds": args.rounds,
        "model_init_seconds": round(init_seconds, 3),
        "inference_seconds": {
            "all": [round(value, 3) for value in inference_times],
            "min": round(min(inference_times), 3),
            "max": round(max(inference_times), 3),
            "mean": round(statistics.fmean(inference_times), 3),
            "median": round(statistics.median(inference_times), 3),
            "stdev": round(statistics.pstdev(inference_times), 3),
        },
        "gpu_memory": read_gpu_memory(paddle, args.device_id),
    }

    args.output_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    output_path = args.output_dir / f"gpu-benchmark-{timestamp}.json"
    output_path.write_text(json.dumps(benchmark, ensure_ascii=False, indent=2), encoding="utf-8")

    print("\n[Benchmark]")
    print(f"Image:           {summary['width']} x {summary['height']}")
    print(f"Layout boxes:    {summary['layout_boxes']}")
    print(f"Parsed blocks:   {summary['parsed_blocks']}")
    print(f"Average:         {benchmark['inference_seconds']['mean']:.3f}s")
    print(f"Min/Max:         {benchmark['inference_seconds']['min']:.3f}s / {benchmark['inference_seconds']['max']:.3f}s")
    print(f"Result JSON:     {output_path}")

    if not args.no_result:
        print_ocr_result(result)

    return 0


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