#!/usr/bin/env bash
# Ubuntu x86_64 GPU hosts. First host: --resolve; second host: same lock.
set -Eeuo pipefail
trap 'printf "ERROR at line %s: environment NOT verified.\n" "$LINENO" >&2' ERR
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
WORKDIR="${ROBOT_WORKDIR:-${HOME}/robot}"
VENV="${WORKDIR}/.venv-lerobot-0.6.1"
LOCK="${SCRIPT_DIR}/requirements-linux-py312.lock"
INPUT="${SCRIPT_DIR}/requirements.in"
RESOLVE=false
case "${1:-}" in
  --resolve) RESOLVE=true ;;
  --help) echo "Usage: bash $0 [--resolve] (Linux x86_64, uv, NVIDIA driver, ffmpeg)"; exit 0 ;;
  "") ;;
  *) echo "Unknown argument: $1" >&2; exit 2 ;;
esac
[[ $# -le 1 ]] || { echo "Too many arguments" >&2; exit 2; }
[[ "$(uname -s)" == Linux && "$(uname -m)" == x86_64 ]] || {
  echo "Only Linux x86_64 GPU hosts are supported; RPi needs a separate environment." >&2; exit 1;
}
for cmd in uv nvidia-smi ffmpeg sha256sum gcc; do
  command -v "$cmd" >/dev/null || { echo "Missing prerequisite: $cmd" >&2; exit 1; }
done
if ! $RESOLVE && [[ ! -s "$LOCK" ]]; then
  echo "Missing lock: use --resolve on 3080 first; transfer the validated lock and .input.sha256 to 3090." >&2; exit 1
fi
if $RESOLVE && [[ -e "$LOCK" || -e "$LOCK.input.sha256" ]]; then
  echo "Refusing to overwrite existing lock. Review and archive it before an intentional update." >&2; exit 1
fi
mkdir -p "$WORKDIR"
REPORT="${WORKDIR}/reports/$(date -u +%Y%m%dT%H%M%SZ)-$$"
mkdir -p "$REPORT"
exec > >(tee "$REPORT/setup.log") 2>&1
echo "[1/5] Prerequisites"
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
uv --version
cat /etc/os-release > "$REPORT/os-release.txt"
uname -srmo > "$REPORT/kernel.txt"
nvidia-smi > "$REPORT/nvidia-smi.txt"
ffmpeg -version > "$REPORT/ffmpeg.txt"
# Set UV_DEFAULT_INDEX explicitly if a mirror is needed. No automatic Hub/token routing.
echo "[2/5] Dedicated Python 3.12 environment"
if [[ ! -d "$VENV" ]]; then uv venv --python 3.12 "$VENV"; fi
PYTHON="${VENV}/bin/python"
"$PYTHON" -c 'import sys; assert sys.version_info[:2] == (3, 12), sys.version'
PYTHON_INCLUDE="$("$PYTHON" -c 'import sysconfig; print(sysconfig.get_path("include"))')"
[[ -f "$PYTHON_INCLUDE/Python.h" ]] || {
  echo "Missing Python development headers: $PYTHON_INCLUDE/Python.h" >&2; exit 1;
}
if $RESOLVE; then
  TEMP_LOCK="$(mktemp "${LOCK}.tmp.XXXXXX")"
  trap 'rm -f -- "$TEMP_LOCK"' EXIT
  uv pip compile "$INPUT" --python "$PYTHON" --generate-hashes --output-file "$TEMP_LOCK"
  mv -- "$TEMP_LOCK" "$LOCK"
  (cd "$SCRIPT_DIR" && sha256sum requirements.in) > "$LOCK.input.sha256"
  trap - EXIT
fi
[[ -s "$LOCK.input.sha256" ]] || { echo "Missing input fingerprint" >&2; exit 1; }
(cd "$SCRIPT_DIR" && sha256sum --check "${LOCK}.input.sha256")
echo "[3/5] Install locked packages (dedicated venv only)"
uv pip sync --python "$PYTHON" --require-hashes "$LOCK"
uv pip check --python "$PYTHON"
echo "[4/5] GPU, policy implementations, video libraries and CLI"
"$PYTHON" - <<'PY'
import importlib
import importlib.metadata
import torch
assert importlib.metadata.version('lerobot') == '0.6.1'
assert torch.cuda.is_available(), 'CUDA unavailable: check driver and torch build'
x = torch.ones((32, 32), device='cuda')
assert (x @ x).sum().item() == 32768
torch.cuda.synchronize()
print('torch:', torch.__version__, 'CUDA:', torch.version.cuda, 'GPU:', torch.cuda.get_device_name(0))
for module in (
    'lerobot.policies.act.modeling_act',
    'lerobot.policies.smolvla.modeling_smolvla',
    'lerobot.robots.so_follower',
    'lerobot.teleoperators.so_leader',
    'av', 'torchcodec',
):
    importlib.import_module(module)
    print('OK:', module)
PY
for cli in lerobot-find-port lerobot-find-cameras lerobot-setup-motors lerobot-calibrate lerobot-teleoperate lerobot-record lerobot-train lerobot-rollout; do
  [[ -x "$VENV/bin/$cli" ]] || { echo "Missing CLI: $cli" >&2; exit 1; }
done
# find-port ignores --help and starts an interactive USB discovery procedure.
# Discovery tools are checked for existence only; never invoke them in setup.
for cli in lerobot-setup-motors lerobot-calibrate lerobot-teleoperate lerobot-record lerobot-train lerobot-rollout; do
  "$VENV/bin/$cli" --help > "$REPORT/$cli-help.txt" 2>&1
done
echo "[5/5] Save verification evidence"
uv pip freeze --python "$PYTHON" > "$REPORT/packages.txt"
"$PYTHON" -VV > "$REPORT/python.txt"
sha256sum "$LOCK" > "$REPORT/lock.sha256"
cp -- "$LOCK" "$LOCK.input.sha256" "$REPORT/"
echo "PASS: software smoke checks only; recording, model training and robot safety remain unverified."
echo "Report: $REPORT"
printf 'Activate: source %q/bin/activate\n' "$VENV"
echo "Use hf auth login only if accessing the Hub; use the official endpoint for authenticated uploads."