#!/usr/bin/env bash
# HAGOP Robot Lab: optional RunPod/XGO-Duck training helper.
# This never provisions a GPU; invoke only inside a Pod you started yourself.
set -Eeuo pipefail

ROOT="${ROBOT_LAB_ROOT:-/workspace/robot-lab}"
REPO="$ROOT/xgoduck_rl"
SRC="https://github.com/LuwuDynamics/xgoduck_rl.git"

usage() {
  printf '%s\n' \
    "Usage: bash bootstrap.sh setup" \
    "       bash bootstrap.sh train [flat|rough|stand|pick|kick|roll] [128|256|512|1024|2048|4096]" \
    "       bash bootstrap.sh tensorboard" \
    "       bash bootstrap.sh status" \
    "The script NEVER creates a GPU Pod and NEVER exposes TensorBoard publicly." \
    "Training has a 55-minute process timeout; stopping the GPU Pod is a separate provider action."
}

check_gpu() {
  if ! command -v nvidia-smi >/dev/null 2>&1; then
    echo "ERROR: No NVIDIA driver/GPU visible. Use an NVIDIA GPU Pod." >&2; exit 2
  fi
  nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
}

setup() {
  check_gpu
  command -v git >/dev/null || { echo "git not found" >&2; exit 2; }
  command -v uv >/dev/null || { echo "uv not found. Install uv from https://docs.astral.sh/uv/getting-started/installation/" >&2; exit 2; }
  mkdir -p "$ROOT"
  if [[ -d "$REPO/.git" ]]; then
    echo "Existing checkout retained; no automatic pull over your work."
  elif [[ -e "$REPO" ]]; then
    echo "ERROR: $REPO exists but is not a git checkout." >&2; exit 2
  else
    git clone --depth 1 "$SRC" "$REPO"
  fi
  cd "$REPO"
  uv python install 3.12
  UV_HTTP_TIMEOUT=600 uv sync --python 3.12 --frozen
  printf 'Source SHA: '; git rev-parse HEAD
  echo "Ready. Next: bash bootstrap.sh train flat 256"
}

resolve_task() {
  case "$1" in
    flat) TASK="Mjlab-Velocity-Flat-XgoDuck" ;;
    rough) TASK="Mjlab-Velocity-Rough-XgoDuck" ;;
    stand) TASK="Mjlab-StandUp-Flat-XgoDuck" ;;
    pick) TASK="Mjlab-GroundPick-Flat-XgoDuck" ;;
    kick) TASK="Mjlab-BallKick-Flat-XgoDuck" ;;
    roll) TASK="Mjlab-Roulade-Flat-XgoDuck" ;;
    *) echo "ERROR: Unsupported task '$1'." >&2; usage; exit 2 ;;
  esac
}

train() {
  [[ -d "$REPO/.git" ]] || { echo "Run 'setup' first." >&2; exit 2; }
  check_gpu
  resolve_task "${1:-flat}"
  case "${2:-256}" in
    128|256|512|1024|2048|4096) ENVS="${2:-256}" ;;
    *) echo "ERROR: environments must be 128, 256, 512, 1024, 2048, or 4096." >&2; exit 2 ;;
  esac
  cd "$REPO"
  mkdir -p "$ROOT/session-logs"
  RUNSTAMP="$(date -u +%Y%m%dT%H%M%SZ)"
  LOG="$ROOT/session-logs/${RUNSTAMP}-${1:-flat}-${ENVS}.log"
  printf 'UTC start: %s\nTask: %s\nEnvironments: %s\nGit SHA: %s\n' \
    "$RUNSTAMP" "$TASK" "$ENVS" "$(git rev-parse HEAD)" | tee "$LOG"
  echo "Checkpoint directory: $REPO/logs/rsl_rl"
  command -v timeout >/dev/null || { echo "GNU timeout is required for the 55-minute pilot process cutoff." >&2; exit 2; }
  echo "Trial training process will be interrupted after 55 minutes (3300s)."
  echo "IMPORTANT: this does NOT stop the billable RunPod Pod; configure its separate 1h auto-stop."
  set +e
  UV_HTTP_TIMEOUT=600 timeout --signal=INT --kill-after=30s 3300s uv run train "$TASK" --env.scene.num-envs "$ENVS" 2>&1 | tee -a "$LOG"
  result=${PIPESTATUS[0]}
  set -e
  if [[ "$result" == 124 ]]; then
    echo "Robot Lab training process reached its 55-minute cutoff. Verify the GPU Pod auto-stop in RunPod."
  elif [[ "$result" != 0 ]]; then
    exit "$result"
  fi
}

tensorboard() {
  [[ -d "$REPO/.git" ]] || { echo "Run 'setup' first." >&2; exit 2; }
  cd "$REPO"
  echo "Local-only TensorBoard at 127.0.0.1:6006. Access via authenticated SSH tunnel."
  uv run --with tensorboard tensorboard --logdir logs/rsl_rl --host 127.0.0.1 --port 6006
}

case "${1:-help}" in
  setup) setup ;;
  train) train "${2:-flat}" "${3:-256}" ;;
  tensorboard) tensorboard ;;
  status) check_gpu; [[ ! -d "$REPO/.git" ]] || { git -C "$REPO" rev-parse --short HEAD; du -sh "$REPO/logs" 2>/dev/null || true; } ;;
  *) usage ;;
esac