#!/bin/bash

# omarchy:summary=Factory-reset this machine back to its freshly-installed state
# omarchy:group=setup
# omarchy:requires-sudo=true
# omarchy:examples=sudo omarchy-system-factory-reset

# Returns the machine to provisioning state — fully installed, no user, first-boot
# setup pending — the state you want before selling or handing the machine on.
#
# On machines installed with an ISO that created the @factory snapshot, the
# running root (@) is swapped for a fresh writable clone of @factory, so
# user-installed packages and /etc drift are gone too. Machines without
# @factory get a degraded reset: the current system is kept and only user
# accounts, home data, and machine identity are wiped.
#
# The heavy lifting happens on the next boot via omarchy-system-factory-reset-finish, staged
# here. On encrypted installs the LUKS volume is re-keyed to a throwaway
# passphrase with an auto-unlock keyfile for the provisioning window; first-boot setup
# re-keys it to the new owner's password and removes the keyfile.
#
# Honest limitations:
#  - On unencrypted disks this is deletion, not forensic erasure. fstrim helps
#    on SSDs; use blkdiscard + a reinstall if you need certainty.
#  - LUKS re-keying changes passphrases, not the volume key. Freed extents
#    remain readable to someone with raw-device access and a valid passphrase.
#    Use cryptsetup reencrypt or a reinstall if you need certainty.

set -euo pipefail

PROVISIONING_DIR=/var/lib/omarchy/provisioning
OMARCHY_PATH="${OMARCHY_PATH:-/usr/share/omarchy}"
TOP_MNT=/run/omarchy-system-factory-reset/top
NEXT_NAME=@omarchy-reset-next
LOG_FILE=/var/log/omarchy-system-factory-reset.log

# Self-elevate so the menu entry (and a bare shell invocation) need no sudo
# prefix. The caller's gum theme vars are forwarded as `env` arguments rather
# than via `sudo -E`, so styling survives even under an env_reset sudoers.
if (( EUID != 0 )); then
  mapfile -t gum_env < <(env | grep '^GUM_' || true)
  exec sudo env "${gum_env[@]}" "$0" "$@"
fi

export PATH="$OMARCHY_PATH/bin:$PATH"

log() {
  echo "$1" | tee -a "$LOG_FILE" >/dev/null
  gum style --foreground 8 "  $1"
}

fail() {
  gum style --foreground 1 "Error: $1"
  exit 1
}

root_device() {
  findmnt -no SOURCE / | sed 's/\[.*\]//'
}

root_is_btrfs_at() {
  [[ $(findmnt -no FSTYPE /) == btrfs ]] && findmnt -no OPTIONS / | grep -q "subvol=/@\(,\|$\)"
}

# Resolve the crypto_LUKS partition backing the root, or return non-zero if
# the root is not on LUKS. Prefers the cmdline cryptdevice= spec (archinstall
# writes PARTUUID=, the pre-mounted path UUID=), and falls back to walking the
# device tree for roots reached via rd.luks/crypttab with a plain /dev/mapper
# root and no cryptdevice=.
luks_device() {
  local spec
  spec=$(grep -o 'cryptdevice=[^ :]*' /proc/cmdline | head -1 | cut -d= -f2-)
  case $spec in
    UUID=*) echo "/dev/disk/by-uuid/${spec#UUID=}"; return 0 ;;
    PARTUUID=*) echo "/dev/disk/by-partuuid/${spec#PARTUUID=}"; return 0 ;;
    LABEL=*) echo "/dev/disk/by-label/${spec#LABEL=}"; return 0 ;;
    PARTLABEL=*) echo "/dev/disk/by-partlabel/${spec#PARTLABEL=}"; return 0 ;;
    /dev/*) echo "$spec"; return 0 ;;
  esac

  local src part
  src=$(findmnt -no SOURCE / | sed 's/\[.*//')
  [[ -n $src ]] || return 1
  part=$(lsblk -nspo NAME,FSTYPE "$src" 2>/dev/null | awk '$2=="crypto_LUKS"{print $1; exit}')
  [[ -n $part ]] && { echo "$part"; return 0; }
  return 1
}

encrypted_install() {
  luks_device >/dev/null 2>&1
}

# No trailing `| head` stage: under pipefail an infinite tr killed by
# SIGPIPE fails the substitution and errexit aborts the whole reset.
generate_passphrase() {
  local raw
  raw=$(head -c 4096 /dev/urandom | tr -dc 'a-zA-Z0-9')
  printf '%s' "${raw:0:48}"
}

cleanup() {
  if (( ${swap_done:-0} == 0 )); then
    rollback_degraded_rekey
  fi
  if mountpoint -q "$TOP_MNT" 2>/dev/null; then
    if [[ -d $TOP_MNT/$NEXT_NAME && ${swap_done:-0} == 0 ]]; then
      btrfs subvolume delete --recursive "$TOP_MNT/$NEXT_NAME" >/dev/null 2>&1 || true
    fi
    umount -R "$TOP_MNT" 2>/dev/null || true
  fi
  rmdir "$TOP_MNT" 2>/dev/null || true
}
trap cleanup EXIT

confirm_reset() {
  local degraded="$1"

  echo
  gum style --bold --foreground 1 "Reset this computer to factory state?"
  echo
  gum style "This will permanently erase:"
  gum style "  • All user accounts and everything in /home"
  if [[ $degraded == true ]]; then
    gum style --foreground 3 "  • Machine identity (network connections, host keys, machine-id)"
    echo
    gum style --foreground 3 "This machine was installed before factory snapshots existed, so the"
    gum style --foreground 3 "current system (installed packages, /etc changes) is kept — only user"
    gum style --foreground 3 "accounts and machine state are wiped."
  else
    gum style "  • All packages and system changes made since installation"
    gum style "  • Machine identity (network connections, host keys, machine-id)"
  fi
  echo
  gum style "The next boot asks for a new user, exactly like a fresh install."
  gum style --foreground 8 "Note: on disks without encryption this is deletion, not secure erasure."
  echo

  local typed
  typed=$(gum input --placeholder "Type 'reset' to continue" --prompt "> ") || exit 1
  [[ $typed == "reset" ]] || fail "Reset not confirmed."
}

# Re-key to a throwaway passphrase whose keyfile auto-unlocks boot during the
# provisioning window. $1 is the root of the system the staging applies to ("" for the
# running root in degraded mode).
stage_luks_rekey() {
  local next="$1" device passphrase

  device=$(luks_device) || fail "encrypted install, but no usable cryptdevice= in /proc/cmdline"
  [[ -e $device ]] || fail "LUKS device $device not found"

  echo
  gum style "Confirm your disk encryption passphrase to authorize the re-key."
  local current
  while true; do
    current=$(gum input --password --placeholder "Current disk encryption passphrase" --prompt "Passphrase> ") || exit 1
    if cryptsetup open --test-passphrase --key-file <(printf '%s' "$current") "$device" 2>/dev/null; then
      break
    fi
    gum style --foreground 1 "That passphrase does not unlock $device. Try again."
  done

  passphrase=$(generate_passphrase)
  cryptsetup luksAddKey --key-file <(printf '%s' "$current") "$device" <(printf '%s' "$passphrase")

  install -d -m 755 "$next$PROVISIONING_DIR"
  printf '%s' "$passphrase" >"$next$PROVISIONING_DIR/luks-key"
  chmod 600 "$next$PROVISIONING_DIR/luks-key"

  install -d -m 755 "$next/etc/omarchy"
  printf '%s' "$passphrase" >"$next/etc/omarchy/provisioning.key"
  chmod 600 "$next/etc/omarchy/provisioning.key"

  install -d "$next/etc/limine-entry-tool.d" "$next/etc/mkinitcpio.conf.d"
  echo 'KERNEL_CMDLINE[default]+=" cryptkey=rootfs:/etc/omarchy/provisioning.key"' \
    >"$next/etc/limine-entry-tool.d/99-omarchy-provisioning-unlock.conf"
  echo 'FILES+=(/etc/omarchy/provisioning.key)' \
    >"$next/etc/mkinitcpio.conf.d/99-omarchy-provisioning-key.conf"
}

install_provisioning_units() {
  local next="$1" unit_src="$2"

  install -Dm644 "$unit_src/omarchy-provision-owner.service" \
    "$next/etc/systemd/system/omarchy-provision-owner.service"
  install -Dm644 "$unit_src/omarchy-system-factory-reset-finish.service" \
    "$next/etc/systemd/system/omarchy-system-factory-reset-finish.service"

  install -d "$next/etc/systemd/system/multi-user.target.wants" \
    "$next/etc/systemd/system/sysinit.target.wants"
  ln -sf /etc/systemd/system/omarchy-provision-owner.service \
    "$next/etc/systemd/system/multi-user.target.wants/omarchy-provision-owner.service"
  ln -sf /etc/systemd/system/omarchy-system-factory-reset-finish.service \
    "$next/etc/systemd/system/sysinit.target.wants/omarchy-system-factory-reset-finish.service"
}

# Start the ESP's limine.conf over from the shipped template and drop the
# previous Omarchy identity's state. limine-entry-tool keys OS entries by
# machine-id, so after a reset gives the machine a fresh identity, the
# previous system's entry (with the now-stale UKI hash) would survive, sort
# first, and make Limine refuse the rebuilt UKI with a hash-mismatch warning.
# Only machine-ids the old limine.conf referenced are removed — a shared ESP
# may hold other installations' machine-id directories.
reset_limine_config() {
  local root="$1" esp="$2" template found=""

  local old_ids=""
  if [[ -f $root$esp/limine.conf ]]; then
    old_ids=$(grep -o 'machine-id=[0-9a-f]\{32\}' "$root$esp/limine.conf" | cut -d= -f2 | sort -u)
  fi

  for template in "$root/usr/share/omarchy/install/assets/limine/limine.conf" \
                  "$root/usr/share/omarchy/default/limine/limine.conf"; do
    if [[ -f $template ]]; then
      cp "$template" "$root$esp/limine.conf"
      found=1
      break
    fi
  done
  [[ -n $found ]] || fail "no limine.conf template in $root/usr/share/omarchy"

  local machine_id old_id
  machine_id=$(cat "$root/etc/machine-id" 2>/dev/null || true)
  for old_id in $old_ids; do
    [[ $old_id == "$machine_id" ]] && continue
    rm -rf "${root:?}$esp/$old_id"
  done
}

# Every UKI entry in limine.conf carries a #blake2b hash of the file it
# points at; a mismatch makes Limine stop at a scary warning. Never hand over
# a staged reset in that state.
verify_limine_hashes() {
  local root="$1" esp="$2" line path hash file
  while IFS= read -r line; do
    path=${line#*boot():}
    path=${path%%#*}
    hash=${line##*#}
    file="$root$esp$path"
    [[ -f $file ]] || fail "limine.conf points at missing $path"
    [[ $(b2sum "$file" | cut -d' ' -f1) == "$hash" ]] ||
      fail "limine.conf hash for $path does not match the rebuilt file"
  done < <(grep -o 'boot():/EFI/Linux/[^#]*#[0-9a-f]*' "$root$esp/limine.conf")
}

# Rebuild the UKI from inside the next root so the deployed kernel/initramfs
# match the factory root's modules (the seller may have updated kernels since
# install) and embed the auto-unlock keyfile on encrypted machines.
rebuild_next_boot() {
  local next="$1"

  local esp_source esp_mount
  esp_mount=$(awk '$3 == "vfat" && $2 ~ /^\/(boot|efi)/ { print $2; exit }' "$next/etc/fstab")
  esp_source=$(awk '$3 == "vfat" && $2 ~ /^\/(boot|efi)/ { print $1; exit }' "$next/etc/fstab")
  [[ -n $esp_mount && -n $esp_source ]] || fail "could not find the ESP in $next/etc/fstab"

  local esp_device=$esp_source
  if [[ $esp_source == UUID=* ]]; then
    esp_device="/dev/disk/by-uuid/${esp_source#UUID=}"
  fi
  [[ -e $esp_device ]] || fail "ESP device $esp_device not found"

  mount "$esp_device" "$next$esp_mount"

  local dir
  for dir in proc sys dev run; do
    mount --rbind "/$dir" "$next/$dir"
    mount --make-rslave "$next/$dir"
  done

  # Btrfs snapshots do not recurse into nested subvolumes, so the factory
  # root carries only an empty /swap directory where the hibernation swapfile
  # lived. Recreate it and refresh the resume offset before the UKI build
  # bakes the cmdline; the empty placeholder directory has to go first so the
  # subvolume can be recreated.
  if [[ -x $next/usr/bin/omarchy-hibernation-setup ]] && grep -q "/swap/swapfile" "$next/etc/fstab"; then
    log "Recreating the hibernation swapfile in the factory system"
    if [[ -d $next/swap ]] && ! btrfs subvolume show "$next/swap" >/dev/null 2>&1; then
      rmdir "$next/swap" 2>/dev/null || true
    fi
    # The factory root still carries the resume drop-ins from install, which
    # make hibernation-setup conclude it has nothing to do; with the swapfile
    # gone they must go so it reconfigures from scratch. The resume-offset
    # drop-in especially: hibernation-setup only recomputes the physical
    # offset when it is absent, so a surviving one bakes the old swapfile's
    # offset into the rebuilt UKI and breaks resume.
    if [[ ! -f $next/swap/swapfile ]]; then
      rm -f "$next/etc/mkinitcpio.conf.d/omarchy_resume.conf" \
        "$next/etc/limine-entry-tool.d/resume.conf"
    fi
    if ! chroot "$next" env OMARCHY_PATH=/usr/share/omarchy \
      /usr/bin/omarchy-hibernation-setup --force --no-rebuild >>"$LOG_FILE" 2>&1; then
      fail "hibernation setup failed in the factory root (see $LOG_FILE)"
    fi
    [[ -f $next/swap/swapfile ]] ||
      fail "hibernation setup did not recreate $next/swap/swapfile (see $LOG_FILE)"
  fi

  reset_limine_config "$next" "$esp_mount"

  log "Rebuilding boot files from the factory system (this can take a minute)"
  if ! chroot "$next" /usr/bin/limine-update >>"$LOG_FILE" 2>&1; then
    fail "limine-update failed in the factory root (see $LOG_FILE)"
  fi

  verify_limine_hashes "$next" "$esp_mount"

  for dir in run dev sys proc; do
    umount -R "$next/$dir" 2>/dev/null || true
  done
  umount "$next$esp_mount"
}

# Remove the seller's account material and machine identity from the retained
# @factory baseline so it can neither be mounted for recovery nor restore the
# seller's account on a future reset. Idempotent (a scrubbed baseline has no
# uid>=1000 accounts left to remove).
sanitize_factory_baseline() {
  local factory="$1" user
  btrfs property set -ts "$factory" ro false

  for user in $(awk -F: '$3 >= 1000 && $3 < 60000 { print $1 }' "$factory/etc/passwd"); do
    userdel --root "$factory" "$user" 2>>"$LOG_FILE" || true
    rm -rf "${factory:?}/home/$user"
  done
  passwd --root "$factory" --lock root >>"$LOG_FILE" 2>&1 || true
  rm -f "$factory"/etc/ssh/ssh_host_*
  rm -f "$factory"/etc/NetworkManager/system-connections/*
  rm -rf "$factory"/var/lib/NetworkManager/* "$factory/var/lib/tailscale" "$factory/var/lib/iwd"
  rm -f "$factory/var/lib/sddm/state.conf" "$factory/etc/sddm.conf.d/autologin.conf"
  : >"$factory/etc/machine-id"

  btrfs property set -ts "$factory" ro true
}

stage_full_reset() {
  local top="$TOP_MNT" next="$TOP_MNT/$NEXT_NAME"

  log "Cloning the factory snapshot"
  [[ -d $next ]] && btrfs subvolume delete --recursive "$next" >/dev/null
  btrfs subvolume snapshot "$top/@factory" "$next" >>"$LOG_FILE"

  local unit_src="$next/usr/share/omarchy/install/provisioning"
  [[ -f $unit_src/omarchy-provision-owner.service && -x $next/usr/bin/omarchy-provision-owner ]] ||
    fail "the factory snapshot predates provisioning support; cannot reset from it"

  log "Scrubbing machine identity from the factory system"
  systemd-id128 new >"$next/etc/machine-id"
  rm -f "$next"/etc/ssh/ssh_host_*
  rm -f "$next"/etc/NetworkManager/system-connections/*
  rm -rf "$next"/var/lib/NetworkManager/* "$next/var/lib/tailscale" "$next/var/lib/iwd"
  rm -f "$next/var/lib/sddm/state.conf" "$next/etc/sddm.conf.d/autologin.conf"

  # A factory snapshot from a normal (normal) install contains the original
  # user account; first-boot setup must start from none. A leftover account
  # would keep its password hash and group memberships (including wheel), so
  # failure here has to abort the reset, not be shrugged off.
  local user
  for user in $(awk -F: '$3 >= 1000 && $3 < 60000 { print $1 }' "$next/etc/passwd"); do
    log "Removing user $user from the factory system"
    userdel --root "$next" "$user" 2>>"$LOG_FILE" ||
      fail "could not remove user $user from the factory system (see $LOG_FILE)"
  done
  passwd --root "$next" --lock root >>"$LOG_FILE" 2>&1 || true

  # @factory itself survives the wipe as the baseline for future resets. If it
  # came from a normal install it still holds the seller's account and
  # /etc/shadow, which the new wheel user could mount and read — and a second
  # reset would restore that account. Scrub it once, in place.
  sanitize_factory_baseline "$top/@factory"

  # Keep the Node tarball reachable for offline first-boot finalization.
  if ! compgen -G "$next$PROVISIONING_DIR/packages/node-v*.tar.gz" >/dev/null; then
    if compgen -G "$PROVISIONING_DIR/packages/node-v*.tar.gz" >/dev/null; then
      install -d -m 755 "$next$PROVISIONING_DIR/packages"
      cp "$PROVISIONING_DIR"/packages/node-v*.tar.gz "$next$PROVISIONING_DIR/packages/"
    fi
  fi

  install -d -m 755 "$next$PROVISIONING_DIR"
  touch "$next$PROVISIONING_DIR/pending" "$next$PROVISIONING_DIR/wipe-pending"

  install_provisioning_units "$next" "$unit_src"

  if encrypted_install; then
    stage_luks_rekey "$next"
  fi

  rebuild_next_boot "$next"

  log "Activating the factory system"
  local old="@omarchy-old-$(date +%s)"
  mv "$top/@" "$top/$old"
  mv "$next" "$top/@"
  swap_done=1
  sync
}

stage_degraded_reset() {
  local unit_src="$OMARCHY_PATH/install/provisioning"
  [[ -f $unit_src/omarchy-provision-owner.service && -x /usr/bin/omarchy-provision-owner ]] ||
    fail "this Omarchy version does not ship provisioning; update first"

  # Fallible staging (re-key, boot rebuild) happens before the wipe is armed:
  # a failure must leave the machine untouched, not schedule a wipe for a
  # reset that never finished staging. The degraded path stages into the LIVE
  # root, so a mid-staging failure must also undo the auto-unlock material it
  # already wrote — otherwise a later boot rebuild would silently embed it.
  install -d -m 755 "$PROVISIONING_DIR"
  if encrypted_install; then
    degraded_rekey_staged=1
    stage_luks_rekey ""
    log "Rebuilding boot files with the provisioning unlock key"
    if ! limine-update >>"$LOG_FILE" 2>&1; then
      fail "limine-update failed (see $LOG_FILE)"
    fi
    degraded_rekey_staged=0
  fi

  install_provisioning_units "" "$unit_src"
  touch "$PROVISIONING_DIR/pending" "$PROVISIONING_DIR/wipe-pending" "$PROVISIONING_DIR/wipe-degraded"

  swap_done=1
}

# Undo the live-root auto-unlock material a failed degraded rekey left behind.
rollback_degraded_rekey() {
  (( ${degraded_rekey_staged:-0} )) || return 0
  rm -f /etc/omarchy/provisioning.key "$PROVISIONING_DIR/luks-key" \
    /etc/limine-entry-tool.d/99-omarchy-provisioning-unlock.conf \
    /etc/mkinitcpio.conf.d/99-omarchy-provisioning-key.conf
  limine-update >>"$LOG_FILE" 2>&1 || true
}

main() {
  omarchy-cmd-present btrfs || { echo "Error: btrfs-progs is required" >&2; exit 1; }

  root_is_btrfs_at || fail "reset requires the standard Omarchy Btrfs layout (subvol=@)"

  touch "$LOG_FILE"
  chmod 600 "$LOG_FILE"
  swap_done=0

  local device
  device=$(root_device)
  [[ -n $device ]] || fail "could not determine the root device"

  mkdir -p "$TOP_MNT"
  mountpoint -q "$TOP_MNT" || mount -o subvolid=5 "$device" "$TOP_MNT"

  local degraded=false
  [[ -d $TOP_MNT/@factory ]] || degraded=true

  confirm_reset "$degraded"

  # The running system's limine-snapper-sync must not rewrite the ESP's
  # limine.conf behind the staged rebuild (subvolume changes below can
  # trigger it). The runtime mask evaporates on the reboot that follows.
  systemctl mask --runtime --now limine-snapper-sync.service >/dev/null 2>&1 || true
  systemctl mask --runtime --now limine-snapper-sync.path >/dev/null 2>&1 || true

  if [[ $degraded == true ]]; then
    stage_degraded_reset
  else
    stage_full_reset
  fi

  umount -R "$TOP_MNT" 2>/dev/null || true

  echo
  gum style --bold "Reset staged. The wipe finishes on the next boot."
  if gum confirm --affirmative "Reboot now" --negative "Reboot later" "Reboot to complete the reset?"; then
    systemctl reboot
  else
    gum style --foreground 3 "Do not keep using this machine — changes made now will be lost."
  fi
}

main "$@"