#!/bin/bash
# wxch installer. Do not run this directly: use the one-liner, which asks for sudo once:
#
#   curl -fsSL @@BASE@@/get | WXCH_TOKEN=<agent token> sh
#
# It installs the agent and the three model runtimes (Ollama, LM Studio, vLLM), downloads the smallest
# model of each, and runs a self-test that loads each model and asks it a question.
# Every step can be re-run safely; a re-run repairs what is missing and skips what is done.
#
# Options (use `sh -s -- <option>` after the one-liner):
#   --reinstall      remove the previous install first (models are kept), then install again
#   --dry-run        show the plan and check the scripts, change nothing (no sudo needed)
#   --no-selftest    skip the final self-test
#   -h, --help       this text
#
# Environment (all optional except the token on a first install):
#   WXCH_TOKEN          agent token from `wxch-server -add-agent NAME` (kept in /etc/wxch-agent.env afterwards)
#   WXCH_USER           user that owns the agent (default: the user who ran sudo)
#   WXCH_SKIP           steps to skip, comma separated: models-disk,agent,ollama,lmstudio,vllm,models,modules,selftest
#   WXCH_MODULES        optional AI modules to install (not installed by default): comma separated from image.generate,
#                       image.edit, video.generate, video.edit, audio.transcribe, or "all" (about 110 GB for all five)
#   WXCH_ONLY           run only these steps
#   WXCH_MODELS_DIR     where models live (default /mnt/models)
#   (external disks are added later, one at a time, with: curl -fsSL <server>/get | sh -s -- add-disk)
#   WXCH_MODEL_OLLAMA / WXCH_MODEL_LMSTUDIO / WXCH_MODEL_VLLM   the starter models to download
#   WXCH_TRUST_REMOTE_CODE  vLLM models (org/name, comma separated) that may run the Python code shipped in their repository;
#                       kept in /etc/wxch/runner.json across re-runs, only the owner can add to it
#   WXCH_VLLM_IMAGE     vLLM container image
#   WXCH_CONTINUE=1     keep going after a failed step

set -uo pipefail
BASE="${WXCH_BASE:?run this through the /get one-liner}"
DRY=0
REINSTALL=0
for a in "$@"; do
  case "$a" in
    --dry-run) DRY=1 ;;
    --reinstall) REINSTALL=1 ;;
    --no-selftest) WXCH_SKIP="${WXCH_SKIP:-},selftest" ;;
    -h|--help) sed -n '2,/^set -uo/p' "$0" | sed '$d' | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) echo "unknown option: $a (see --help)" >&2; exit 2 ;;
  esac
done
export WXCH_BASE="$BASE" WXCH_SKIP="${WXCH_SKIP:-}"

if [ "$DRY" = 0 ]; then
  [ "$(id -u)" = 0 ] || { echo "must run as root: use the /get one-liner (it asks for sudo once)" >&2; exit 1; }
  mkdir -p /var/log
  exec > >(tee -a /var/log/wxch-setup.log) 2>&1
fi
echo "== wxch setup  $(date '+%F %T')  server: $BASE  $([ "$DRY" = 1 ] && echo '(dry run: nothing is changed)')"

ALL_STEPS="preflight models-disk agent ollama lmstudio vllm models modules selftest"
enabled() {
  if [ -n "${WXCH_ONLY:-}" ]; then case ",$WXCH_ONLY," in *",$1,"*) return 0 ;; *) return 1 ;; esac; fi
  case ",${WXCH_SKIP:-}," in *",$1,"*) return 1 ;; esac
  [ -e "/run/wxch-auto-skip-$1" ] && return 1
  return 0
}

WORK="$(mktemp -d /tmp/wxch-steps.XXXXXX)"
trap 'rm -rf "$WORK"' EXIT

# fetch every script first: a broken download must stop us before anything has been changed
for s in lib $ALL_STEPS; do
  curl -fsSL --connect-timeout 20 "$BASE/steps/$s.sh" -o "$WORK/$s.sh" || { echo "cannot download step '$s' from $BASE" >&2; exit 1; }
  bash -n "$WORK/$s.sh" || { echo "step '$s' has a syntax error" >&2; exit 1; }
done
if [ "$REINSTALL" = 1 ]; then
  curl -fsSL "$BASE/uninstall.sh" -o "$WORK/uninstall.sh" || { echo "cannot download uninstall.sh" >&2; exit 1; }
  bash -n "$WORK/uninstall.sh" || exit 1
fi

if [ "$DRY" = 1 ]; then
  echo "plan:"
  [ "$REINSTALL" = 1 ] && echo "  0. uninstall the previous install (models are kept)"
  i=1
  for s in $ALL_STEPS; do
    if enabled "$s"; then echo "  $i. $s"; else echo "  $i. $s   (skipped)"; fi
    i=$((i+1))
  done
  echo "all scripts downloaded and syntax-checked; nothing was changed."
  exit 0
fi

RESULTS=""
FAILED=0
if [ "$REINSTALL" = 1 ]; then
  echo; echo "== step: uninstall (previous install)"
  bash "$WORK/uninstall.sh" -y || { echo "uninstall failed" >&2; exit 1; }   # asking for --reinstall is the confirmation; models are kept
fi

for s in $ALL_STEPS; do
  if ! enabled "$s"; then RESULTS="$RESULTS\n  skipped   $s"; continue; fi
  echo; echo "== step: $s"
  if bash "$WORK/$s.sh"; then
    RESULTS="$RESULTS\n  ok        $s"
  else
    RESULTS="$RESULTS\n  FAILED    $s"
    FAILED=1
    if [ "${WXCH_CONTINUE:-0}" != 1 ]; then
      echo; echo "stopped at step '$s'. Fix the problem and run the same command again: finished steps are skipped."
      printf "\nsummary:$RESULTS\n"
      exit 1
    fi
  fi
done

echo
printf "summary:$RESULTS\n"
echo
[ -x /opt/wxch-agent/wxch-agent ] && sudo -u "${WXCH_USER:-${SUDO_USER:-root}}" -H /opt/wxch-agent/wxch-agent status
echo
if [ "$FAILED" = 0 ]; then
  echo "DONE. Open $BASE and pick this machine; the Models page lists what is installed."
else
  echo "FINISHED WITH ERRORS (see above). The log is in /var/log/wxch-setup.log"
  exit 1
fi