#!/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 has no
# NixOS module system to pick that config up, so this edits nix.conf by hand.
#
# Two modes depending on who runs it:
#
#   root (multi-user / daemon install):
#     Writes /etc/nix/nix.conf, /etc/ssh/ssh_known_hosts, restarts nix-daemon.
#     Requires /etc/nix/nix.conf to already exist (i.e. nix-daemon is set up).
#     Run as: sudo ./configure-nix-cache-client.sh [options]
#
#   non-root (single-user install):
#     Writes ~/.config/nix/nix.conf, ~/.ssh/known_hosts. No daemon to restart.
#     Run as: ./configure-nix-cache-client.sh [options]
#
# 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 the running user's default SSH identity
# (root: /root/.ssh/id_ed25519, other user: ~/.ssh/id_ed25519). That key
# must be listed in vars.remoteBuilderAuthorizedKeys in this repo and
# nix-cache rebuilt before remote building works.
#
# Usage:
#   ./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 AAAAC3NzaC1lZDI1NTE5AAAAICuHUxGNH6ei3BZD+EfZs3l4X8uJNcjQiOsM/G4yo4O/ lxc-nix-cache}"
: "${REMOTE_BUILDER_USER:=nixremote}"

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

MARKER_BEGIN="# BEGIN nix-cache client config (configure-nix-cache-client.sh)"
MARKER_END="# END nix-cache client config"

# Mode: root uses system-wide paths and restarts the daemon; non-root uses
# user-level paths and has no daemon to restart.
if [[ "$EUID" -eq 0 ]]; then
  install_mode="multi"
  NIX_CONF="/etc/nix/nix.conf"
  KNOWN_HOSTS="/etc/ssh/ssh_known_hosts"
  : "${REMOTE_BUILDER_KEY:=/root/.ssh/id_ed25519}"
else
  install_mode="single"
  NIX_CONF="${XDG_CONFIG_HOME:-$HOME/.config}/nix/nix.conf"
  KNOWN_HOSTS="$HOME/.ssh/known_hosts"
  : "${REMOTE_BUILDER_KEY:=$HOME/.ssh/id_ed25519}"
fi

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,37p' "$0"
      exit 0
      ;;
    *)
      echo "ERROR: unknown argument: $arg" >&2
      exit 1
      ;;
  esac
done

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 [[ "$install_mode" == "multi" && ! -f "$NIX_CONF" ]]; then
  echo "ERROR: $NIX_CONF not found -- expected an existing multi-user Nix install." >&2
  exit 1
fi

# Single-user: create the config file if it doesn't exist yet.
if [[ "$install_mode" == "single" && "$dry_run" -eq 0 ]]; then
  mkdir -p "$(dirname "$NIX_CONF")"
  [[ -f "$NIX_CONF" ]] || touch "$NIX_CONF"
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 ($NIX_CONF) =="
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

# Only restart the daemon for multi-user installs -- single-user has no daemon.
if [[ "$dry_run" -eq 0 && "$restart_daemon" -eq 1 && "$install_mode" == "multi" ]]; 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