#!/usr/bin/env nix-shell
#!nix-shell -i bash -p nettools

SELF=$(hostname)

usage() {
  echo "deploy: deploy a nix config to a remote machine, possibly activating it"
  echo ""
  echo "usage: deploy [options] [host] [host2 ...]"
  echo "options:"
  echo "- --build: only build; don't copy or deploy"
  echo "- --copy: only build + copy files, nothing more"
  echo "- --switch  (default)"
  echo "- --test: switch to the new configuration, but do not make it bootable"
  echo "- --dry-run: show what would be done without actually doing it"
  echo "- --pre: alias for --action copy --variant all all"
  echo "- --reboot: reboot the target machine after deploying (if deployed with no errors)"
  echo "- --reboot-force: reboot the target machine after deploying (even on error)"
  echo "- --variant light|min|''|all  (default: '')"
  echo "- --wireguard always|never|opportunistic: deploy over wireguard"
  echo "- --ip <address>: deploy to the specific IP address"
  echo "- --deriv /nix/store/...: prebuilt store path (or .drv to realize) to deploy instead of (re-)building the default target"
  echo ""
  echo "common idioms:"
  echo "- deploy all: deploy all hosts, sequentially"
  echo "- deploy --pre: build and copy all hosts"
  echo "- deploy desko lappy: build and deploy just those hosts"
  echo "- deploy: deploy the local host"
  exit 1
}

info() {
  echo "[deploy]" "$@"
}

action=switch
hosts=()
ip=
defaultHost="$SELF"
variants=()
defaultVariant=
# by default, don't ship builds to servo. i'd guess this can be overriden by passing --builders servo
nixArgs=(--builders "")
doReboot=
doRebootForce=
dryRun=
wireguard=opportunistic
storePath=
addHost() {
  if [ "$1" = all ]; then
    # order matters:
    hosts+=(moby lappy desko servo)
  else
    hosts+=("$1")
  fi
}
addVariant() {
  if [ "$1" = all ]; then
    variants+=("-min" "-light" "" "-next-min" "-next-light" "-next")
  elif [ -n "$1" ]; then
    variants+=("-$1")
  else
    # "full" variant
    variants+=("")
  fi
}
parseArgs() {
  while [ "$#" -ne 0 ]; do
    local arg=$1
    shift
    case "$arg" in
      (--build|--copy|--switch|--test)
        action=${arg/--/}
        ;;
      (--deriv)
        storePath="$1"
        shift
        ;;
      (--dry-run)
        dryRun=1
        ;;
      (--help)
        usage
        ;;
      (--ip)
        ip="$1"
        shift
        ;;
      (--pre)
        action=copy
        defaultVariant=all
        defaultHost=all
        ;;
      (--reboot)
        doReboot=1
        ;;
      (--reboot-force)
        doReboot=1
        doRebootForce=1
        ;;
      (--variant)
        addVariant "$1"
        shift
        ;;
      (--wireguard)
        wireguard="$1"
        shift
        ;;
      (all|crappy|desko|lappy|moby|servo)
        addHost "$arg"
        ;;
      (*)
        nixArgs+=("$arg")
        ;;
    esac
  done

  if [ "${#hosts[@]}" -eq 0 ] && [ -n "$defaultHost" ]; then
    addHost "$defaultHost"
  fi
  if [ "${#variants[@]}" -eq 0 ]; then
    addVariant "$defaultVariant"
  fi
}

destructive() {
  if [ -z "$dryRun" ]; then
    if [ -n "$ECHO_CMD" ]; then
      echo "$@"
    fi
    "$@"
  else
    echo "dry-run: $@"
  fi
}

# return "$1" or "$1-hn", based on if wireguard was requested or not
resolveHost() {
  if [ -n "$ip" ]; then
    echo "$ip"
  else
    case "$wireguard-$1" in
      (opportunistic-moby)
        echo "$1-hn"
        ;;
      (opportunistic-*)
        echo "$1"
        ;;
      (never-*)
        echo "$1"
        ;;
      (always-*)
        echo "$1-hn"
        ;;
      (*)
        echo "$1-hn"
        ;;
    esac
  fi
}

# return the number of seconds to allot to `nix copy` when copying the given variant.
# this is done to avoid stalled copies from blocking the entire process, while hopefully not breaking the copies that are actually important
timeoutFor() {
  case $1 in
    (-min|-light|-next)
      echo 3600
      ;;
    (-next-min|-next-light)
      echo 1800
      ;;
    (*)
      # this catches the normal variant
      echo 14400
      ;;
  esac
}

runOnTarget() {
  local host="$1"
  shift
  # run the command ($@) on the machine we're deploying to.
  # if that's a remote machine, then do it via ssh, else local shell.
  if [ -n "$host" ] && [ "$host" != "$SELF" ]; then
    info "running on remote ($host):" "$@"
    ssh "$host" "$@"
  else
    info "running locally ($SELF):" "$@"
    "$@"
  fi
}

# deployOneHost $host $variant
deployOneHost() {
  local host="$1"
  local variant="$2"

  local timeout=$(timeoutFor "$variant")

  # storePath is allowed to be either a realized derivation,
  # or the path to a .drv file itself
  local myStorePath="$storePath"
  if [ -z "$myStorePath" ]; then
    # `nix-build -A foo` evals and then realizes foo, but it never unloads the memory used to eval foo.
    # my exprs are heavyweight, we need that memory for building, so do the evals separately from the realizations:
    info "evaluating $host$variant..."
    myStorePath=$(nix eval --raw -f . "hosts.$host$variant.toplevel.drvPath")
  fi

  if [[ "$myStorePath" == *.drv ]]; then
    info "building $host$variant ($drvPath)"
    myStorePath=$(destructive nix-store --realize "$myStorePath" "${nixArgs[@]}")
    if [ -z "$myStorePath" ]; then
      return 1
    fi
    info "built $host$variant -> $myStorePath"
  fi

  # mimic `nixos-rebuild --target-host`, in effect:
  # - nix-copy-closure ...
  # - nix-env --set ...
  # - switch-to-configuration <boot|dry-activate|switch|test|>
  # avoid the actual `nixos-rebuild` for a few reasons:
  # - fewer nix evals
  # - more introspectability and debuggability
  # - sandbox friendliness (especially: `git` doesn't have to be run as root)

  local netHost=$(resolveHost "$host")

  case "$action" in
    (copy|switch|test)
      if [ -n "$host" ] && [ "$host" != "$SELF" ]; then
        if [ -e /run/secrets/nix_signing_key ]; then
          info "signing store paths ..."
          destructive sudo nix store sign -r -k /run/secrets/nix_signing_key "$myStorePath"
        else
          info "not signing store paths: /run/secrets/nix_signing_key does not exist"
        fi
        # add more `-v` for more verbosity (up to 5).
        # builders-use-substitutes false: optimizes so that the remote machine doesn't try to get paths from its substituters.
        #   we already have all paths here, and the remote substitution is slow to check and SERIOUSLY flaky on moby in particular.
        ECHO_CMD=1 destructive timeout "$timeout" nix copy -vv --option builders-use-substitutes false --to "ssh-ng://$netHost" "$myStorePath" || return 1
      fi
      ;;
  esac

  case "$action" in
    (switch|test)
      info "activating profile... "
      destructive runOnTarget "$netHost" sudo nix-env -p /nix/var/nix/profiles/system --set "$myStorePath" || return 1
      destructive runOnTarget "$netHost" sudo "$myStorePath/bin/switch-to-configuration" "$action"
      local rc=$?
      if [[ -n "$doReboot" && ("$rc" -eq 0 || -n "$doRebootForce") ]]; then
        info "rebooting $host"
        destructive runOnTarget "$netHost" sane-reboot "$host" || return 1
      fi
      return $rc
      ;;
  esac
}

failedDeploys=()
deployHosts() {
  local hosts=("$@")
  for v in "${variants[@]}"; do
    for h in "${hosts[@]}"; do
      deployOneHost "$h" "$v" || \
        failedDeploys+=("$h$v")
    done
  done
}

parseArgs "$@"

# i care e.g. that full moby is deployed before crappy:
earlyHosts=()
lateHosts=()
for host in "${hosts[@]}"; do
  case $host in
    (crappy)
      lateHosts+=("$host")
      ;;
    (*)
      earlyHosts+=("$host")
      ;;
  esac
done

deployHosts "${earlyHosts[@]}"
deployHosts "${lateHosts[@]}"

if [ "${#failedDeploys[@]}" -ne 0 ]; then
  echo "FAILED DEPLOYMENT:"
  for d in "${failedDeploys[@]}"; do
    echo "- $d"
  done
  exit 1
else
  echo "SUCCESS"
fi