#!/usr/bin/env bash
set -eufo pipefail
set -x

showUsage() {
  cat <<USAGE
Usage: $0 [options] ssh-host

Options:

* -f, --flake flake
  set the flake to install the system from
* --arg name value
  pass value to nix-build. can be used to set disk-names for example
* --argstr name value
  pass value to nix-build as string
* --kexec url
  use another kexec tarball to bootstrap NixOS
USAGE
}

abort() {
  echo "aborted: $*" >&2
  exit 1
}

nix_args=()
kexec_url=https://github.com/nix-community/nixos-images/releases/download/nixos-22.05/nixos-kexec-installer-x86_64-linux.tar.gz

while [[ $# -gt 0 ]]; do
  case "$1" in
    -f | --flake)
      flake=$2
      shift
      ;;
    --argstr | --arg)
      nix_args+=("$1" "$2" "$3")
      shift
      shift
      ;;
    --help)
      showUsage
      exit 0
      ;;
    --kexec)
      kexec_url=$2
      shift
      ;;
    --no-ssh-copy-id)
      no_ssh_copy=y
      ;;
    *)
      if [ -z ${ssh_connection+x} ]; then
        ssh_connection=$1
      else
        showUsage
        exit 1
      fi
      ;;
  esac
  shift
done

# ssh wrapper
timeout_ssh_() {
  timeout 10 ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no "$ssh_connection" "$@"
}
ssh_() {
  ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no "$ssh_connection" "$@"
}
nixCopy() {
  NIX_SSHOPTS='-o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no' nix copy --extra-experimental-features nix-command "$@"
}
nix_build() {
  nix \
    --experimental-features flakes build \
    --extra-experimental-features nix-command \
    --no-write-lock-file \
    --print-out-paths \
    "$@"
}

nixos_system=""
# parse flake nixos-install style syntax, get the system attr
if [[ ! -z "${flake+x}" ]]; then
  if [[ $flake =~ ^(.*)\#([^\#\"]*)$ ]]; then
   flake="${BASH_REMATCH[1]}"
   flakeAttr="${BASH_REMATCH[2]}"
  fi
  if [[ -z "$flakeAttr" ]]; then
    echo "Please specify the name of the NixOS configuration to be installed, as a URI fragment in the flake-uri."
    echo "For example, to use the output nixosConfigurations.foo from the flake.nix, append \"#foo\" to the flake-uri."
    exit 1
  fi
  disko_script=$(nix_build "${flake}#nixosConfigurations.${flakeAttr}.config.system.build.disko")
  nixos_system=$(nix_build "${flake}#nixosConfigurations.${flakeAttr}.config.system.build.toplevel")
else
  abort "flake must be set"
fi


# wait for machine to become reachable (possibly forever)
if [ ${no_ssh_copy-n} != "y" ]; then
  until ssh-copy-id "$ssh_connection"; do sleep 5; done
else
  until ssh_ -o ConnectTimeout=10 -- exit 0; do sleep 5; done
fi

# first check if the remote system is kexec booted
if $(ssh_ -- test -e /etc/is_kexec); then
  is_kexec=y
fi

if [ ${is_kexec-n} != "y" ]; then
  # TODO we probably need an architecture detection here
  ssh_ << SSH
set -efux
os=\$(uname)
if [[ "\$os" != "Linux" ]]; then
  echo "This script requires Linux as the operating system, but got \${os}" >&2
  exit 1
fi
if ! command -v tar >/dev/null 2>&1; then
  echo "no tar command found, but required to unpack kexec tarball" >&2
  exit 1
fi

rm -rf /root/kexec
mkdir -p /root/kexec
SSH

if [[ -e "$kexec_url" ]]; then
  cat "$kexec_url" | ssh_ 'tar -C /root/kexec -xvzf-'
else
  ssh_ << SSH
fetch(){
  if command -v curl >/dev/null 2>&1; then
    curl --fail -Ss -L "\$1"
  elif command -v wget >/dev/null 2>&1; then
    wget "\$1" -O-
  else
    echo "no downloader (curl or wget) found, bailing out"
    exit 1
  fi
}
fetch "$kexec_url" | tar -C /root/kexec -xvzf-
SSH
fi

ssh_ << SSH
export TMPDIR=/root/kexec
setsid /root/kexec/kexec/run
SSH

  # wait for machine to become unreachable
  while timeout_ssh_ -- exit 0; do sleep 1; done

  # watiting for machine to become available again
  until ssh_ -o ConnectTimeout=10 -- exit 0; do sleep 5; done
fi

nixCopy --to "ssh://$ssh_connection" "$disko_script"
ssh_ $disko_script

nixCopy --to "ssh://$ssh_connection?remote-store=local?root=/mnt" "$nixos_system"
ssh_ << SSH
set -efu
nixos-install --no-root-passwd --no-channel-copy --system "$nixos_system"
reboot
SSH