"""Run PP-OCRv6 parameter combinations against data/ and write a Markdown report."""

from __future__ import annotations

import argparse
import itertools
import json
import shutil
import statistics
import subprocess
import sys
import time
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent
DEFAULT_DATA = ROOT / "data"
DEFAULT_BENCHMARK_ROOT = ROOT / "benchmarks" / "parameter-runs"
MODEL_SIZES = ("tiny", "small", "medium")
PROFILE_OPTIONS = {
    "fast": {
        "doc_orientation": False,
        "doc_unwarping": False,
        "textline_orientation": False,
        "description": "关闭文档方向、去畸变和文本行方向，优先速度",
    },
    "standard": {
        "doc_orientation": True,
        "doc_unwarping": False,
        "textline_orientation": True,
        "description": "开启文档方向和文本行方向，关闭去畸变",
    },
    "robust": {
        "doc_orientation": True,
        "doc_unwarping": True,
        "textline_orientation": True,
        "description": "开启方向处理和文档去畸变，优先复杂输入适应性",
    },
}


@dataclass(frozen=True)
class Scenario:
    model_size: str
    profile: str
    dpi: int
    threads: int | None
    det_limit_side_len: int = 64
    det_limit_type: str = "min"
    det_thresh: float = 0.3
    det_box_thresh: float = 0.6
    det_unclip_ratio: float = 1.5
    rec_score_thresh: float = 0.0
    rec_batch_size: int = 6

    @staticmethod
    def _slug(value: float) -> str:
        return f"{value:g}".replace("-", "m").replace(".", "p")

    @property
    def name(self) -> str:
        thread_name = "auto" if self.threads is None else str(self.threads)
        return (
            f"{self.model_size}-{self.profile}-dpi{self.dpi}-th{thread_name}"
            f"-s{self.det_limit_side_len}{self.det_limit_type}"
            f"-d{self._slug(self.det_thresh)}-b{self._slug(self.det_box_thresh)}"
            f"-u{self._slug(self.det_unclip_ratio)}-r{self._slug(self.rec_score_thresh)}"
            f"-rb{self.rec_batch_size}"
        )


@dataclass
class ScenarioResult:
    scenario: Scenario
    status: str
    exit_code: int
    wall_seconds: float
    model_init_seconds: float | None
    pure_ocr_seconds: float
    image_inference_seconds: float
    pdf_ocr_seconds: float
    processed_files: int
    image_files: int
    pdf_files: int
    ocr_pages: int
    recognized_lines: int
    mean_confidence: float | None
    units_per_second: float | None
    output_dir: str
    log_file: str
    command: list[str]
    error: str | None = None


def parse_threads(values: list[str]) -> list[int | None]:
    parsed: list[int | None] = []
    for value in values:
        normalized = value.strip().lower()
        if normalized == "auto":
            item = None
        else:
            item = int(normalized)
            if item < 1:
                raise ValueError("线程数必须大于等于 1")
        if item not in parsed:
            parsed.append(item)
    return parsed


def build_scenarios(
    model_sizes: list[str],
    profiles: list[str],
    dpis: list[int],
    threads: list[int | None],
    det_limit_side_lens: list[int] | None = None,
    det_limit_types: list[str] | None = None,
    det_thresholds: list[float] | None = None,
    det_box_thresholds: list[float] | None = None,
    det_unclip_ratios: list[float] | None = None,
    rec_score_thresholds: list[float] | None = None,
    rec_batch_sizes: list[int] | None = None,
) -> list[Scenario]:
    dimensions = itertools.product(
        model_sizes,
        profiles,
        dpis,
        threads,
        det_limit_side_lens or [64],
        det_limit_types or ["min"],
        det_thresholds or [0.3],
        det_box_thresholds or [0.6],
        det_unclip_ratios or [1.5],
        rec_score_thresholds or [0.0],
        rec_batch_sizes or [6],
    )
    return [Scenario(*values) for values in dimensions]


def _bool_flag(name: str, enabled: bool) -> str:
    return f"--{name}" if enabled else f"--no-{name}"


def build_command(
    scenario: Scenario,
    *,
    data_dir: Path,
    output_dir: Path,
    device: str,
    device_id: int,
    warmup: int,
    rounds: int,
    pdf_mode: str,
) -> list[str]:
    profile = PROFILE_OPTIONS[scenario.profile]
    command = [
        sys.executable,
        str(ROOT / "ocr.py"),
        str(data_dir),
        "--recursive",
        "--device",
        device,
        "--device-id",
        str(device_id),
        "--model-size",
        scenario.model_size,
        "--pdf-mode",
        pdf_mode,
        "--dpi",
        str(scenario.dpi),
        "--warmup",
        str(warmup),
        "--rounds",
        str(rounds),
        "--text-det-limit-side-len",
        str(scenario.det_limit_side_len),
        "--text-det-limit-type",
        scenario.det_limit_type,
        "--text-det-thresh",
        str(scenario.det_thresh),
        "--text-det-box-thresh",
        str(scenario.det_box_thresh),
        "--text-det-unclip-ratio",
        str(scenario.det_unclip_ratio),
        "--text-rec-score-thresh",
        str(scenario.rec_score_thresh),
        "--text-recognition-batch-size",
        str(scenario.rec_batch_size),
        "--output",
        str(output_dir),
        "--overwrite",
        "--no-result",
        _bool_flag("doc-orientation-classify", profile["doc_orientation"]),
        _bool_flag("doc-unwarping", profile["doc_unwarping"]),
        _bool_flag("textline-orientation", profile["textline_orientation"]),
    ]
    if scenario.threads is not None:
        command.extend(["--threads", str(scenario.threads)])
    return command


def _read_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _collect_confidences(payload: dict[str, Any]) -> list[float]:
    scores: list[float] = []
    for line in payload.get("lines", []):
        value = line.get("score")
        if isinstance(value, (int, float)):
            scores.append(float(value))
    return scores


def collect_scenario_metrics(
    scenario: Scenario,
    *,
    output_dir: Path,
    log_file: Path,
    command: list[str],
    exit_code: int,
    wall_seconds: float,
    error: str | None,
) -> ScenarioResult:
    batch_manifests = sorted((output_dir / "batches").glob("*.json"))
    batch = _read_json(batch_manifests[-1]) if batch_manifests else {}
    image_benchmarks = list((output_dir / "images").rglob("benchmark.json"))
    image_results = list((output_dir / "images").rglob("result.json"))
    pdf_manifests = list((output_dir / "pdfs").rglob("manifest.json"))
    pdf_page_results = list((output_dir / "pdfs").rglob("pages/page-*.json"))

    image_inference = 0.0
    model_init_values: list[float] = []
    recognized_lines = 0
    confidences: list[float] = []
    for path in image_benchmarks:
        payload = _read_json(path)
        image_inference += sum(float(value) for value in payload.get("inference_seconds", {}).get("all", []))
        value = payload.get("model_init_seconds")
        if isinstance(value, (int, float)):
            model_init_values.append(float(value))
    for path in image_results:
        payload = _read_json(path)
        recognized_lines += int(payload.get("summary", {}).get("non_empty_lines", 0))
        confidences.extend(_collect_confidences(payload))

    pdf_ocr_seconds = 0.0
    ocr_pages = 0
    for path in pdf_manifests:
        payload = _read_json(path)
        summary = payload.get("summary", {})
        ocr_pages += int(summary.get("ocr_pages", 0))
        metadata = payload.get("run_metadata", {})
        value = metadata.get("model_init_seconds")
        if isinstance(value, (int, float)):
            model_init_values.append(float(value))
        for page in payload.get("pages", {}).values():
            pdf_ocr_seconds += float(page.get("ocr_seconds", 0.0) or 0.0)
            if page.get("source_type") == "text":
                recognized_lines += int(page.get("detected_lines", 0) or 0)
    for path in pdf_page_results:
        payload = _read_json(path)
        if "lines" in payload:
            recognized_lines += int(payload.get("summary", {}).get("non_empty_lines", 0))
            confidences.extend(_collect_confidences(payload))

    pure_ocr_seconds = image_inference + pdf_ocr_seconds
    image_files = len(image_results)
    pdf_files = len(pdf_manifests)
    units = image_files + ocr_pages
    units_per_second = units / pure_ocr_seconds if pure_ocr_seconds > 0 else None
    status = "completed" if exit_code == 0 else "failed"
    if exit_code == 0 and not batch:
        status = "incomplete"
        error = error or "未找到批处理 manifest"

    return ScenarioResult(
        scenario=scenario,
        status=status,
        exit_code=exit_code,
        wall_seconds=wall_seconds,
        model_init_seconds=max(model_init_values) if model_init_values else None,
        pure_ocr_seconds=pure_ocr_seconds,
        image_inference_seconds=image_inference,
        pdf_ocr_seconds=pdf_ocr_seconds,
        processed_files=int(batch.get("completed_files", image_files + pdf_files)),
        image_files=image_files,
        pdf_files=pdf_files,
        ocr_pages=ocr_pages,
        recognized_lines=recognized_lines,
        mean_confidence=statistics.fmean(confidences) if confidences else None,
        units_per_second=units_per_second,
        output_dir=str(output_dir),
        log_file=str(log_file),
        command=command,
        error=error,
    )


def _format_seconds(value: float | None) -> str:
    return "-" if value is None else f"{value:.3f}"


def _format_float(value: float | None, digits: int = 4) -> str:
    return "-" if value is None else f"{value:.{digits}f}"


def render_report(
    results: list[ScenarioResult],
    *,
    data_dir: Path,
    device: str,
    pdf_mode: str,
    warmup: int,
    rounds: int,
    started_at: datetime,
    finished_at: datetime,
) -> str:
    completed = [result for result in results if result.status == "completed"]
    by_ocr = sorted(completed, key=lambda item: item.pure_ocr_seconds or float("inf"))
    by_wall = sorted(completed, key=lambda item: item.wall_seconds)
    lines = [
        "# PP-OCRv6 参数测试报告",
        "",
        "## 测试信息",
        "",
        f"- 数据目录：`{data_dir}`",
        f"- 运行设备：`{device}`",
        f"- PDF 模式：`{pdf_mode}`",
        f"- 图片预热轮数：`{warmup}`",
        f"- 图片推理轮数：`{rounds}`",
        f"- 场景数量：`{len(results)}`",
        f"- 开始时间：`{started_at.astimezone().isoformat()}`",
        f"- 结束时间：`{finished_at.astimezone().isoformat()}`",
        f"- 总耗时：`{(finished_at - started_at).total_seconds():.3f}s`",
        "",
        "> 纯 OCR 耗时为图片推理耗时与 PDF OCR 页面耗时之和，不含模型初始化、文件扫描、PDF 文本层提取和结果导出。首次下载模型会显著增加墙钟耗时，应优先参考缓存模型后的结果。",
        "",
        "## 汇总对比",
        "",
        "| 排名 | 场景 | 状态 | 墙钟耗时(s) | 模型初始化(s) | 纯OCR耗时(s) | 图片推理(s) | PDF OCR(s) | OCR单位/秒 | 识别行数 | 平均置信度 |",
        "|---:|---|---|---:|---:|---:|---:|---:|---:|---:|---:|",
    ]
    rank_map = {result.scenario.name: index + 1 for index, result in enumerate(by_ocr)}
    for result in results:
        lines.append(
            "| {rank} | `{name}` | {status} | {wall} | {init} | {pure} | {image} | {pdf} | {rate} | {count} | {confidence} |".format(
                rank=rank_map.get(result.scenario.name, "-"),
                name=result.scenario.name,
                status=result.status,
                wall=_format_seconds(result.wall_seconds),
                init=_format_seconds(result.model_init_seconds),
                pure=_format_seconds(result.pure_ocr_seconds),
                image=_format_seconds(result.image_inference_seconds),
                pdf=_format_seconds(result.pdf_ocr_seconds),
                rate=_format_float(result.units_per_second, 3),
                count=result.recognized_lines,
                confidence=_format_float(result.mean_confidence),
            )
        )

    lines.extend(["", "## 结论", ""])
    if by_ocr:
        fastest = by_ocr[0]
        lines.append(
            f"- 纯 OCR 耗时最短：`{fastest.scenario.name}`，耗时 `{fastest.pure_ocr_seconds:.3f}s`，吞吐 `{_format_float(fastest.units_per_second, 3)}` OCR 单位/秒。"
        )
    if by_wall:
        fastest_wall = by_wall[0]
        lines.append(
            f"- 墙钟耗时最短：`{fastest_wall.scenario.name}`，耗时 `{fastest_wall.wall_seconds:.3f}s`。"
        )
    if len(by_ocr) > 1 and by_ocr[-1].pure_ocr_seconds > 0:
        speedup = by_ocr[-1].pure_ocr_seconds / by_ocr[0].pure_ocr_seconds
        lines.append(f"- 最快与最慢完成场景的纯 OCR 速度差约为 `{speedup:.2f}x`。")
    lines.append("- 平均置信度仅作为结果稳定性的辅助观察值，不等价于有标注数据集上的识别准确率。")

    lines.extend(["", "## 参数说明", ""])
    for name, profile in PROFILE_OPTIONS.items():
        lines.append(f"- `{name}`：{profile['description']}。")
    lines.extend(
        [
            "",
            "## 场景明细",
            "",
        ]
    )
    for result in results:
        scenario = result.scenario
        lines.extend(
            [
                f"### {scenario.name}",
                "",
                f"- 状态：`{result.status}`；退出码：`{result.exit_code}`",
                f"- 模型规格：`{scenario.model_size}`",
                f"- 预处理配置：`{scenario.profile}`",
                f"- PDF DPI：`{scenario.dpi}`",
                f"- CPU 线程：`{'auto' if scenario.threads is None else scenario.threads}`",
                f"- 检测边长：`{scenario.det_limit_side_len}`；限制方式：`{scenario.det_limit_type}`",
                f"- 检测阈值：`{scenario.det_thresh}`；文本框阈值：`{scenario.det_box_thresh}`；扩张比例：`{scenario.det_unclip_ratio}`",
                f"- 识别阈值：`{scenario.rec_score_thresh}`；识别 Batch Size：`{scenario.rec_batch_size}`",
                f"- 完成文件：`{result.processed_files}`（图片 `{result.image_files}`，PDF `{result.pdf_files}`，OCR 页 `{result.ocr_pages}`）",
                f"- 墙钟耗时：`{result.wall_seconds:.3f}s`",
                f"- 纯 OCR 耗时：`{result.pure_ocr_seconds:.3f}s`",
                f"- 输出目录：`{result.output_dir}`",
                f"- 运行日志：`{result.log_file}`",
                f"- 命令：`{' '.join(result.command)}`",
            ]
        )
        if result.error:
            lines.append(f"- 错误：`{result.error}`")
        lines.append("")
    return "\n".join(lines).rstrip() + "\n"


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="对 data 目录执行 PP-OCRv6 参数矩阵测试并生成 Markdown 报告",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--data", type=Path, default=DEFAULT_DATA, help="测试数据目录")
    parser.add_argument("--device", choices=("cpu", "gpu"), default="cpu", help="测试设备")
    parser.add_argument("--device-id", type=int, default=0, help="GPU 编号")
    parser.add_argument("--models", nargs="+", choices=MODEL_SIZES, default=list(MODEL_SIZES), help="模型规格列表")
    parser.add_argument("--profiles", nargs="+", choices=tuple(PROFILE_OPTIONS), default=["fast", "standard"], help="预处理配置列表")
    parser.add_argument("--dpis", nargs="+", type=int, default=[144], help="PDF 渲染 DPI 列表")
    parser.add_argument("--threads", nargs="+", default=["auto"], help="CPU 线程列表，可使用 auto")
    parser.add_argument("--det-limit-side-lens", nargs="+", type=int, default=[64], help="文本检测边长列表")
    parser.add_argument("--det-limit-types", nargs="+", choices=("min", "max"), default=["min"], help="文本检测边长限制方式列表")
    parser.add_argument("--det-thresholds", nargs="+", type=float, default=[0.3], help="文本检测像素阈值列表")
    parser.add_argument("--det-box-thresholds", nargs="+", type=float, default=[0.6], help="文本框阈值列表")
    parser.add_argument("--det-unclip-ratios", nargs="+", type=float, default=[1.5], help="文本框扩张比例列表")
    parser.add_argument("--rec-score-thresholds", nargs="+", type=float, default=[0.0], help="识别置信度阈值列表")
    parser.add_argument("--rec-batch-sizes", nargs="+", type=int, default=[6], help="文本识别 Batch Size 列表")
    parser.add_argument("--warmup", type=int, default=1, help="每个场景首张图片预热轮数")
    parser.add_argument("--rounds", type=int, default=1, help="每张图片推理轮数")
    parser.add_argument("--pdf-mode", choices=("ocr", "hybrid", "text"), default="ocr", help="PDF 测试模式；速度对比推荐 ocr")
    parser.add_argument("--output", type=Path, default=None, help="Markdown 报告路径")
    parser.add_argument("--work-dir", type=Path, default=None, help="各场景原始输出目录")
    parser.add_argument("--overwrite", action="store_true", help="允许删除已存在的 work-dir")
    parser.add_argument("--fail-fast", action="store_true", help="任一场景失败后停止")
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    data_dir = args.data.expanduser().resolve()
    if not data_dir.is_dir():
        print(f"测试数据目录不存在: {data_dir}", file=sys.stderr)
        return 2
    if args.warmup < 0 or args.rounds < 1:
        print("--warmup 必须 >= 0，--rounds 必须 >= 1", file=sys.stderr)
        return 2
    if any(dpi < 72 or dpi > 600 for dpi in args.dpis):
        print("--dpis 必须在 72 到 600 之间", file=sys.stderr)
        return 2
    if any(value < 1 for value in args.det_limit_side_lens):
        print("--det-limit-side-lens 必须大于等于 1", file=sys.stderr)
        return 2
    if any(value < 1 for value in args.rec_batch_sizes):
        print("--rec-batch-sizes 必须大于等于 1", file=sys.stderr)
        return 2
    try:
        thread_values = parse_threads(args.threads)
    except ValueError as exc:
        print(f"无效线程参数: {exc}", file=sys.stderr)
        return 2

    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    work_dir = (args.work_dir or DEFAULT_BENCHMARK_ROOT / timestamp).expanduser().resolve()
    report_path = (args.output or ROOT / "benchmarks" / f"参数测试报告-{timestamp}.md").expanduser().resolve()
    if work_dir.exists():
        if not args.overwrite:
            print(f"测试输出目录已存在: {work_dir}；请使用 --overwrite", file=sys.stderr)
            return 2
        shutil.rmtree(work_dir)
    work_dir.mkdir(parents=True, exist_ok=True)
    report_path.parent.mkdir(parents=True, exist_ok=True)

    scenarios = build_scenarios(
        args.models,
        args.profiles,
        args.dpis,
        thread_values,
        args.det_limit_side_lens,
        args.det_limit_types,
        args.det_thresholds,
        args.det_box_thresholds,
        args.det_unclip_ratios,
        args.rec_score_thresholds,
        args.rec_batch_sizes,
    )
    print(f"计划执行 {len(scenarios)} 个参数场景，报告将写入: {report_path}")
    started_at = datetime.now().astimezone()
    results: list[ScenarioResult] = []
    for index, scenario in enumerate(scenarios, 1):
        scenario_dir = work_dir / scenario.name
        output_dir = scenario_dir / "outputs"
        log_file = scenario_dir / "run.log"
        scenario_dir.mkdir(parents=True, exist_ok=True)
        command = build_command(
            scenario,
            data_dir=data_dir,
            output_dir=output_dir,
            device=args.device,
            device_id=args.device_id,
            warmup=args.warmup,
            rounds=args.rounds,
            pdf_mode=args.pdf_mode,
        )
        print(f"[{index}/{len(scenarios)}] {scenario.name}")
        wall_started = time.perf_counter()
        error = None
        try:
            completed = subprocess.run(
                command,
                cwd=ROOT,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                check=False,
            )
            raw_output = completed.stdout or b""
            log_file.write_bytes(raw_output)
            exit_code = completed.returncode
            if exit_code != 0:
                text = raw_output.decode("utf-8", errors="replace")
                error = next((line for line in reversed(text.splitlines()) if line.strip()), f"退出码 {exit_code}")
        except Exception as exc:
            exit_code = 1
            error = f"{type(exc).__name__}: {exc}"
            log_file.write_text(error + "\n", encoding="utf-8")
        wall_seconds = time.perf_counter() - wall_started
        result = collect_scenario_metrics(
            scenario,
            output_dir=output_dir,
            log_file=log_file,
            command=command,
            exit_code=exit_code,
            wall_seconds=wall_seconds,
            error=error,
        )
        results.append(result)
        print(
            f"  状态={result.status} 墙钟={result.wall_seconds:.3f}s "
            f"纯OCR={result.pure_ocr_seconds:.3f}s"
        )
        if result.status != "completed" and args.fail_fast:
            break

    finished_at = datetime.now().astimezone()
    report = render_report(
        results,
        data_dir=data_dir,
        device=args.device,
        pdf_mode=args.pdf_mode,
        warmup=args.warmup,
        rounds=args.rounds,
        started_at=started_at,
        finished_at=finished_at,
    )
    report_path.write_text(report, encoding="utf-8")
    summary_path = work_dir / "summary.json"
    summary_path.write_text(
        json.dumps(
            [
                {
                    **asdict(result),
                    "scenario": asdict(result.scenario),
                }
                for result in results
            ],
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    print(f"测试完成，Markdown 报告: {report_path}")
    print(f"原始汇总 JSON: {summary_path}")
    return 0 if all(result.status == "completed" for result in results) else 3


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