#!/usr/bin/env bash
# 权重模块的统一入口（薄封装）。
#
#   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
#   bash ~/ros2/scripts/weights.sh doctor --check-engines
#
# 本脚本只做三件事：准备环境（ROS + TensorRT 动态库）→ 挑一个能用的 python → 转发。
# 真正的逻辑在 scripts/weights.py。
#
# 完全等价的老命令 export_yolo_trt.sh 仍然可用（它是本脚本的转发壳）。

set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

# 用户显式指定解释器时优先用它；否则试 ros_flexiv。
# 本模块需要的重依赖（ultralytics / onnx / tensorrt / onnxruntime）都在那里。
pick_python() {
  if [[ -n "${WEIGHTS_PY:-}" ]]; then
    echo "${WEIGHTS_PY}"; return 0
  fi
  local cands=(
    "${HOME}/anaconda3/envs/ros_flexiv/bin/python"
    "$(command -v python3 || true)"
  )
  for c in "${cands[@]}"; do
    [[ -n "${c}" && -x "${c}" ]] || continue
    # list/labels/show 只需要 yaml；convert/doctor 才需要 tensorrt。
    # 这里统一要求 yaml，避免半路才发现解释器不对。
    if "${c}" -c 'import yaml' >/dev/null 2>&1; then
      echo "${c}"; return 0
    fi
  done
  return 1
}

PY="$(pick_python)" || {
  echo "错误: 找不到可用的 python（需要能 import yaml）。" >&2
  echo "      可用 WEIGHTS_PY=/path/to/python 指定。" >&2
  exit 1
}

# TensorRT 的 pip 版把 .so 放在非标准位置，TensorRT 的 Python API 和运行时都要用。
# export_yolo_trt.sh 里也是这么做的；这里保持同样做法。
TRT_LIBS="$("${PY}" - <<'PY' 2>/dev/null || true
import pathlib
try:
    import tensorrt
except Exception:
    raise SystemExit(0)
p = pathlib.Path(tensorrt.__file__).resolve().parent
cands = [p, p.parent / "tensorrt_libs", p.parent / "nvidia" / "tensorrt" / "lib"]
print(":".join(str(x) for x in cands if x.exists()))
PY
)"
[[ -n "${TRT_LIBS}" ]] && export LD_LIBRARY_PATH="${TRT_LIBS}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"

# 不让 Ultralytics 自己去 pip 装东西（会卡住或破坏环境）。
export YOLO_AUTOINSTALL=false

# site.yaml 里的 yolo.imgsz 作为 convert 的默认尺寸。
if [[ -f "${ROOT}/scripts/lib/load_site_config.sh" ]]; then
  # shellcheck disable=SC1091
  source "${ROOT}/scripts/lib/load_site_config.sh" >/dev/null 2>&1 || true
fi

exec "${PY}" "${ROOT}/scripts/weights.py" "$@"