#!/usr/bin/env bash
# Manages host-keys/ + .sops.yaml + secrets/*.yaml recipients together, so
# a flake target's SSH host key and its sops registration never drift out
# of sync with each other or with the flake itself.
#
#   sync-host-keys.sh --all                Generate/register every flake
#                                           target missing a key.
#   sync-host-keys.sh <target>              Same, for just one target.
#   sync-host-keys.sh --remove              Interactively remove one
#                                           locally-managed key.
#   sync-host-keys.sh --regenerate-all-keys Remove and freshly regenerate
#                                           every locally-managed key.
#
# "Generate/register" is idempotent and additive only: an existing
# host-keys/ file is never touched, and .sops.yaml only ever gains an
# anchor/alias it doesn't already have -- safe to re-run any time, e.g.
# right after adding a new host to flake.nix.
#
# --remove and --regenerate-all-keys only ever operate on anchors that have
# a corresponding host-keys/<name>_ssh_host_ed25519_key file. Anchors
# without one (&admin, and any anchor for an already-deployed host whose
# real /etc/ssh key was registered by hand, e.g. &docker/&server/&nix-cache
# today) are never listed, removed, or regenerated -- this tooling only
# ever touches keys it itself manages.
set -euo pipefail

repo_root="$(cd "$(dirname "$0")/../.." && pwd)"
sops_yaml="${repo_root}/.sops.yaml"
keydir="${repo_root}/host-keys"
editor="${repo_root}/scripts/lib/sync-host-keys-edit-sops.py"

# shellcheck source=../env.sh
source "${repo_root}/scripts/env.sh"
# shellcheck source=../lib/nix-eval.sh
source "${repo_root}/scripts/lib/nix-eval.sh"
# shellcheck source=../lib/ssh-host-keys.sh
source "${repo_root}/scripts/lib/ssh-host-keys.sh"
# shellcheck source=../lib/sops-age.sh
source "${repo_root}/scripts/lib/sops-age.sh"
# shellcheck source=../lib/confirm.sh
source "${repo_root}/scripts/lib/confirm.sh"
# shellcheck source=../lib/clan-vars.sh
source "${repo_root}/scripts/lib/clan-vars.sh"

mkdir -p "$keydir"

usage() {
  cat <<EOF
Usage: $0 --all [--dry-run]
       $0 <flake-target> [--dry-run]
       $0 --remove [--dry-run]
       $0 --regenerate-all-keys [--dry-run]

  --all                    Generate + register a host key for every flake
                            target that's missing one.
  <flake-target>            Same, for just one target (e.g. lxc-server).
                            Reports if it already has one.
  --remove                  Interactively pick one locally-managed key to
                            remove from .sops.yaml and host-keys/.
  --regenerate-all-keys     Remove every locally-managed key and generate
                            fresh replacements for every current flake
                            target. Destructive -- requires typed
                            confirmation.
  --dry-run                 Combine with any of the above: print what would
                            change (host-keys/ files, .sops.yaml anchors and
                            key_groups, which secrets/*.yaml would be
                            re-encrypted) without touching anything. No keys
                            generated, no files written, no sops calls,
                            no prompts for confirmation.
EOF
}

# --- step 0: make sure we can actually decrypt anything at all -------------
#
# Registering a host means editing .sops.yaml and then running
# `sops updatekeys`, which has to decrypt each secrets file with an
# existing recipient's key before it can re-encrypt it for the new one.
# Check this before doing anything else, the same order sops/age itself
# resolves a usable key in: SOPS_AGE_KEY (inline), then SOPS_AGE_KEY_FILE,
# then the XDG default path.
ensure_admin_decrypt_key() {
  if [[ -n "${SOPS_AGE_KEY:-}" ]]; then
    echo "Using SOPS_AGE_KEY from the environment."
    return
  fi

  local key_file="$DEFAULT_SOPS_AGE_KEY_FILE"

  if [[ -s "$key_file" ]]; then
    echo "Found existing sops age key at ${key_file}."
    return
  fi

  if [[ "$dry_run" -eq 1 ]]; then
    echo "[dry-run] No sops age decryption key found (checked \$SOPS_AGE_KEY, \$SOPS_AGE_KEY_FILE, ${key_file})."
    echo "[dry-run] Would generate a new one here -- continuing the dry run without one; any"
    echo "[dry-run] 'would re-encrypt' output below couldn't actually run for real yet."
    return
  fi

  echo "No sops age decryption key found (checked \$SOPS_AGE_KEY, \$SOPS_AGE_KEY_FILE, ${key_file})."
  echo "Generating a new one at ${key_file}..."
  mkdir -p "$(dirname "$key_file")"
  nix-shell "${NIX_OPTS[@]}" -p age --run "age-keygen -o '${key_file}'" 2>&1 | grep -v "^Public key:" || true
  local new_pub
  new_pub="$(age_pubkey_from_identity_file "$key_file")"

  cat <<EOF

A brand-new age key was just generated -- it cannot decrypt anything that
already exists in secrets/*.yaml, since nothing was ever encrypted for it.
That trust can't be bootstrapped automatically (nobody can decrypt a file
for a recipient that didn't exist when it was last encrypted).

To actually use this key:
  1. Have someone who currently CAN decrypt replace the &admin entry in
     .sops.yaml with this public key:
       ${new_pub}
  2. They re-encrypt every secrets/*.yaml:
       sops updatekeys --yes secrets/common.yaml
       sops updatekeys --yes secrets/nix-cache.yaml
       sops updatekeys --yes secrets/server.yaml
  3. Re-run this script.

Exiting without making any other changes.
EOF
  exit 1
}

discover_targets() {
  # installer is the one nixosConfigurations target that doesn't import
  # sops-nix at all (see CLAUDE.md's "Security Notes" -- hardcoded login
  # password instead) -- config.sops.secrets doesn't exist for it.
  list_flake_targets "$repo_root" | grep -v '^installer$'
}

locally_managed_hosts() {
  {
    for f in "$keydir"/*_ssh_host_ed25519_key.pub; do
      [[ -e "$f" ]] || continue
      basename "$f" _ssh_host_ed25519_key.pub
    done
    local d
    for d in "${repo_root}/vars/per-machine"/*/openssh/ssh_host_ed25519_key/secret; do
      [[ -f "$d" ]] || continue
      basename "$(dirname "$(dirname "$(dirname "$d")")")"
    done
  } | sort -u
}

add_keys_json="[]"
add_aliases_json="[]"
dry_run=0

queue_host_sync() {
  local host="$1"
  local keyfile="${keydir}/${host}_ssh_host_ed25519_key"
  local has_local_key=0 has_clan_key=0 has_anchor=0
  [[ -f "$keyfile" ]] && has_local_key=1
  clan_ssh_key_exists "$host" "$repo_root" && has_clan_key=1
  grep -qE "^  - &${host} age1" "$sops_yaml" && has_anchor=1

  if [[ "$has_local_key" -eq 0 && "$has_clan_key" -eq 0 && "$has_anchor" -eq 1 ]]; then
    echo "SKIP ${host}: .sops.yaml already has an &${host} anchor, but"
    echo "     neither host-keys/${host}_ssh_host_ed25519_key nor"
    echo "     vars/per-machine/${host}/openssh/ exist locally."
    echo "     Not generating a replacement -- it wouldn't match whatever's"
    echo "     already registered (and possibly deployed). Remove the"
    echo "     &${host} line from .sops.yaml first if you really want a"
    echo "     fresh key, then re-run."
    return 1
  fi

  if [[ "$has_local_key" -eq 0 && "$has_clan_key" -eq 0 ]]; then
    if [[ "$dry_run" -eq 1 ]]; then
      echo "[dry-run] ${host}: would generate host key via clan vars"
    else
      echo "==> ${host}: generating host key via clan vars"
      clan_generate_ssh_key "$host" "$repo_root"
      has_clan_key=1
    fi
  elif [[ "$has_clan_key" -eq 1 ]]; then
    echo "==> ${host}: clan-managed SSH host key already present"
  else
    echo "==> ${host}: host key already present (host-keys/)"
  fi

  if [[ "$has_anchor" -eq 0 ]]; then
    local age_pub
    if [[ "$dry_run" -eq 1 ]]; then
      age_pub="dry-run-placeholder-not-a-real-key"
    elif [[ "$has_clan_key" -eq 1 ]]; then
      age_pub="$(ssh_pubkey_to_age "$(clan_ssh_pubkey_path "$host" "$repo_root")")"
    else
      age_pub="$(ssh_pubkey_to_age "${keyfile}.pub")"
    fi
    add_keys_json="$(jq --arg host "$host" --arg key "$age_pub" \
      '. + [{host: $host, age_key: $key}]' <<<"$add_keys_json")"
  fi

  echo "==> ${host}: checking which secrets files it references"
  local basenames
  mapfile -t basenames < <(
    nix eval --json --no-use-registries --no-accept-flake-config \
      "${repo_root}#nixosConfigurations.${host}.config.sops.secrets" \
      --apply 'builtins.mapAttrs (n: v: baseNameOf v.sopsFile)' \
    | jq -r '[.[]] | unique | .[]'
  )
  local basename
  for basename in "${basenames[@]}"; do
    add_aliases_json="$(jq --arg host "$host" --arg basename "$basename" \
      '. + [{host: $host, basename: $basename}]' <<<"$add_aliases_json")"
  done
}

# In dry-run, this runs the exact same edit logic (so idempotency/what's-
# actually-new is determined for real, not guessed) but against a scratch
# copy of .sops.yaml that's discarded afterward -- the real file is never
# opened for writing, and `sops updatekeys` never runs.
apply_edit_plan() {
  local plan="$1"
  local target="$sops_yaml"
  local tmpfile=""
  if [[ "$dry_run" -eq 1 ]]; then
    tmpfile="$(mktemp)"
    cp "$sops_yaml" "$tmpfile"
    target="$tmpfile"
  fi

  local result
  result="$(echo "$plan" | nix-shell "${NIX_OPTS[@]}" -p python3 --run "python3 '${editor}' '${target}'")"
  [[ -n "$tmpfile" ]] && rm -f "$tmpfile"

  local added removed changed
  added="$(jq -r '.added_keys[]?' <<<"$result")"
  removed="$(jq -r '.removed_keys[]?' <<<"$result")"
  changed="$(jq -r '.changed_secrets_files[]?' <<<"$result")"

  if [[ -z "$added" && -z "$removed" && -z "$changed" ]]; then
    echo "Nothing changed in .sops.yaml."
    return
  fi

  local prefix=""
  [[ "$dry_run" -eq 1 ]] && prefix="[dry-run] would "

  [[ -n "$added" ]] && echo "${prefix}Add .sops.yaml anchors: $(tr '\n' ' ' <<<"$added")"
  [[ -n "$removed" ]] && echo "${prefix}Remove .sops.yaml anchors: $(tr '\n' ' ' <<<"$removed")"

  if [[ -n "$changed" ]]; then
    if [[ "$dry_run" -eq 1 ]]; then
      echo "[dry-run] would re-encrypt:"
      while IFS= read -r basename; do
        [[ -z "$basename" ]] && continue
        echo "  secrets/${basename}"
      done <<<"$changed"
    else
      echo "Re-encrypting affected secrets files..."
      while IFS= read -r basename; do
        [[ -z "$basename" ]] && continue
        echo "==> secrets/${basename}"
        sops_updatekeys "${repo_root}/secrets/${basename}"
      done <<<"$changed"
    fi
  fi
}

flush_additions() {
  if [[ "$add_keys_json" == "[]" && "$add_aliases_json" == "[]" ]]; then
    echo "Nothing to do -- every requested target already has a fully registered host key."
    return
  fi
  echo
  echo "Applying .sops.yaml edits..."
  local plan
  plan="$(jq -n --argjson add_keys "$add_keys_json" --argjson add_aliases "$add_aliases_json" \
    '{add_keys: $add_keys, add_aliases: $add_aliases}')"
  apply_edit_plan "$plan"
  echo
  if [[ "$dry_run" -eq 1 ]]; then
    echo "[dry-run] Nothing was changed. Re-run without --dry-run to apply this."
  else
    echo "Done. Review the .sops.yaml / secrets/*.yaml diff, then commit and push --"
    echo "the flake build the installer uses has to see the new recipient(s) before"
    echo "any of these hosts can decrypt their secrets on first boot."
  fi
}

cmd_all() {
  echo "Discovering flake targets..."
  local targets
  mapfile -t targets < <(discover_targets)
  local host
  for host in "${targets[@]}"; do
    queue_host_sync "$host" || true
  done
  flush_additions
}

cmd_target() {
  local host="$1"
  local targets
  mapfile -t targets < <(discover_targets)
  if ! printf '%s\n' "${targets[@]}" | grep -qxF "$host"; then
    echo "ERROR: '${host}' is not a current nixosConfigurations target." >&2
    echo "Current targets:" >&2
    printf '  %s\n' "${targets[@]}" >&2
    exit 1
  fi
  queue_host_sync "$host" || exit 1
  flush_additions
}

cmd_remove() {
  local hosts
  mapfile -t hosts < <(locally_managed_hosts)
  if [[ "${#hosts[@]}" -eq 0 ]]; then
    echo "No locally-managed keys in host-keys/ -- nothing to remove."
    return
  fi

  echo "Locally-managed keys:"
  local i=1 host
  for host in "${hosts[@]}"; do
    local registered="not registered in .sops.yaml"
    grep -qE "^  - &${host} age1" "$sops_yaml" && registered="registered in .sops.yaml"
    local where="host-keys/"
    clan_ssh_key_exists "$host" "$repo_root" && where="clan-vars"
    printf '  %d) %s [%s, %s]\n' "$i" "$host" "$where" "$registered"
    i=$((i + 1))
  done

  local choice
  read -rp "Remove which one? (number, or blank to cancel): " choice
  if [[ -z "$choice" ]]; then
    echo "Cancelled."
    return
  fi
  if ! [[ "$choice" =~ ^[0-9]+$ ]] || (( choice < 1 || choice > ${#hosts[@]} )); then
    echo "ERROR: invalid selection." >&2
    exit 1
  fi
  local target="${hosts[$((choice - 1))]}"

  if [[ "$dry_run" -ne 1 ]]; then
    read -rp "Really remove '${target}'? Its key files will be deleted and it will lose access to every secrets file it can currently decrypt. (y/N): " confirm
    if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
      echo "Cancelled."
      return
    fi
  fi

  local plan
  plan="$(jq -n --arg host "$target" \
    '{remove_keys: [$host], remove_aliases_for_hosts: [$host]}')"
  apply_edit_plan "$plan"

  if [[ "$dry_run" -eq 1 ]]; then
    echo "[dry-run] would delete host-keys/${target}_ssh_host_ed25519_key(.pub) if present."
    echo "[dry-run] would delete vars/per-machine/${target}/openssh/ if present."
    echo "[dry-run] Nothing was changed. Re-run without --dry-run to apply this."
  else
    rm -f "${keydir}/${target}_ssh_host_ed25519_key" "${keydir}/${target}_ssh_host_ed25519_key.pub"
    rm -rf "${repo_root}/vars/per-machine/${target}/openssh"
    echo "Removed key for ${target} (host-keys/ and/or vars/per-machine/ as applicable)."
    echo
    echo "Review the diff, then commit and push."
  fi
}

cmd_regenerate_all() {
  local hosts
  mapfile -t hosts < <(locally_managed_hosts)
  if [[ "${#hosts[@]}" -eq 0 ]]; then
    echo "No locally-managed keys in host-keys/ -- nothing to regenerate."
    return
  fi

  echo "This will remove and freshly regenerate ALL locally-managed keys:"
  printf '  %s\n' "${hosts[@]}"
  echo
  echo "After regenerating, each host needs its new key before it can decrypt secrets:"
  echo "  • Already running: push the key before rebuilding:"
  echo "      scripts/secrets/push-host-keys.sh --all"
  echo "  • Not yet deployed: rebuild the install image with the new keys baked in"
  echo "      (see docs/auto-installer.md)."

  if [[ "$dry_run" -ne 1 ]]; then
    if ! confirm_typed "REGENERATE" "Type REGENERATE to confirm: "; then
      echo "Cancelled."
      return
    fi
  fi

  echo
  local hosts_json
  hosts_json="$(printf '%s\n' "${hosts[@]}" | jq -R . | jq -s .)"
  local plan
  plan="$(jq -n --argjson hosts "$hosts_json" \
    '{remove_keys: $hosts, remove_aliases_for_hosts: $hosts}')"
  apply_edit_plan "$plan"

  if [[ "$dry_run" -eq 1 ]]; then
    echo "[dry-run] would delete ${#hosts[@]} key pair(s) from host-keys/ and/or vars/per-machine/."
    echo "[dry-run] would then generate fresh clan vars replacements for the same hosts"
    echo "[dry-run] (not simulated further here -- run without --dry-run, or"
    echo "[dry-run] preview a specific target with: $0 <target> --dry-run)."
    echo
    echo "[dry-run] Nothing was changed. Re-run without --dry-run to apply this."
    return
  fi

  echo "Removing existing keys..."
  local host
  for host in "${hosts[@]}"; do
    rm -f "${keydir}/${host}_ssh_host_ed25519_key" "${keydir}/${host}_ssh_host_ed25519_key.pub"
    rm -rf "${repo_root}/vars/per-machine/${host}/openssh"
  done
  echo "Removed ${#hosts[@]} key pair(s)."

  echo
  echo "Regenerating fresh keys for every current flake target..."
  cmd_all

  echo
  echo "Next steps:"
  echo "  1. Commit and push .sops.yaml + secrets/ so the remote flake is current."
  echo "  2. Push the new host key to each already-running managed host:"
  echo "       scripts/secrets/push-host-keys.sh --all"
  echo "     (this also prompts to commit/push if step 1 wasn't done yet)"
  echo "  3. Run nixos-rebuild switch on each updated host."
  echo "  4. For hosts not yet deployed, rebuild the install image (see"
  echo "     docs/auto-installer.md)."
}

main() {
  local args=()
  local arg
  for arg in "$@"; do
    if [[ "$arg" == "--dry-run" ]]; then
      dry_run=1
    else
      args+=("$arg")
    fi
  done
  set -- "${args[@]+"${args[@]}"}"

  if [[ "$dry_run" -eq 1 ]]; then
    echo "[dry-run] no changes will be made"
    echo
  fi

  nix_extra_opts
  ensure_admin_decrypt_key

  case "${1:-}" in
    --all)
      cmd_all
      ;;
    --remove)
      cmd_remove
      ;;
    --regenerate-all-keys)
      cmd_regenerate_all
      ;;
    -h | --help | "")
      usage
      ;;
    --*)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 1
      ;;
    *)
      cmd_target "$1"
      ;;
  esac
}

main "$@"