setup-host.sh

#!/usr/bin/env bash
# setup-host.sh — Avalanche NixOS flake host setup
# Usage: sudo ./setup-host.sh
set -euo pipefail

# ── Colours ───────────────────────────────────────────────────────────────────
BOLD=$'\033[1m'; RED=$'\033[0;31m'; GREEN=$'\033[0;32m'
YELLOW=$'\033[1;33m'; BLUE=$'\033[0;34m'; RESET=$'\033[0m'
info() { echo -e "  ${BLUE}→${RESET} $*"; }
ok()   { echo -e "  ${GREEN}✓${RESET} $*"; }
warn() { echo -e "  ${YELLOW}!${RESET} $*"; }
die()  { echo -e "\n  ${RED}✗${RESET} $*\n" >&2; exit 1; }
ask()  { read -rp "  ${BOLD}${1}${RESET} " "${2}"; }

# ── Sanity checks ─────────────────────────────────────────────────────────────
[[ $EUID -ne 0 ]]                    && die "Must run as root: sudo ./setup-host.sh"
command -v git     >/dev/null 2>&1   || die "git not found — run: nix-shell -p git"
command -v python3 >/dev/null 2>&1   || die "python3 not found"

FLAKE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
[[ -f "$FLAKE_DIR/flake.nix" ]]      || die "flake.nix not found in $FLAKE_DIR"
[[ -d "$FLAKE_DIR/.git" ]]           || die "Not inside a git repo — git add will not work"

HOSTS_DIR="$FLAKE_DIR/hosts"
HM_DIR="$FLAKE_DIR/home/hosts"
cd "$FLAKE_DIR"

# ── Parse hosts already declared in flake.nix ────────────────────────────────
mapfile -t FLAKE_HOSTS < <(
  grep -E '^\s+[a-z][a-z0-9-]* = mkHost' flake.nix \
  | sed -E 's/\s+([a-z][a-z0-9-]*) = mkHost.*/\1/'
)
[[ ${#FLAKE_HOSTS[@]} -eq 0 ]] && die "Could not parse any mkHost entries from flake.nix"

# ── Module picker ─────────────────────────────────────────────────────────────
# Discovers all subdirs of modules/ except 'common' (always included).
# Populates global SELECTED_MODULES array.
SELECTED_MODULES=()
pick_modules() {
  mapfile -t OPTIONAL_MODULES < <(
    find "$FLAKE_DIR/modules" -mindepth 1 -maxdepth 1 -type d \
      ! -name common -printf '%f\n' | sort
  )

  if [[ ${#OPTIONAL_MODULES[@]} -eq 0 ]]; then
    return
  fi

  echo ""
  echo "  Module selection — common is always included, toggle the rest:"
  echo "  Press a number to toggle, 0 to confirm."
  echo ""

  local -a sel=()
  for _ in "${OPTIONAL_MODULES[@]}"; do sel+=(false); done

  while true; do
    for i in "${!OPTIONAL_MODULES[@]}"; do
      local mark="○"
      [[ "${sel[$i]}" == "true" ]] && mark="${GREEN}●${RESET}"
      echo -e "    $((i+1))) ${mark} ${OPTIONAL_MODULES[$i]}"
    done
    echo    "    0) Done"
    read -rp "  Toggle [0-${#OPTIONAL_MODULES[@]}]: " PICK
    # Move cursor up to redraw the list cleanly
    local lines=$(( ${#OPTIONAL_MODULES[@]} + 1 ))
    for (( _j=0; _j<lines; _j++ )); do echo -ne "\033[1A\033[2K"; done
    echo -ne "\033[1A\033[2K"  # the prompt line

    if [[ "$PICK" == "0" ]]; then
      # Print final state before exiting
      for i in "${!OPTIONAL_MODULES[@]}"; do
        local mark="○"
        [[ "${sel[$i]}" == "true" ]] && mark="${GREEN}●${RESET}"
        echo -e "    $((i+1))) ${mark} ${OPTIONAL_MODULES[$i]}"
      done
      break
    elif [[ "$PICK" =~ ^[0-9]+$ ]] && (( PICK >= 1 && PICK <= ${#OPTIONAL_MODULES[@]} )); then
      local idx=$(( PICK - 1 ))
      [[ "${sel[$idx]}" == "true" ]] && sel[$idx]=false || sel[$idx]=true
    else
      warn "Enter a number between 0 and ${#OPTIONAL_MODULES[@]}."
    fi
  done

  for i in "${!OPTIONAL_MODULES[@]}"; do
    [[ "${sel[$i]}" == "true" ]] && SELECTED_MODULES+=("${OPTIONAL_MODULES[$i]}")
  done
}

# ── Template: hosts/<hostname>/default.nix ───────────────────────────────────
write_host_nix() {
  local hostname="$1"

  local extra_imports=""
  for mod in "${SELECTED_MODULES[@]}"; do
    extra_imports+="    ../../modules/${mod}"$'\n'
  done

  # Write the file in parts to avoid heredoc quoting conflicts with Nix syntax
  {
    echo "# hosts/${hostname}/default.nix — generated by setup-host.sh"
    echo "{ config, pkgs, inputs, ... }: {"
    echo "  imports = ["
    echo "    ./hardware-configuration.nix"
    echo "    ../../modules/common"
    printf '%s' "$extra_imports"
    echo "    inputs.noctalia.nixosModules.default"
    echo "  ];"
    echo ""
    echo "  networking.hostName = \"${hostname}\";"
    echo ""
    echo ""
    echo "  boot.loader.systemd-boot.enable = true;"
    echo "  boot.loader.efi.canTouchEfiVariables = true;"
    echo "  boot.loader.grub.enable = false;"
    echo "  services.printing.enable = true;"
    echo ""
    echo ""
    echo "  system.stateVersion = \"25.11\";"
    echo "}"
  } > "$HOSTS_DIR/${hostname}/default.nix"

  ok "hosts/${hostname}/default.nix written."
}

# ── Template: home/hosts/<hostname>.nix ──────────────────────────────────────
write_home_nix() {
  local hostname="$1"
  cat > "$HM_DIR/$hostname.nix" << NIXEOF
# home/hosts/$hostname.nix — generated by setup-host.sh
{ config, pkgs, ... }: {
  imports = [ ../common/core.nix ];

  # $hostname-specific packages
  home.packages = with pkgs; [

  ];

  home.file.".config/niri/config.kdl".source = ./niri.kdl;
}
NIXEOF
  ok "home/hosts/$hostname.nix written."
}

# ── Patch flake.nix: insert new mkHost entry ──────────────────────────────────
add_to_flake() {
  local hostname="$1"
  python3 - "$hostname" << 'PYEOF'
import sys

hostname = sys.argv[1]
with open("flake.nix") as f:
    lines = f.readlines()

if any("= mkHost" in l and hostname in l for l in lines):
    print(f"  ! {hostname} already in flake.nix — skipping")
    sys.exit(0)

indices = [i for i, l in enumerate(lines) if "= mkHost" in l]
if not indices:
    print("  ✗ Could not locate any mkHost entries in flake.nix", file=sys.stderr)
    sys.exit(1)

new_line = f'        {hostname} = mkHost {{ hostname = "{hostname}"; }};\n'
lines.insert(max(indices) + 1, new_line)

with open("flake.nix", "w") as f:
    f.writelines(lines)
print(f"  ✓ {hostname} added to flake.nix")
PYEOF
}

# ═════════════════════════════════════════════════════════════════════════════
echo ""
echo -e "${BOLD}❄  Avalanche — Host Setup${RESET}"
echo "────────────────────────────────────────────"
echo "  1) Assign this machine to an existing flake host"
echo "  2) Register this machine as a brand-new host"
echo "────────────────────────────────────────────"
ask "Option [1/2]:" OPTION

# ── Option 1 — existing flake host ───────────────────────────────────────────
if [[ "$OPTION" == "1" ]]; then
  echo ""
  info "Hosts declared in flake.nix:"
  PS3="  Select host: "
  select HOSTNAME in "${FLAKE_HOSTS[@]}"; do
    [[ -n "$HOSTNAME" ]] && break
    warn "Invalid selection — try again."
  done

  HOST_DIR="$HOSTS_DIR/$HOSTNAME"
  HM_FILE="$HM_DIR/$HOSTNAME.nix"

  if [[ ! -d "$HOST_DIR" ]]; then
    warn "'$HOSTNAME' has no host directory yet — scaffolding now..."
    pick_modules
    mkdir -p "$HOST_DIR"
    info "Generating hardware-configuration.nix..."
    nixos-generate-config --show-hardware-config > "$HOST_DIR/hardware-configuration.nix"
    ok "hardware-configuration.nix generated."
    write_host_nix "$HOSTNAME"
  else
    ask "Regenerate hardware-configuration.nix for this hardware? [y/N]:" REGEN
    if [[ "$REGEN" =~ ^[Yy]$ ]]; then
      info "Regenerating hardware-configuration.nix..."
      nixos-generate-config --show-hardware-config > "$HOST_DIR/hardware-configuration.nix"
      ok "hardware-configuration.nix updated."
    fi
  fi

  if [[ ! -f "$HM_FILE" ]]; then
    warn "home/hosts/$HOSTNAME.nix not found — creating stub..."
    write_home_nix "$HOSTNAME"
  fi

# ── Option 2 — brand-new host ─────────────────────────────────────────────────
elif [[ "$OPTION" == "2" ]]; then
  ask "Hostname (lowercase, hyphens allowed):" HOSTNAME
  [[ -z "$HOSTNAME" ]]             && die "Hostname cannot be empty."
  [[ "$HOSTNAME" =~ [^a-z0-9-] ]] && die "Hostname must be lowercase alphanumeric + hyphens only."

  for h in "${FLAKE_HOSTS[@]}"; do
    [[ "$h" == "$HOSTNAME" ]] && die "'$HOSTNAME' is already declared in flake.nix — use option 1."
  done
  [[ -d "$HOSTS_DIR/$HOSTNAME" ]] && die "Directory hosts/$HOSTNAME already exists."

  pick_modules

  mkdir -p "$HOSTS_DIR/$HOSTNAME"
  info "Generating hardware-configuration.nix..."
  nixos-generate-config --show-hardware-config > "$HOSTS_DIR/$HOSTNAME/hardware-configuration.nix"
  ok "hardware-configuration.nix generated."
  write_host_nix  "$HOSTNAME"
  write_home_nix  "$HOSTNAME"
  add_to_flake    "$HOSTNAME"

else
  die "Invalid option — expected 1 or 2."
fi

# ── Stage all new/modified files ─────────────────────────────────────────────
info "Staging files..."
git add \
  "$HOSTS_DIR/$HOSTNAME/" \
  "$HM_DIR/$HOSTNAME.nix" \
  "$FLAKE_DIR/flake.nix" 2>/dev/null || true
ok "Files staged."

# ── SOPS reminder ─────────────────────────────────────────────────────────────
echo ""
warn "SOPS: your age key must be present at /home/mntn/.config/sops/age/keys.txt"
warn "      before deploying, or secret decryption will fail at activation."
warn "      To derive the key from an existing age identity:"
warn "      age-keygen -o /home/mntn/.config/sops/age/keys.txt"

# ── Deploy ────────────────────────────────────────────────────────────────────
echo ""
echo "────────────────────────────────────────────"
echo "  1) nixos-install  — fresh disk mounted at /mnt"
echo "  2) nixos-rebuild  — update a running system"
echo "  3) Exit           — deploy manually later"
echo "────────────────────────────────────────────"
ask "Deploy [1/2/3]:" DEPLOY

case "$DEPLOY" in
  1)
    echo ""
    warn "The disk must be partitioned and mounted at /mnt before proceeding."
    warn "Docs: https://nixos.org/manual/nixos/stable/#sec-installation-manual-partitioning"
    ask "Proceed with nixos-install? [y/N]:" CONFIRM
    [[ "$CONFIRM" =~ ^[Yy]$ ]] || { info "Aborted."; exit 0; }
    info "Running nixos-install --flake .#$HOSTNAME ..."
    nixos-install --flake "$FLAKE_DIR#$HOSTNAME" --no-root-passwd
    ;;
  2)
    info "Running nixos-rebuild switch --flake .#$HOSTNAME ..."
    nixos-rebuild switch --flake "$FLAKE_DIR#$HOSTNAME"
    ;;
  *)
    echo ""
    ok "All files are staged. Deploy when ready:"
    echo ""
    echo -e "    ${BOLD}# Fresh install:${RESET}"
    echo -e "    nixos-install --flake $FLAKE_DIR#$HOSTNAME --no-root-passwd"
    echo ""
    echo -e "    ${BOLD}# Running system:${RESET}"
    echo -e "    nixos-rebuild switch --flake $FLAKE_DIR#$HOSTNAME"
    echo ""
    ;;
esac