#!/usr/bin/env bash
# Points a non-NixOS Debian machine's Nix install at nix-cache: adds it as
# a substituter (with cache.nixos.org kept as fallback) and, once the
# remote-builder private key is installed, as a distributed-build machine
# too.
#
# This is the non-NixOS equivalent of modules/nix-cache/client.nix +
# modules/nix-cache/remote-builder-client.nix -- those two only apply to
# hosts built from this flake. A plain Debian box with Nix installed
# (single- or multi-user install, nix-daemon running) has no NixOS module
# system to pick that config up, so this edits /etc/nix/nix.conf by hand
# instead. Run this ON the target Debian machine, as root.
#
# The values below mirror variables.nix / modules/nix-cache/client.nix in
# this repo -- update both if nix-cache is ever rebuilt with a new host
# key or the cache signing key is rotated (see docs/nix-cache.md).
#
# REMOTE_BUILDER_KEY defaults to this machine's own default root SSH
# identity (matches modules/nix-cache/remote-builder-client.nix's
# convention for real NixOS clients: authenticate as nixremote with the
# host's own default key, added individually to
# vars.remoteBuilderAuthorizedKeys, rather than a separately-named or
# shared keypair) -- generate one with
# `ssh-keygen -t ed25519 -N '' -f /root/.ssh/id_ed25519` if this machine
# doesn't have one yet, then add its .pub to vars.remoteBuilderAuthorizedKeys
# and rebuild nix-cache.
#
# Usage:
#   sudo ./configure-nix-cache-client.sh [--dry-run] [--no-remote-builder] [--no-restart]
#
# Env overrides (defaults match variables.nix):
#   NIX_CACHE_HOST, NIX_CACHE_HOST_KEY, REMOTE_BUILDER_USER, REMOTE_BUILDER_KEY

set -euo pipefail

: "${NIX_CACHE_HOST:=nix-cache}"
: "${NIX_CACHE_HOST_KEY:=ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIPeWgMsdaiz4axT/deFc1+0B5bN+GX/NOeW9bbQ0c/IT lxc-nix-cache}"
: "${REMOTE_BUILDER_USER:=nixremote}"
: "${REMOTE_BUILDER_KEY:=/root/.ssh/id_ed25519}"

CACHE_PUB_KEY="cache.local-1:usoWYanY3Kpq2+kDIS2nhWoLZiRxanmdysdzqCFBHW4="
FALLBACK_URL="https://cache.nixos.org/"
FALLBACK_PUB_KEY="cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY="

NIX_CONF="/etc/nix/nix.conf"
KNOWN_HOSTS="/etc/ssh/ssh_known_hosts"
MARKER_BEGIN="# BEGIN nix-cache client config (configure-nix-cache-client.sh)"
MARKER_END="# END nix-cache client config"

dry_run=0
with_remote_builder=1
restart_daemon=1

for arg in "$@"; do
  case "$arg" in
    --dry-run) dry_run=1 ;;
    --no-remote-builder) with_remote_builder=0 ;;
    --no-restart) restart_daemon=0 ;;
    -h|--help)
      sed -n '2,20p' "$0"
      exit 0
      ;;
    *)
      echo "ERROR: unknown argument: $arg" >&2
      exit 1
      ;;
  esac
done

if [[ "$dry_run" -eq 0 && "$EUID" -ne 0 ]]; then
  echo "ERROR: must run as root (writes $NIX_CONF and, unless --no-remote-builder, $KNOWN_HOSTS)." >&2
  exit 1
fi

if ! command -v nix >/dev/null 2>&1; then
  echo "ERROR: no 'nix' binary on PATH -- install the Nix package manager first." >&2
  exit 1
fi

if [[ ! -f "$NIX_CONF" ]]; then
  echo "ERROR: $NIX_CONF not found -- expected an existing multi-user Nix install." >&2
  exit 1
fi

builder_line=""
if [[ "$with_remote_builder" -eq 1 ]]; then
  if [[ -f "$REMOTE_BUILDER_KEY" ]]; then
    case "$(uname -m)" in
      x86_64) nix_system="x86_64-linux" ;;
      aarch64) nix_system="aarch64-linux" ;;
      *)
        echo "WARNING: unrecognized architecture '$(uname -m)' -- skipping remote builder, keeping substituter config." >&2
        with_remote_builder=0
        ;;
    esac
    if [[ "$with_remote_builder" -eq 1 ]]; then
      builder_line="builders = ssh://${REMOTE_BUILDER_USER}@${NIX_CACHE_HOST} ${nix_system} ${REMOTE_BUILDER_KEY} 4 2 big-parallel,kvm,nixos-test,benchmark"
    fi
  else
    echo "WARNING: $REMOTE_BUILDER_KEY not found -- skipping remote builder config (substituter still configured)." >&2
    echo "  See docs/nix-cache.md 'Remote builder SSH keys' for how to install it, then re-run this script." >&2
    with_remote_builder=0
  fi
fi

block="$(cat <<EOF
$MARKER_BEGIN
extra-substituters = http://${NIX_CACHE_HOST} ${FALLBACK_URL}
extra-trusted-public-keys = ${CACHE_PUB_KEY} ${FALLBACK_PUB_KEY}
EOF
)"
if [[ "$with_remote_builder" -eq 1 ]]; then
  block="${block}
builders-use-substitutes = true
${builder_line}"
fi
block="${block}
$MARKER_END"

echo "== nix.conf block to install =="
echo "$block"
echo "================================"

if [[ "$dry_run" -eq 1 ]]; then
  echo "(--dry-run: not writing $NIX_CONF)"
else
  tmp_conf="$(mktemp)"
  trap 'rm -f "$tmp_conf"' EXIT

  if grep -qF "$MARKER_BEGIN" "$NIX_CONF"; then
    awk -v begin="$MARKER_BEGIN" -v end="$MARKER_END" -v block="$block" '
      $0 == begin { print block; skip = 1; next }
      $0 == end   { skip = 0; next }
      skip        { next }
      { print }
    ' "$NIX_CONF" > "$tmp_conf"
  else
    cp "$NIX_CONF" "$tmp_conf"
    printf '\n%s\n' "$block" >> "$tmp_conf"
  fi

  cp "$NIX_CONF" "${NIX_CONF}.bak.$(date +%Y%m%d%H%M%S)"
  install -m 0644 "$tmp_conf" "$NIX_CONF"
  echo "Updated $NIX_CONF (backup saved alongside it)."
fi

if [[ "$with_remote_builder" -eq 1 ]]; then
  known_hosts_line="${NIX_CACHE_HOST} ${NIX_CACHE_HOST_KEY}"
  if [[ "$dry_run" -eq 1 ]]; then
    echo "(--dry-run: would ensure this line is present in $KNOWN_HOSTS)"
    echo "  $known_hosts_line"
  else
    mkdir -p "$(dirname "$KNOWN_HOSTS")"
    touch "$KNOWN_HOSTS"
    if ! grep -qF "$known_hosts_line" "$KNOWN_HOSTS" 2>/dev/null; then
      echo "$known_hosts_line" >> "$KNOWN_HOSTS"
      echo "Added nix-cache's SSH host key to $KNOWN_HOSTS."
    fi
  fi
fi

if [[ "$dry_run" -eq 0 && "$restart_daemon" -eq 1 ]]; then
  if command -v systemctl >/dev/null 2>&1 && systemctl is-active --quiet nix-daemon 2>/dev/null; then
    systemctl restart nix-daemon
    echo "Restarted nix-daemon to pick up the new config."
  else
    echo "nix-daemon not managed by systemd (or not running) -- restart it manually to pick up the new config."
  fi
fi

cat <<EOF

Done. Verify with:
  curl http://${NIX_CACHE_HOST}/nix-cache-info
  nix show-config | grep -E 'substituters|trusted-public-keys|builders'
EOF
if [[ "$with_remote_builder" -eq 1 ]]; then
  cat <<EOF
  ssh -i ${REMOTE_BUILDER_KEY} ${REMOTE_BUILDER_USER}@${NIX_CACHE_HOST} nix-store --version
  nix build nixpkgs#hello -L
EOF
fi