#!/usr/bin/env bash
# ============================================================
#  LeRobot 一键环境 — 3080 主家 / 3090 训练机（两台 Ubuntu 都跑一遍）
#  用法:
#    nano ~/robot/setup_lerobot.sh   # 粘入本文件
#    bash ~/robot/setup_lerobot.sh
#  跑完把【整段输出】粘回给助手，我带你走 P1。
# ============================================================
set -uo pipefail

WORKDIR="${HOME}/robot"
LEROBOT_DIR="${WORKDIR}/lerobot"
VENV="${WORKDIR}/.venv"
PY=3.12
# LeRobot extras：feetech(SO-101) + lekiwi(底盘) + dataset(采数) + smolvla(主力VLA) + pi0(π0/π0-FAST)
EXTRAS="feetech,lekiwi,dataset,smolvla,pi0"

hr(){ echo; echo "========== $* =========="; }

hr "[1/7] 环境探测"
if [ -r /etc/os-release ]; then . /etc/os-release; echo "OS: ${PRETTY_NAME:-?}"; fi
if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
else
  echo "!! 未检测到 nvidia-smi —— GPU 驱动可能没装（3090 应能识别）。先装驱动再继续。"
fi
echo "系统 python3: $(python3 --version 2>&1)"
echo "有 python3.12: $([ -x "$(command -v python3.12)" ] && echo 是 || echo 否)"
echo "有 uv:          $([ -x "$(command -v uv)" ] && echo 是 || echo 否)"
echo "有 git:         $(git --version 2>&1)"
echo "有 ffmpeg:      $([ -x "$(command -v ffmpeg)" ] && echo 是 || echo 否)"

hr "[2/7] 网络/镜像策略（直连 vs 国内镜像）"
PIP_URL="https://pypi.org/simple"
if curl -fsS --max-time 6 https://pypi.org >/dev/null 2>&1; then
  echo "pypi 直连可达 → 官方源"
else
  echo "pypi 直连不可达 → 清华镜像"; PIP_URL="https://pypi.tuna.tsinghua.edu.cn/simple"
fi
if curl -fsS --max-time 6 https://huggingface.co >/dev/null 2>&1; then
  echo "HF 直连可达 → 官方"; HF_ENDPOINT="https://huggingface.co"
else
  echo "HF 直连不可达 → hf-mirror"; HF_ENDPOINT="https://hf-mirror.com"
fi
export HF_ENDPOINT
echo "→ PIP_URL=$PIP_URL"; echo "→ HF_ENDPOINT=$HF_ENDPOINT"

hr "[3/7] Python ${PY} venv (${VENV})"
mkdir -p "${WORKDIR}"
if [ -d "${VENV}" ]; then echo "venv 已存在，复用"; else
  if command -v uv >/dev/null 2>&1; then
    uv venv --python ${PY} "${VENV}"
  elif command -v python3.12 >/dev/null 2>&1; then
    python3.12 -m venv "${VENV}"
  elif command -v python3.11 >/dev/null 2>&1; then
    echo "无 3.12，用 3.11 兜底（LeRobot 要 3.12，可能装不上，建议装 uv 或 deadsnakes 3.12）"
    python3.11 -m venv "${VENV}"
  else
    echo "!! 没有 3.12/3.11/uv。装 uv:  curl -LsSf https://astral.sh/uv/install.sh | sh   然后重跑"; exit 1
  fi
fi
# shellcheck disable=SC1091
source "${VENV}/bin/activate"
python -m pip install -U pip -i "${PIP_URL}"
echo "venv python: $(python --version)  ($(which python))"

hr "[4/7] clone LeRobot → ${LEROBOT_DIR}"
if [ -d "${LEROBOT_DIR}/.git" ]; then
  echo "已存在 → git pull 更新"; ( cd "${LEROBOT_DIR}" && git pull ) || true
else
  if curl -fsS --max-time 6 https://github.com >/dev/null 2>&1; then
    git clone https://github.com/huggingface/lerobot.git "${LEROBOT_DIR}"
  else
    echo "github 直连慢/不可达 → 走 ghproxy 加速"
    git clone https://ghproxy.com/https://github.com/huggingface/lerobot.git "${LEROBOT_DIR}"
  fi
fi
cd "${LEROBOT_DIR}"
echo "当前 commit: $(git log -1 --oneline 2>/dev/null)"

hr "[5/7] 安装 LeRobot [${EXTRAS}]"
pip install -e ".[${EXTRAS}]" -i "${PIP_URL}" 2>&1 | tail -n 25
# ffmpeg（LeRobot 编视频需要）
if ! command -v ffmpeg >/dev/null 2>&1; then
  echo "装 ffmpeg ..."; ( sudo apt-get install -y ffmpeg 2>&1 || conda install -y -c conda-forge ffmpeg 2>&1 ) | tail -n 5
fi

hr "[6/7] 验证"
python - <<'PY'
import importlib, torch
print("torch", torch.__version__,
      "| cuda", torch.cuda.is_available(),
      "| gpu", (torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A"),
      "| vram_GB", round(torch.cuda.get_device_properties(0).total_memory/1e9,1) if torch.cuda.is_available() else 0)
try:
    import lerobot; print("lerobot", getattr(lerobot, "__version__", "?"))
except Exception as e:
    print("lerobot import 失败:", e)
for m in ["lerobot.policies.smolvla","lerobot.policies.pi0","lerobot.policies.act"]:
    try: importlib.import_module(m); print("OK ", m)
    except Exception as e: print("MISS", m, "->", e)
PY
echo "--- CLI 命令 ---"
for c in lerobot-find-port lerobot-find-cameras lerobot-calibrate lerobot-teleoperate lerobot-record lerobot-train lerobot-eval lerobot-rollout; do
  command -v "$c" >/dev/null 2>&1 && echo "OK   $c" || echo "MISS $c"
done

hr "[7/7] 完成 — 之后每次用 3090 前先:"
cat <<EOF
  source ${VENV}/bin/activate
  cd ${LEROBOT_DIR}
  export HF_ENDPOINT=${HF_ENDPOINT}
  huggingface-cli login      # 首次: 输你的 HF token
EOF
echo "✅ 把上面 [1]~[7] 的完整输出粘回给助手。"
echo "   若全 OK，我们就进 P1：插第一条 SO-101 → lerobot-find-port → 校准 → 遥操。"