#!/usr/bin/env python3
"""权重模块的统一命令行入口：转换 / 标签 / 列表 / 自检。

通常不直接调用本文件，而是走薄封装（它会先把环境准备好）：

    bash ~/ros2/scripts/weights.sh list
    bash ~/ros2/scripts/weights.sh labels detect_best_0
    bash ~/ros2/scripts/weights.sh labels detect_best_0 --set "工件,通止规"
    bash ~/ros2/scripts/weights.sh convert staging/best_v3.pt --name workpiece --group yolo

把"纯逻辑"放在 ``arm_application.weights_layout`` / ``arm_application.model_labels``
里，本文件只负责：解析参数、打印人话、真正调用 PyTorch/Ultralytics/TensorRT。

为什么单独做这个模块：原先只有 ``export_yolo_trt.sh``，只能"就地"把
``<同名>.onnx`` 变成 ``<同名>.engine``，既不能改名、不能归位、也不管标签。
实际用起来必须手工搬文件 + 手改 site.yaml，而**标签在 engine 那一环会静默丢失**
（原因见 model_labels 模块头）。这里把三件事一次做完：
转换 → 规范改名 → 放到权重目录 → 顺带把标签落成可编辑的文件。
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
from pathlib import Path
from typing import Dict, List, Optional, Sequence

# ---------------------------------------------------------------------------
# 让脚本能在没 colcon build / 没 source install 的情况下直接 import
# ---------------------------------------------------------------------------


def _bootstrap_import_path() -> Optional[Path]:
    """把 ``src/arm_application`` 插进 sys.path，返回仓库根（找不到返回 None）。

    ``arm_application/__init__.py`` 里只有文档字符串、没有 ROS 依赖，
    所以这样 import 不会把 ROS 拖进来。
    """
    here = Path(__file__).resolve()
    for cand in [here, *here.parents]:
        pkg = cand / "src" / "arm_application"
        if (pkg / "arm_application" / "weights_layout.py").is_file():
            if str(pkg) not in sys.path:
                sys.path.insert(0, str(pkg))
            return cand
    return None


_REPO_ROOT = _bootstrap_import_path()

from arm_application.model_labels import (  # noqa: E402
    LABELS_SUFFIX,
    SOURCE_LABELS_ZH,
    SOURCE_PLACEHOLDER,
    SOURCE_SIDECAR,
    describe,
    format_labels_yaml,
    labels_from_model,
    labels_from_onnx,
    labels_path_for,
    parse_label_list,
    read_label_text_file,
    read_labels,
    write_labels,
)
from arm_application.weights_layout import (  # noqa: E402
    CONVERTIBLE_INPUT_EXTS,
    STAGING_DIRNAME,
    WEIGHT_EXTS,
    ModelEntry,
    canonical_path,
    conversion_chain,
    default_weights_root,
    discover,
    find_model,
    groups,
    model_labels,
    plan_conversion,
    staging_dir,
    staging_files,
    validate_name,
)

# 默认转换输入尺寸。site.yaml 的 yolo.imgsz 优先（由 weights.sh 通过环境变量传入）。
DEFAULT_IMGSZ = 640


# ---------------------------------------------------------------------------
# 小工具
# ---------------------------------------------------------------------------


def _root(args) -> Path:
    return Path(args.root).expanduser() if getattr(args, "root", None) else default_weights_root()


def _imgsz(args) -> int:
    """输入尺寸：命令行 > site.yaml（环境变量） > 默认 640。"""
    if getattr(args, "imgsz", None):
        return int(args.imgsz)
    env = os.environ.get("SITE_YOLO_IMGSZ")
    if env:
        try:
            return int(env)
        except ValueError:
            pass
    return DEFAULT_IMGSZ


def _err(msg: str) -> None:
    print(f"错误: {msg}", file=sys.stderr)


def _ok(msg: str) -> None:
    print(f"✓ {msg}")


def _warn(msg: str) -> None:
    print(f"! {msg}")


def _hr(char: str = "─", n: int = 74) -> None:
    print(char * n)


def _print_table(headers: Sequence[str], rows: Sequence[Sequence[str]]) -> None:
    """左对齐打印一张表。宽度按显示宽度（东亚字符按 2 算）估算。"""

    def w(s: str) -> int:
        return sum(2 if ord(c) > 0x2E80 else 1 for c in s)

    widths = [w(h) for h in headers]
    for r in rows:
        for i, c in enumerate(r):
            widths[i] = max(widths[i], w(c))

    def pad(s: str, width: int) -> str:
        return s + " " * max(0, width - w(s))

    print("  ".join(pad(h, widths[i]) for i, h in enumerate(headers)))
    print("  ".join("-" * widths[i] for i in range(len(headers))))
    for r in rows:
        print("  ".join(pad(c, widths[i]) for i, c in enumerate(r)))


def _mode_of(p: Path) -> str:
    """文件可读性/大小的人话描述，用于列表显示。"""
    try:
        return f"{p.stat().st_size / 1e6:.1f}MB"
    except OSError:
        return "?"


# ---------------------------------------------------------------------------
# list
# ---------------------------------------------------------------------------


def cmd_list(args) -> int:
    root = _root(args)
    if not root.is_dir():
        _err(f"权重目录不存在: {root}")
        return 2

    entries = discover(root)
    if not entries:
        _warn(f"{root} 下没找到任何权重文件（后缀: {', '.join(WEIGHT_EXTS)}）")
        return 0

    payload = []
    rows = []
    for e in entries:
        ls = model_labels(e, deep=args.deep)
        if ls.source == SOURCE_PLACEHOLDER:
            label_cell = "— 缺标签"
        else:
            label_cell = f"{', '.join(ls.sorted_names())}"
        rows.append(
            [
                e.group,
                e.name,
                " ".join(e.exts) or "—",
                label_cell,
                SOURCE_LABELS_ZH.get(ls.source, ls.source),
            ]
        )
        payload.append(
            {
                "group": e.group,
                "name": e.name,
                "formats": e.exts,
                "files": {k: str(v) for k, v in e.files.items()},
                "labels": ls.names,
                "labels_source": ls.source,
                "labels_path": str(ls.path) if ls.path else None,
            }
        )

    if args.json:
        print(json.dumps(payload, ensure_ascii=False, indent=2))
        return 0

    print(f"权重目录: {root}")
    _hr()
    _print_table(["分组", "模型名", "已有格式", "标签", "标签来源"], rows)
    _hr()
    missing = [e for e in entries if model_labels(e).source == SOURCE_PLACEHOLDER]
    print(f"共 {len(entries)} 个模型，其中 {len(missing)} 个缺标签。")
    if missing and not args.deep:
        print("提示: 加 --deep 可再试着从 .pt/.engine 自身读标签（会加载模型，慢）。")
    if missing:
        print("补标签: bash scripts/weights.sh labels <模型名> --from-model")
    return 0


# ---------------------------------------------------------------------------
# labels
# ---------------------------------------------------------------------------


def _resolve_entry(args, root: Path) -> Optional[ModelEntry]:
    entry = find_model(root, args.model, getattr(args, "group", None))
    if entry is None:
        # 提示写成 ``分组/模型名``：这正是 ``list`` 里显示的样子，可以直接复制粘贴，
        # 比只列裸名字更不容易再踩一次坑。
        avail = ", ".join(f"{e.group}/{e.name}" for e in discover(root)) or "(无)"
        _err(
            f"找不到模型 '{args.model}'。现有模型（可整串复制）: {avail}\n"
            f"     也可以只写裸名字，或加 --group 指定分组。全部列出: "
            f"bash ~/ros2/scripts/weights.sh list"
        )
    return entry


def cmd_labels(args) -> int:
    root = _root(args)
    entry = _resolve_entry(args, root)
    if entry is None:
        return 2

    # 修改类操作 ------------------------------------------------------------
    if args.clear:
        p = entry.labels_path
        if p.is_file():
            p.unlink()
            _ok(f"已删除标签文件 {p}")
        else:
            _warn(f"本来就没有标签文件 {p}")
        return 0

    if args.set is not None and args.set_file:
        _err("--set 与 --set-file 不能同时用")
        return 2

    new_names: Optional[Dict[int, str]] = None
    if args.set is not None:
        new_names = parse_label_list(args.set)
        if not new_names:
            _err("--set 解析不出任何类别名")
            return 2
    elif args.set_file:
        src = Path(args.set_file).expanduser()
        if not src.is_file():
            _err(f"标签文件不存在: {src}")
            return 2
        new_names = read_label_text_file(src)
        if not new_names:
            _err(f"{src} 里没读到任何类别名（一行一个）")
            return 2

    if new_names is None and args.from_model:
        # 从模型内嵌 metadata 抄一份出来，目的是"变成可编辑的旁挂文件"。
        for p in entry.sorted_files():
            got = None
            if p.suffix.lower() == ".onnx":
                from arm_application.model_labels import labels_from_onnx

                got = labels_from_onnx(p)
            if got is None:
                got = labels_from_model(p)
            if got:
                new_names = got[0]
                print(f"从 {p.name} 读到 {len(new_names)} 个类别名")
                break
        if new_names is None:
            _err(f"{entry.name} 的模型文件里读不到标签，只能用 --set 手工指定")
            return 2

    if new_names is not None:
        dst = write_labels(entry.path(entry.exts[0]) or entry.labels_path, new_names)
        _ok(f"已写入 {dst}")
        print(f"   类别({len(new_names)}): {', '.join(new_names[k] for k in sorted(new_names))}")
        return 0

    # 读取（默认动作）------------------------------------------------------
    ls = read_labels(entry.path(entry.exts[0]), allow_model_metadata=args.deep) if entry.exts else None
    if ls is None or ls.is_placeholder:
        ls = model_labels(entry, deep=args.deep)

    if args.json:
        print(
            json.dumps(
                {
                    "model": entry.name,
                    "group": entry.group,
                    "labels": ls.names,
                    "source": ls.source,
                    "source_zh": SOURCE_LABELS_ZH.get(ls.source, ls.source),
                    "source_file": str(ls.path) if ls.path else None,
                    "expected_labels_path": str(entry.labels_path),
                },
                ensure_ascii=False,
                indent=2,
            )
        )
        return 0

    print(f"模型: {entry.group}/{entry.name}")
    print(f"格式: {' '.join(entry.exts) or '—'}")
    _hr()
    if ls.is_placeholder:
        _warn(f"没有标签。{describe(ls)}")
        print(f"   期望的标签文件位置: {entry.labels_path}")
        print("   修复方式（三选一）:")
        print(f"     1) 从模型内嵌 metadata 抄一份出来再改: weights.sh labels {entry.name} --from-model")
        print(f'     2) 直接指定类别名:                     weights.sh labels {entry.name} --set "工件,通止规"')
        print(f"     3) 手工创建 {entry.labels_path.name}")
        return 1

    print(f"标签来源: {SOURCE_LABELS_ZH.get(ls.source, ls.source)}")
    if ls.path:
        print(f"来源文件: {ls.path}")
    _hr()
    for k in sorted(ls.names):
        print(f"  {k}: {ls.names[k]}")
    return 0


# ---------------------------------------------------------------------------
# convert
# ---------------------------------------------------------------------------


def _export_onnx(src: Path, imgsz: int, dst: Path, *, simplify: bool = True) -> Path:
    """YOLO .pt/.pth → ONNX，直接产出到 ``dst``。

    **关键：先把 .pt 复制到临时目录再导出。**
    Ultralytics 会把 onnx 写在**源文件旁边**（同名同目录），我们对源 .pt 直接导出时，
    会把它旁边已有的 ``<同名>.onnx`` 静默覆盖掉 —— 那是用户自己的产物，
    可能还是不同 imgsz 导出的。更糟的是随后的搬运会把它彻底删掉。
    实测就发生过：转换 detect_hc90.pt 时，models/yolo/detect_hc90.onnx 被覆盖后又被搬走。
    在临时目录里导出，源目录一个字节都不会动。
    """
    import tempfile

    os.environ.setdefault("YOLO_AUTOINSTALL", "false")
    from ultralytics import YOLO

    dst.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="weights_export_") as td:
        work = Path(td) / src.name
        shutil.copy2(src, work)
        model = YOLO(str(work))
        out = Path(
            str(
                model.export(
                    format="onnx",
                    imgsz=int(imgsz),
                    device=0,
                    simplify=simplify,
                    dynamic=False,
                )
            )
        ).expanduser()
        if not out.is_file():
            raise RuntimeError(f"Ultralytics 说导出了 {out}，但文件不存在")
        # 临时目录马上会被删，所以这里必须是 move。
        if out.resolve() != dst.resolve():
            shutil.move(str(out), str(dst))
    return dst


def _build_engine(onnx_src: Path, dst: Path, *, fp16: bool = True, workspace_gb: int = 4) -> Path:
    """ONNX → TensorRT engine，直接写到 dst。

    这里刻意绕开 Ultralytics 的 engine 导出：它对 .onnx 会按发行包名检查
    onnxruntime，环境装的是 onnxruntime-gpu 时会触发 pip AutoUpdate（可能卡死）。
    直接用 TensorRT 的 Python API，行为可控、错误信息也更清楚。
    """
    import tensorrt as trt

    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    # TRT 10 里 EXPLICIT_BATCH 已是默认行为，但这个 flag 仍然接受；
    # 保持和 export_yolo_trt.sh 一致，避免不同版本行为差异。
    flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    network = builder.create_network(flag)
    parser = trt.OnnxParser(network, logger)

    if not parser.parse(onnx_src.read_bytes()):
        msgs = [str(parser.get_error(i)) for i in range(parser.num_errors)]
        raise RuntimeError("ONNX 解析失败:\n  " + "\n  ".join(msgs))

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, int(workspace_gb) << 30)
    if fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)

    serialized = builder.build_serialized_network(network, config)
    if serialized is None:
        raise RuntimeError("TensorRT 构建失败（返回空）。可试着 --no-fp16 或换 imgsz。")

    dst.parent.mkdir(parents=True, exist_ok=True)
    dst.write_bytes(bytes(serialized))
    return dst


def _materialize_labels(src: Path, dest_weights: Path, explicit: Optional[str]) -> Optional[Path]:
    """转换后把标签落到目标权重旁边。

    这是防止"engine 丢标签"再次发生的**关键一步**：转换时就把标签固化成
    可编辑的旁挂文件，而不是等到运行时才发现读不到名字。

    优先级：命令行 --labels > 源文件/中间产物的内嵌 metadata。
    已经有旁挂文件时不覆盖（避免把用户手改的内容冲掉）。
    """
    dst_labels = labels_path_for(dest_weights)
    if dst_labels.is_file():
        _warn(f"标签文件已存在，保持不变: {dst_labels.name}")
        return dst_labels

    if explicit:
        names = parse_label_list(explicit)
        if names:
            return write_labels(dest_weights, names)

    for cand in (src, src.with_suffix(".onnx"), dest_weights):
        if not cand.is_file():
            continue
        if cand.suffix.lower() == ".onnx":
            from arm_application.model_labels import labels_from_onnx

            got = labels_from_onnx(cand)
        else:
            got = labels_from_model(cand)
        if got:
            _ok(f"标签来自 {cand.name}: {', '.join(got[0][k] for k in sorted(got[0]))}")
            return write_labels(dest_weights, got[0])
    return None


def cmd_convert(args) -> int:
    root = _root(args)
    src = Path(args.input).expanduser()
    if not src.is_file():
        # 相对路径优先按**权重目录**解析，而不是当前工作目录。
        # 否则 `weights.sh convert staging/x.pt` 会去 <cwd>/staging 找，
        # 而用户心里想的是 models/staging/x.pt —— 报"文件不存在"很费解。
        cands = []
        if not src.is_absolute():
            cands.append(root / args.input)              # models/staging/x.pt
            cands.append(staging_dir(root) / args.input)  # 只写文件名: models/staging/x.pt
            cands.append(Path(args.input))               # 退一步：相对 cwd
        found = next((c for c in cands if c.is_file()), None)
        if found is None:
            _err(
                f"输入文件不存在: {args.input}\n"
                f"      找过: " + ", ".join(str(c) for c in cands) +
                f"\n      把待转换的文件放进 {staging_dir(root)} 再试。"
            )
            return 2
        src = found

    name = (args.name or src.stem).strip()
    bad = validate_name(name)
    if bad:
        _err(bad)
        return 2

    group = args.group
    chain, _outs = plan_conversion(root, src, name=name, group=group, target=args.to)

    print(f"输入: {src}")
    print(f"输出: {root / group}/{name}.<ext>")
    print(f"计划: {chain.describe()}")
    print(f"尺寸: imgsz={_imgsz(args)}")
    if args.to != "all":
        print(f"目标: 只产出 {args.to}")

    in_staging = src.parent.resolve() == staging_dir(root).resolve()
    # 归位路径会把源文件本身搬到目标去；那种情况下它已经不在 staging 了，
    # 下面不能再"归档"一次（否则会 FileNotFoundError）。
    src_consumed = False

    if chain.is_noop:
        ext = src.suffix.lower().lstrip(".")
        if ext in WEIGHT_EXTS:
            dst = canonical_path(root, group, name, ext)
            if dst.resolve() == src.resolve():
                _warn(f"{src.name} 已经是最终产物，且就在目标位置，无需处理")
                return 0
            print(f"{ext} 无需转换，直接归位 → {dst}")
            if args.dry_run:
                return 0
            dst.parent.mkdir(parents=True, exist_ok=True)
            # 归位本身就是"移动源文件"。若用户要求保留 staging 原件，
            # 就改成复制，否则源文件会连同 staging 一起消失。
            if in_staging and args.keep_staging:
                shutil.copy2(str(src), str(dst))
            else:
                shutil.move(str(src), str(dst))
                src_consumed = True
            _ok(f"已归位 {dst}")
            _materialize_labels(dst, dst, args.labels)
            if in_staging and not src_consumed:
                _ok(f"staging 原件保留在 {src}（--keep-staging）")
            elif in_staging:
                _ok(f"staging 已清空该文件（它已变成 {dst.name}）")
            return 0
        _err(
            f"不认识的输入格式 '{src.suffix}'。"
            f"可转换的输入: {', '.join('.' + e for e in CONVERTIBLE_INPUT_EXTS)}；"
            f"已是最终产物的: .engine/.trt/.plan"
        )
        return 2

    if args.dry_run:
        print("(--dry-run，未实际执行)")
        return 0

    # 真正执行 --------------------------------------------------------------
    produced: Dict[str, Path] = {}
    try:
        for ext in chain.produce:
            dst = canonical_path(root, group, name, ext)
            if dst.exists() and not args.force:
                raise FileExistsError(f"{dst} 已存在（加 --force 覆盖）")
            if ext == "onnx":
                if src.suffix.lower() == ".onnx":
                    # 源本身就是 onnx：复制而不是移动 —— 用户给的文件不该消失。
                    dst.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(src), str(dst))
                else:
                    _export_onnx(src, _imgsz(args), dst)
                produced["onnx"] = dst
                _ok(f"ONNX  → {dst}  ({_mode_of(dst)})")
            elif ext == "engine":
                onnx_in = produced.get("onnx") or src
                _build_engine(onnx_in, dst, fp16=not args.no_fp16)
                produced["engine"] = dst
                _ok(f"engine→ {dst}  ({_mode_of(dst)})")
    except Exception as exc:  # noqa: BLE001
        _err(f"转换失败: {type(exc).__name__}: {exc}")
        return 1

    # 标签 ------------------------------------------------------------------
    anchor = produced.get("onnx") or produced.get("engine") or src
    lp = _materialize_labels(src, anchor, args.labels)
    if lp:
        _ok(f"标签  → {lp}")
    else:
        _warn("没有可用的标签（源模型里也没有）。请手动补：")
        print(f"     bash scripts/weights.sh labels {name} --set \"类别1,类别2\"")

    # staging 收尾 ----------------------------------------------------------
    if in_staging and not args.keep_staging and not src_consumed:
        done = staging_dir(root) / "done"
        done.mkdir(parents=True, exist_ok=True)
        target = done / src.name
        # 同名已存在时加个序号，避免静默覆盖上一次的源文件。
        n = 1
        while target.exists():
            target = done / f"{src.stem}.{n}{src.suffix}"
            n += 1
        shutil.move(str(src), str(target))
        _ok(f"已把 staging 里的源文件挪到 {target}（加 --keep-staging 可保留）")

    print()
    _ok(f"完成。共产出 {len(produced)} 个文件。")
    print(f"下一步: 把 config/site.yaml 的 paths.yolo_weights 指向 {anchor}")
    print("        或直接跑 bash scripts/weights.sh list 确认标签读得出来")
    return 0


# ---------------------------------------------------------------------------
# doctor
# ---------------------------------------------------------------------------


def _read_site_yaml() -> Dict:
    if _REPO_ROOT is None:
        return {}
    p = _REPO_ROOT / "config" / "site.yaml"
    if not p.is_file():
        return {}
    try:
        import yaml

        return yaml.safe_load(p.read_text(encoding="utf-8")) or {}
    except Exception:  # noqa: BLE001
        return {}


def _engine_usable(p: Path) -> Optional[str]:
    """试着反序列化 engine。返回 None 表示可用，否则返回错误摘要。

    引擎和 TensorRT 版本是绑死的：换过 TensorRT 之后旧引擎会直接加载失败。
    这个检查就是为了把这种"静默失效"提前暴露出来。
    """
    try:
        import tensorrt as trt
    except Exception as exc:  # noqa: BLE001
        return f"无法 import tensorrt: {exc}"
    try:
        rt = trt.Runtime(trt.Logger(trt.Logger.ERROR))
        eng = rt.deserialize_cuda_engine(p.read_bytes())
        if eng is None:
            return "反序列化返回空（多半是 TensorRT 版本不匹配）"
        return None
    except Exception as exc:  # noqa: BLE001
        msg = str(exc)
        if "incompatible serialization version" in msg or "magicTag" in msg:
            return "TensorRT 版本不匹配（需重新导出：weights.sh convert <onnx> --name ...）"
        return f"{type(exc).__name__}: {msg[:160]}"


def cmd_doctor(args) -> int:
    root = _root(args)
    problems = 0
    print(f"权重目录: {root}")
    _hr("═")

    # 1) 逐模型体检 --------------------------------------------------------
    entries = discover(root)
    print(f"[1/4] 模型（{len(entries)} 个）")
    for e in entries:
        ls = model_labels(e, deep=args.deep)
        tag = "OK " if ls.source != SOURCE_PLACEHOLDER else "缺标签"
        if ls.source == SOURCE_PLACEHOLDER:
            problems += 1
        print(f"  [{tag}] {e.group}/{e.name}  格式: {' '.join(e.exts)}")
        if ls.source == SOURCE_PLACEHOLDER:
            print(f"          期望: {e.labels_path.name}")
        # 只有 .pt 没有 onnx/engine 时，two_stage/onnx 后端会用不了
        if "onnx" not in e.files and any(x in e.files for x in ("engine", "trt")):
            print("          注意: 只有 engine 没有 onnx，CPU/ORT 后端无处可用")

    # 2) engine 版本 -------------------------------------------------------
    print()
    print("[2/4] engine 可用性（TensorRT 版本绑定检查）")
    engines = [
        p for e in entries for ext in ("engine", "trt") if (p := e.path(ext)) is not None
    ]
    if not engines:
        print("  (没有 engine)")
    for p in engines:
        if args.check_engines:
            err = _engine_usable(p)
            if err:
                problems += 1
                print(f"  [失效] {p.name}: {err}")
            else:
                print(f"  [可用] {p.name}")
        else:
            print(f"  [跳过] {p.name}  (加 --check-engines 实测)")

    # 3) 配置引用 ----------------------------------------------------------
    print()
    print("[3/4] config/site.yaml 的引用")
    site = _read_site_yaml()
    paths = site.get("paths") or {}
    refs: List[tuple] = []
    if paths.get("yolo_weights"):
        refs.append(("paths.yolo_weights", paths["yolo_weights"]))
    if paths.get("hole_weights"):
        refs.append(("paths.hole_weights", paths["hole_weights"]))
    if paths.get("hole_workpiece_weights"):
        refs.append(("paths.hole_workpiece_weights", paths["hole_workpiece_weights"]))
    hd = site.get("hole_detect") or {}
    if hd.get("weights"):
        refs.append(("hole_detect.weights", hd["weights"]))
    if hd.get("workpiece_weights"):
        refs.append(("hole_detect.workpiece_weights", hd["workpiece_weights"]))

    if not refs:
        print("  (没读到引用)")
    for key, val in refs:
        p = Path(str(val)).expanduser()
        if not p.is_file():
            if p.suffix == ".pt" and p.with_suffix(".onnx").is_file() is False:
                problems += 1
                print(f"  [悬空] {key} → {p}  (文件不存在)")
            else:
                problems += 1
                print(f"  [悬空] {key} → {p}  (文件不存在)")
            continue
        ls = read_labels(p, allow_model_metadata=False)
        origin = SOURCE_LABELS_ZH.get(ls.source, ls.source)
        print(f"  [OK ] {key} → {p.name}  标签: {origin}")
        if ls.is_placeholder:
            problems += 1
            print(f"         → 该模型读不到标签，运行时类名会变成 class_0")

    backend = (site.get("yolo") or {}).get("backend")
    if backend == "tensorrt":
        yw = Path(str(paths.get("yolo_weights") or "")).expanduser()
        if yw.name and not yw.with_suffix(".engine").is_file():
            problems += 1
            print(
                f"  [风险] yolo.backend=tensorrt 但 {yw.with_suffix('.engine').name} 不存在，"
                "启动会直接失败"
            )

    # 4) staging -----------------------------------------------------------
    print()
    print("[4/4] staging（待转换）")
    st = staging_files(root)
    if not st:
        print("  (空)")
    for p in st:
        ext = p.suffix.lower().lstrip(".")
        if ext in CONVERTIBLE_INPUT_EXTS:
            print(f"  [待转换] {p.name}  → weights.sh convert {STAGING_DIRNAME}/{p.name} --name <模型名>")
        else:
            print(f"  [需人工] {p.name}  格式 .{ext} 不在支持列表（{', '.join(CONVERTIBLE_INPUT_EXTS)}）")

    _hr("═")
    if problems:
        print(f"发现 {problems} 处问题。")
        return 1
    print("没有发现问题。")
    return 0


# ---------------------------------------------------------------------------
# show：把模块约定直接打出来，省得翻文档
# ---------------------------------------------------------------------------


def cmd_show(args) -> int:
    root = _root(args)
    print("权重模块约定")
    _hr()
    print(f"权重根目录      : {root}")
    print(f"仓库根          : {_REPO_ROOT or '(未找到)'}")
    print(f"待转换目录      : {staging_dir(root)}")
    print(f"标签文件后缀    : {LABELS_SUFFIX}   （<模型名>{LABELS_SUFFIX}，放权重旁边，进 git）")
    print(f"权重后缀        : {', '.join('.' + e for e in WEIGHT_EXTS)}   （二进制，不进 git）")
    print(f"已存在的分组    : {', '.join(groups(root)) or '(无)'}")
    print()
    print("命令")
    _hr()
    print("  weights.sh list [--deep] [--json]")
    print("  weights.sh labels <模型名> [--group G] [--set a,b,c] [--set-file f.txt]")
    print("                             [--from-model] [--clear] [--deep] [--json]")
    print("  weights.sh convert <文件> --name <模型名> [--group yolo] [--to all|onnx|engine]")
    print("                            [--imgsz 640] [--labels a,b] [--force] [--dry-run]")
    print("                            [--no-fp16] [--keep-staging]")
    print("  weights.sh doctor [--deep] [--check-engines]")
    print()
    print("标签优先级: 旁挂文件 > ONNX 内嵌 metadata > 模型内嵌 metadata > 占位符(class_0)")
    return 0


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="weights.sh",
        description="模型权重模块：转换（PT→ONNX→TensorRT）、标签读写、列表、自检",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "示例:\n"
            '  weights.sh list\n'
            '  weights.sh labels detect_best_0\n'
            '  weights.sh labels detect_best_0 --set "工件,通止规"\n'
            '  weights.sh convert staging/best_v3.pt --name workpiece\n'
            '  weights.sh convert models/yolo/detect_best_0.onnx --name detect_best_0 --to engine\n'
        ),
    )
    p.add_argument("--root", help="权重根目录（默认自动探测 models/）")
    sub = p.add_subparsers(dest="cmd", required=True)

    s = sub.add_parser("list", help="列出所有权重及其标签")
    s.add_argument("--deep", action="store_true", help="也尝试从 .pt/.engine 自身读标签（慢）")
    s.add_argument("--json", action="store_true", help="以 JSON 输出")
    s.set_defaults(func=cmd_list)

    s = sub.add_parser("labels", help="查看/修改模型标签")
    s.add_argument("model", help="模型名（权重目录里的文件名主干）")
    s.add_argument("--group", help="限定分组")
    s.add_argument("--set", help='直接指定类别名，如 "工件,通止规"（改了立刻生效）')
    s.add_argument("--set-file", dest="set_file", help="从纯文本文件导入（一行一个类名）")
    s.add_argument("--from-model", dest="from_model", action="store_true",
                   help="从模型内嵌 metadata 抄一份到旁挂文件（变成可编辑）")
    s.add_argument("--clear", action="store_true", help="删除旁挂标签文件")
    s.add_argument("--deep", action="store_true", help="读取时也加载 .pt/.engine 自身")
    s.add_argument("--json", action="store_true", help="以 JSON 输出")
    s.set_defaults(func=cmd_labels)

    s = sub.add_parser("convert", help="转换权重并归位")
    s.add_argument("input", help="输入文件（可以是 staging/ 里的相对名）")
    s.add_argument("--name", help="输出模型名（默认沿用输入文件名主干）")
    s.add_argument("--group", default="yolo", help="输出分组子目录（默认 yolo）")
    s.add_argument("--to", default="all", choices=["all", "onnx", "engine"], help="产出目标")
    s.add_argument("--imgsz", type=int, help="输入尺寸（默认取 site.yaml 的 yolo.imgsz）")
    s.add_argument("--labels", help='顺带写标签，如 "工件,通止规"')
    s.add_argument("--force", action="store_true", help="覆盖已存在的产物")
    s.add_argument("--no-fp16", dest="no_fp16", action="store_true", help="engine 不用 FP16")
    s.add_argument("--dry-run", dest="dry_run", action="store_true", help="只打印计划")
    s.add_argument("--keep-staging", dest="keep_staging", action="store_true",
                   help="转换后不把 staging 里的源文件挪走")
    s.set_defaults(func=cmd_convert)

    s = sub.add_parser("doctor", help="自检：缺标签 / 悬空引用 / engine 失效")
    s.add_argument("--deep", action="store_true", help="也加载 .pt/.engine 读标签")
    s.add_argument("--check-engines", dest="check_engines", action="store_true",
                   help="实测反序列化每个 engine（探测 TensorRT 版本不匹配）")
    s.set_defaults(func=cmd_doctor)

    s = sub.add_parser("show", help="打印模块约定与命令")
    s.set_defaults(func=cmd_show)

    return p


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        return int(args.func(args))
    except KeyboardInterrupt:
        _err("被中断")
        return 130


if __name__ == "__main__":
    sys.exit(main())