#!/usr/bin/env bash
# Pushes newly-generated SSH host keys from host-keys/ to already-running
# NixOS hosts, so they can decrypt sops secrets after a nixos-rebuild
# following scripts/secrets/sync-host-keys.sh --regenerate-all-keys.
#
# Before pushing any key, verifies that .sops.yaml and secrets/*.yaml are
# committed and pushed to the remote -- hosts rebuild from the remote Gitea
# flake, so recipient changes must land there before any rebuild, not just
# before the key push.
#
#   push-host-keys.sh --all [--dry-run] [--skip-git-check]
#   push-host-keys.sh <target> [--dry-run] [--skip-git-check]
#
#   --all              Push to every reachable managed host. Default when no
#                      target is given.
#   <target>           Push to one flake target only (e.g. lxc-server).
#   --dry-run          Print what would be done; write nothing.
#   --skip-git-check   Skip the commit/push check. Use only when the remote
#                      already has the current .sops.yaml/secrets/*.yaml.
#
# SSH: connects as SSH_USER@<hostname> (default: nixos, the user with the
# admin authorized key), then installs files via sudo. You will be prompted
# for the sudo password once per host.
# Hosts are reached at their bare hostname (relies on LAN DNS/mDNS).
set -euo pipefail

repo_root="$(cd "$(dirname "$0")/../.." && pwd)"
keydir="${repo_root}/host-keys"

# shellcheck source=../env.sh
source "${repo_root}/scripts/env.sh"
# shellcheck source=../lib/nix-eval.sh
source "${repo_root}/scripts/lib/nix-eval.sh"

: "${SSH_USER:=nixos}"
SSH_OPTS=(-o StrictHostKeyChecking=no -o BatchMode=yes -o ConnectTimeout=5)

dry_run=0
skip_git_check=0

usage() {
  cat <<EOF
Usage: $0 [--all | <target>] [--dry-run] [--skip-git-check]

  --all              Push to every reachable managed host. Default when no
                     target is given.
  <target>           Push to one flake target only (e.g. lxc-server).
  --dry-run          Print what would be done; write nothing.
  --skip-git-check   Skip the check that .sops.yaml/secrets/*.yaml are
                     committed and pushed to the remote repo.

Environment:
  SSH_USER           SSH username (default: nixos).
EOF
}

locally_managed_hosts() {
  for f in "${keydir}"/*_ssh_host_ed25519_key.pub; do
    [[ -e "$f" ]] || continue
    basename "$f" _ssh_host_ed25519_key.pub
  done
}

# --- git state check/fix --------------------------------------------------
# Hosts rebuild from the remote Gitea flake:
#   nixos-rebuild switch --flake "git+https://<gitea>/nixos.git#<target>"
# so .sops.yaml (updated recipients) and secrets/*.yaml (re-encrypted DEKs)
# must be committed and pushed before any rebuild can succeed. This check
# catches the common case where --regenerate-all-keys was just run but the
# resulting diff hasn't been committed/pushed yet.
ensure_remote_current() {
  [[ "$skip_git_check" -eq 1 ]] && return

  cd "$repo_root"

  local dirty_unstaged dirty_staged
  dirty_unstaged="$(git diff --name-only -- .sops.yaml secrets/ 2>/dev/null || true)"
  dirty_staged="$(git diff --cached --name-only -- .sops.yaml secrets/ 2>/dev/null || true)"

  if [[ -n "$dirty_unstaged" || -n "$dirty_staged" ]]; then
    echo "Uncommitted changes in sops-managed files:"
    [[ -n "$dirty_unstaged" ]] && sed 's/^/  (unstaged) /' <<<"$dirty_unstaged"
    [[ -n "$dirty_staged"   ]] && sed 's/^/  (staged)   /' <<<"$dirty_staged"
    echo
    if [[ "$dry_run" -eq 1 ]]; then
      echo "[dry-run] would prompt to commit .sops.yaml/secrets/ before continuing."
    else
      read -rp "Commit .sops.yaml + secrets/ now? [y/N]: " ans
      if [[ "$ans" =~ ^[Yy]$ ]]; then
        git add -- .sops.yaml secrets/
        git commit -m "secrets: update recipients and re-encrypt for host key changes"
        echo "Committed."
      else
        echo "Continuing with uncommitted changes -- the remote won't have the"
        echo "updated recipients until you commit and push."
      fi
    fi
    echo
  fi

  # Check if we're ahead of the remote tracking branch
  local ahead
  ahead="$(git rev-list --count '@{upstream}..HEAD' 2>/dev/null || echo "")"
  if [[ -z "$ahead" ]]; then
    echo "NOTE: no remote tracking branch found -- skipping push check."
    echo "      Ensure the remote has the current .sops.yaml/secrets/ before"
    echo "      triggering nixos-rebuild on any host."
    echo
    return
  fi

  if [[ "$ahead" -gt 0 ]]; then
    echo "Local branch is ${ahead} commit(s) ahead of remote."
    if [[ "$dry_run" -eq 1 ]]; then
      echo "[dry-run] would prompt to push before continuing."
    else
      read -rp "Push to remote now? [y/N]: " ans
      if [[ "$ans" =~ ^[Yy]$ ]]; then
        git push
        echo "Pushed."
      else
        echo "Continuing without pushing -- remember to push before running"
        echo "nixos-rebuild on any of these hosts."
      fi
    fi
    echo
  fi
}

# --- key installation (shared) -------------------------------------------
_do_push() {
  local hostname="$1" target="$2"
  local keyfile="${keydir}/${target}_ssh_host_ed25519_key"
  local pubfile="${keyfile}.pub"

  if [[ "$dry_run" -eq 1 ]]; then
    echo "  [dry-run] would scp host-keys/${target}_ssh_host_ed25519_key{,.pub} to /tmp/"
    echo "  [dry-run] would: sudo install -m 0600/0644 to /etc/ssh/ and rm /tmp copies"
    return
  fi

  # Upload to /tmp (writable as nixos, no privilege needed)
  scp -o StrictHostKeyChecking=no \
    "$keyfile" "${SSH_USER}@${hostname}:/tmp/push_ed25519_key"
  scp -o StrictHostKeyChecking=no \
    "$pubfile" "${SSH_USER}@${hostname}:/tmp/push_ed25519_key.pub"

  # Install with correct permissions in one interactive sudo session
  echo "  (sudo password may be required)"
  ssh -t -o StrictHostKeyChecking=no "${SSH_USER}@${hostname}" \
    "sudo bash -s" <<'REMOTE'
install -m 0600 /tmp/push_ed25519_key     /etc/ssh/ssh_host_ed25519_key
install -m 0644 /tmp/push_ed25519_key.pub /etc/ssh/ssh_host_ed25519_key.pub
rm -f /tmp/push_ed25519_key /tmp/push_ed25519_key.pub
echo "  [ok] host key installed"
REMOTE

  # Drop the stale known_hosts entry for this host (public key just changed)
  ssh-keygen -R "$hostname" 2>/dev/null || true

  echo "  Done. Run nixos-rebuild switch on ${hostname} to activate."
}

# --- single named target --------------------------------------------------
push_target() {
  local target="$1"
  local keyfile="${keydir}/${target}_ssh_host_ed25519_key"

  if [[ ! -f "$keyfile" ]]; then
    echo "ERROR: host-keys/${target}_ssh_host_ed25519_key not found." >&2
    echo "       This target may not be locally managed (e.g. &${target} was" >&2
    echo "       registered from the host's real SSH key, not generated here)." >&2
    exit 1
  fi

  local hostname
  hostname="$(flake_target_hostname "$repo_root" "$target")"
  if [[ -z "$hostname" ]]; then
    echo "ERROR: cannot resolve hostname for '${target}' from the flake." >&2
    exit 1
  fi

  echo "==> ${target} (→ ${hostname})"

  if ! ssh "${SSH_OPTS[@]}" "${SSH_USER}@${hostname}" true 2>/dev/null; then
    echo "  SKIP: ${SSH_USER}@${hostname} unreachable."
    return
  fi

  # Sanity-check that /etc/flake-target on the host agrees
  local live_target
  live_target="$(ssh "${SSH_OPTS[@]}" "${SSH_USER}@${hostname}" \
    "cat /etc/flake-target 2>/dev/null || true")"
  if [[ -n "$live_target" && "$live_target" != "$target" ]]; then
    echo "  WARN: host reports /etc/flake-target='${live_target}', not '${target}'."
    echo "       Pushing the key you specified (${target}) anyway."
  fi

  _do_push "$hostname" "$target"
}

# --- all managed hosts ----------------------------------------------------
# For each unique hostname derived from managed targets, SSHes in and reads
# /etc/flake-target to determine which key to push -- handles the case where
# multiple targets share a hostname (e.g. lxc-server and proxmox-server both
# resolve to "server"; only one is actually running).
push_all() {
  mapfile -t managed < <(locally_managed_hosts)
  if [[ "${#managed[@]}" -eq 0 ]]; then
    echo "No managed keys in host-keys/ -- nothing to push."
    return
  fi

  echo "Pushing to all reachable managed hosts..."
  echo

  declare -A seen_hostnames=()
  local t hostname
  for t in "${managed[@]}"; do
    hostname="$(flake_target_hostname "$repo_root" "$t" 2>/dev/null || true)"
    [[ -z "$hostname" ]] && continue
    [[ -n "${seen_hostnames[$hostname]+x}" ]] && continue
    seen_hostnames["$hostname"]=1

    echo "==> checking ${hostname}"

    if ! ssh "${SSH_OPTS[@]}" "${SSH_USER}@${hostname}" true 2>/dev/null; then
      echo "  SKIP: ${SSH_USER}@${hostname} unreachable."
      continue
    fi

    # Ask the host which flake target it actually is
    local live_target
    live_target="$(ssh "${SSH_OPTS[@]}" "${SSH_USER}@${hostname}" \
      "cat /etc/flake-target 2>/dev/null || true")"

    if [[ -z "$live_target" ]]; then
      echo "  SKIP: no /etc/flake-target on host -- can't determine which key to push."
      continue
    fi

    local live_keyfile="${keydir}/${live_target}_ssh_host_ed25519_key"
    if [[ ! -f "$live_keyfile" ]]; then
      echo "  SKIP: host is '${live_target}' but no host-keys/${live_target}_... (hand-registered key, not managed here)."
      continue
    fi

    echo "  target: ${live_target}"
    _do_push "$hostname" "$live_target"
  done
}

# --- main -----------------------------------------------------------------
mode="all"
target_arg=""
extra_args=()

for arg in "$@"; do
  case "$arg" in
    --dry-run)        dry_run=1 ;;
    --skip-git-check) skip_git_check=1 ;;
    --all)            mode="all" ;;
    -h|--help)        usage; exit 0 ;;
    --*)              echo "Unknown option: $arg" >&2; usage >&2; exit 1 ;;
    *)                extra_args+=("$arg") ;;
  esac
done

if [[ "${#extra_args[@]}" -gt 1 ]]; then
  echo "ERROR: specify at most one target (or --all)." >&2
  usage >&2; exit 1
elif [[ "${#extra_args[@]}" -eq 1 ]]; then
  mode="single"
  target_arg="${extra_args[0]}"
fi

[[ "$dry_run" -eq 1 ]] && { echo "[dry-run] no changes will be made"; echo; }

nix_extra_opts
ensure_remote_current

if [[ "$mode" == "single" ]]; then
  push_target "$target_arg"
else
  push_all
fi

echo
if [[ "$dry_run" -eq 1 ]]; then
  echo "[dry-run] Nothing was changed. Re-run without --dry-run to apply."
else
  echo "Key push complete. For each updated host, run nixos-rebuild switch to"
  echo "apply the config and let sops-nix decrypt secrets with the new key."
fi