#!/usr/bin/env bash
# ╔══════════════════════════════════════════════════════════════════════════════╗
# ║  setup/reset-arch.sh — Wipe and reinstall the OS while keeping user data  ║
# ║                                                                              ║
# ║  PURPOSE:                                                                    ║
# ║    Performs a clean OS reinstall on an existing btrfs-on-LUKS system       ║
# ║    WITHOUT destroying user home directories or passwords. Useful when      ║
# ║    the system is broken, corrupted, or needs a clean slate.                ║
# ║                                                                              ║
# ║  WHAT THIS DOES:                                                             ║
# ║    1. Detects LUKS encryption; unlocks via FIDO2 token and/or passphrase  ║
# ║    2. Saves user credentials and system config from the old @ subvolume    ║
# ║    3. Clears app configs (~/.config) from @home, preserving auth keys     ║
# ║       (Yubico/ is kept so FIDO2 login works after the reset)              ║
# ║    4. Deletes and recreates the @ (root) btrfs subvolume                  ║
# ║    5. Reinstalls the base system via pacstrap                              ║
# ║    6. Restores credentials, PAM, fstab, mkinitcpio, GRUB config          ║
# ║    7. Regenerates initramfs and GRUB menu from chroot                     ║
# ║                                                                              ║
# ║  REQUIREMENTS:                                                               ║
# ║    - Run from an Arch Linux live environment (archiso USB)                 ║
# ║    - Target disk uses btrfs with @ and @home subvolumes                   ║
# ║    - Optional: LUKS2 encryption on the root partition                     ║
# ║                                                                              ║
# ║  WARNING: The root subvolume (@ i.e. /etc, /usr, /var) is DESTROYED.     ║
# ║           User home data (@home) is preserved. Run with extreme caution.  ║
# ╚══════════════════════════════════════════════════════════════════════════════╝

# reset-arch.sh — Reset the root btrfs subvolume while preserving user home data.
#
# What this does:
#   1. Detects LUKS encryption; unlocks via FIDO2 token and/or passphrase
#   2. Saves user credentials and system config from the old @ subvolume
#   3. Clears app configs (~/.config) from @home, preserving auth keys (Yubico/)
#   4. Deletes and recreates the @ (root) btrfs subvolume
#   5. Reinstalls the base system via pacstrap
#   6. Restores credentials, PAM, fstab, mkinitcpio, GRUB config
#   7. Regenerates initramfs and GRUB menu from chroot so the system boots cleanly

set -euo pipefail
# -e: abort immediately on any error — wrong moves here could corrupt the system
# -u: error on unset variable references
# -o pipefail: catch failures in pipes

# ── Temporary working directory ───────────────────────────────────────────────
# All temporary files (saved configs, mount points) go here.
# Cleaned up automatically on EXIT, INT, or error via trap.
TMPDIR=$(mktemp -d /tmp/arch-reset.XXXXXX)
trap 'rm -rf "$TMPDIR"' EXIT

# Simple "pause and wait for user" helper
pause() { read -rp "Press ENTER to continue..."; }

# ── Banner ────────────────────────────────────────────────────────────────────
echo "======================================="
echo "  M-Archy System Reset"
echo "======================================="
echo "This will:"
echo "  • Delete and recreate the root (@) btrfs subvolume"
echo "  • Reinstall base system packages from scratch"
echo "  • Clear all user ~/.config directories (auth keys preserved)"
echo "  • Preserve home directories, passwords, and FIDO2 login keys"
echo ""

# ── Required tools in live environment ──────────────────────────────────────
# WHY: The archiso live environment has minimal packages. We need these specific
#      tools that may not be pre-installed on the ISO:
#      - cryptsetup: to detect and unlock LUKS partitions
#      - btrfs-progs: to manage btrfs subvolumes
#      - jq: used by some sub-scripts
#      - libfido2: for FIDO2/token-based LUKS unlock
# -yd: sync database, download-only skip (don't install dependencies blindly)
pacman -Syd --noconfirm cryptsetup btrfs-progs jq libfido2

# ── Drive selection ──────────────────────────────────────────────────────────
# List all block devices so the user can see what's available
lsblk
echo ""
read -rp "Enter drive to reset (e.g., /dev/sda): " DRIVE

# Assume standard partition layout from arch-autoinstall.sh:
#   partition 1 = EFI/boot
#   partition 2 = root (possibly LUKS)
# NOTE: This assumes a non-NVMe drive. NVMe would need ${DRIVE}p2 etc.
ROOT_PART="${DRIVE}2"
EFI_PART="${DRIVE}1"

echo ""
echo "WARNING: The root subvolume on $ROOT_PART will be DELETED and reinstalled."
echo "         User home directories will be preserved."
echo "         App configs (~/.config) will be wiped (Yubico auth keys excepted)."
echo ""

# Require explicit "YES" (uppercase) to prevent accidental data loss
read -rp "Type YES to continue: " _CONFIRM
[[ "$_CONFIRM" == "YES" ]] || { echo "Aborted."; exit 1; }

# ── LUKS detection and unlock ────────────────────────────────────────────────
# Default assumption: root partition is NOT encrypted — MAPPER_DEV points to it directly.
# If LUKS is detected, MAPPER_DEV is updated to /dev/mapper/cryptroot after unlock.
MAPPER_DEV="$ROOT_PART"

# cryptsetup isLuks: exits 0 if the partition has a LUKS header, non-zero if not.
if cryptsetup isLuks "$ROOT_PART" 2>/dev/null; then
    echo ""
    echo "Partition $ROOT_PART is LUKS2-encrypted."
    echo "Select unlock method:"
    echo "  1) Try enrolled token (FIDO2/TPM2) first, fall back to passphrase  [recommended]"
    echo "  2) Passphrase only"
    echo "  3) Enrolled token only (FIDO2/TPM2)"
    read -rp "Choice [1]: " _UNLOCK
    _UNLOCK="${_UNLOCK:-1}"  # Default to option 1 if user just presses Enter

    case "$_UNLOCK" in
        1)
            # Best of both worlds: try hardware token first, fall back to passphrase.
            # This works whether or not the user has their FIDO2 key with them.
            echo "Insert FIDO2 key if using one, then press ENTER..."
            pause
            # --token-only: use enrolled tokens (FIDO2, TPM2) without prompting passphrase
            # If token unlock fails (key not present / not enrolled), fall back to password
            if ! cryptsetup open --token-only "$ROOT_PART" cryptroot 2>/dev/null; then
                echo "Token unlock failed — enter passphrase..."
                cryptsetup open "$ROOT_PART" cryptroot
            fi
            ;;
        2)
            # Traditional passphrase-only unlock — no hardware key needed
            cryptsetup open "$ROOT_PART" cryptroot
            ;;
        3)
            # Token-only mode — will fail if key isn't present/enrolled
            echo "Insert FIDO2 key and press ENTER..."
            pause
            cryptsetup open --token-only "$ROOT_PART" cryptroot
            ;;
        *)
            echo "Invalid choice, using passphrase..."
            cryptsetup open "$ROOT_PART" cryptroot
            ;;
    esac
    # After unlock, the decrypted device is available at /dev/mapper/cryptroot
    MAPPER_DEV="/dev/mapper/cryptroot"
    echo "Partition unlocked."
fi

# ── Detect installed kernel from EFI partition ───────────────────────────────
# WHY: We need to know which kernel to reinstall (linux, linux-lts, linux-zen, etc.).
# HOW: Mount the EFI partition temporarily and look for vmlinuz-* kernel images.
#      The filename suffix is the package name (e.g. vmlinuz-linux-lts → linux-lts).
TMPBOOT=$(mktemp -d /tmp/arch-reset-boot.XXXXXX)
mount "$EFI_PART" "$TMPBOOT"
KERNEL_PKG="linux"  # Default to mainline kernel
for _img in "$TMPBOOT"/vmlinuz-*; do
    # Loop over glob; first match wins (strip the vmlinuz- prefix)
    [[ -f "$_img" ]] && KERNEL_PKG=$(basename "$_img" | sed 's/^vmlinuz-//') && break
done
umount "$TMPBOOT"; rmdir "$TMPBOOT"
echo "Detected kernel: $KERNEL_PKG"

# ── Detect GPU from live hardware ────────────────────────────────────────────
# WHY: Different GPU vendors need different Xorg/DRM driver packages.
#      lspci reads the PCI device list; grep filters for display-related devices
#      (VGA=traditional GPU, 3D=compute GPU/integrated graphics).
# GPU_PKGS will be passed to pacstrap as an extra package.
GPU_INFO=$(lspci 2>/dev/null | grep -E "VGA|3D" || true)
GPU_PKGS=""
if   echo "$GPU_INFO" | grep -qi nvidia; then GPU_PKGS="nvidia-open"    # NVIDIA open-source driver
elif echo "$GPU_INFO" | grep -qi amd;    then GPU_PKGS="xf86-video-amdgpu"  # AMD open-source driver
elif echo "$GPU_INFO" | grep -qi intel;  then GPU_PKGS="xf86-video-intel"   # Intel driver
fi
# If no GPU detected, GPU_PKGS remains empty and the kernel's DRM will handle it

# ── Mount btrfs top-level ────────────────────────────────────────────────────
# WHY: We need to access the raw btrfs filesystem (not via a subvolume) so we
#      can manipulate subvolumes directly.
# HOW: subvolid=5 is always the btrfs top-level (root of all subvolumes).
BTRFS_MNT="$TMPDIR/btrfs"
mkdir -p "$BTRFS_MNT"
mount -o subvolid=5 "$MAPPER_DEV" "$BTRFS_MNT"

# ── Save critical configuration from @ ──────────────────────────────────────
# WHY: We're about to destroy the @ subvolume, which contains /etc.
#      We need to preserve critical system files so the reinstalled system
#      has the same users, passwords, locales, hostname, and boot config.
echo "Saving system configuration..."
SAVED="$TMPDIR/saved"

# _save: helper to copy a file/dir from the old @ subvolume to our save area.
# Skips silently if the source doesn't exist (not all systems have all configs).
_save() {
    local src="$BTRFS_MNT/@/etc/$1"
    local dst="$SAVED/etc/$1"
    [[ -e "$src" ]] || return 0  # Skip if source doesn't exist
    mkdir -p "$(dirname "$dst")"
    cp -a "$src" "$dst"  # -a: archive mode (preserves permissions, timestamps, symlinks)
}

# User and group databases — contain usernames, UIDs, group memberships
_save passwd
_save shadow    # Hashed passwords (mode 000 — only root can read)
_save group
_save gshadow   # Group passwords (mode 000)

# Privilege escalation config
_save sudoers
_save sudoers.d

# PAM (Pluggable Authentication Modules) — contains FIDO2 PAM config
# WHY: Must preserve this so FIDO2 login works after reset
_save pam.d

# System identity and locale
_save hostname
_save locale.conf
_save locale.gen
_save vconsole.conf   # Console keymap and font settings

# Boot configuration
_save fstab           # Mount point definitions (UUIDs etc.)
_save mkinitcpio.conf # Initramfs hooks config (encryption hooks must be preserved)
_save mkinitcpio.conf.d
_save default/grub    # GRUB cmdline (has cryptdevice= or rd.luks.name= parameters)

# Network connection profiles
_save NetworkManager  # Contains /etc/NetworkManager/system-connections/ WiFi passwords

# Save timezone symlink target as plain text because symlinks can't cross roots.
# The symlink /etc/localtime → /usr/share/zoneinfo/Europe/Vienna would be broken
# if we copied the symlink itself into a different root.
{ readlink "$BTRFS_MNT/@/etc/localtime" 2>/dev/null || echo "/usr/share/zoneinfo/UTC"; } \
    > "$SAVED/timezone"

# ── Clear ~/.config in @home (preserve auth-critical subdirs) ────────────────
# WHY: After a system reset, old app configs may be incompatible with the
#      freshly installed versions. Clearing them gives apps a clean start.
# EXCEPTION: Yubico/ and pam-u2f/ contain FIDO2/U2F PAM key files.
#      Deleting these would break FIDO2 login for all users after the reset.
echo "Clearing user app configs..."
PRESERVED_CONFIG_DIRS=("Yubico" "pam-u2f")

# Iterate over all user home directories in the @home subvolume
for _homedir in "$BTRFS_MNT/@home"/*/; do
    [[ -d "$_homedir" ]] || continue  # Skip if glob matched nothing
    _user=$(basename "$_homedir")
    _cfgdir="$_homedir/.config"
    [[ -d "$_cfgdir" ]] || continue   # Skip users with no .config
    echo "  Clearing ~/.config for: $_user"

    # Delete all top-level files in .config (not directories — handled separately)
    # -mindepth 1: don't delete .config itself
    # -maxdepth 1: only direct children
    # ! -type d: only files, not directories
    find "$_cfgdir" -mindepth 1 -maxdepth 1 ! -type d -delete

    # Delete subdirectories, except preserved ones
    # Using process substitution + read with NUL delimiter for paths with spaces
    while IFS= read -r -d '' _subdir; do
        _dname=$(basename "$_subdir")
        _skip=false
        # Check if this directory is in the preserved list
        for _keep in "${PRESERVED_CONFIG_DIRS[@]}"; do
            [[ "$_dname" == "$_keep" ]] && _skip=true && break
        done
        $_skip || rm -rf "$_subdir"  # Delete if not in preserved list
    done < <(find "$_cfgdir" -mindepth 1 -maxdepth 1 -type d -print0)
done

# ── Delete @ and recreate fresh ──────────────────────────────────────────────
# WHY: btrfs subvolume delete removes the entire @ filesystem tree instantly.
#      Creating a new empty @ gives us a clean slate for pacstrap.
echo "Deleting root subvolume @..."
btrfs subvolume delete "$BTRFS_MNT/@"
echo "Creating fresh root subvolume @..."
btrfs subvolume create "$BTRFS_MNT/@"

# ── Mount for installation ───────────────────────────────────────────────────
# Mount the new empty @ subvolume and the existing @home subvolume for pacstrap.
umount "$BTRFS_MNT"
mount -o subvol=@ "$MAPPER_DEV" /mnt
mkdir -p /mnt/home
mount -o subvol=@home "$MAPPER_DEV" /mnt/home
mkdir -p /mnt/boot
mount "$EFI_PART" /mnt/boot

# ── Pacstrap base system ─────────────────────────────────────────────────────
# WHY: pacstrap installs packages into the new root at /mnt.
#      We install only the minimal base needed to boot and manage the system.
#      The full package set is reinstalled via tui-install.sh after first boot.
echo "Reinstalling base system (this will take a while)..."
# shellcheck disable=SC2086 — GPU_PKGS is intentionally word-split here
pacstrap /mnt \
    base base-devel "$KERNEL_PKG" linux-firmware vim zsh git networkmanager grub efibootmgr \
    btrfs-progs cryptsetup libfido2 pam-u2f sudo less jq $GPU_PKGS

# ── Restore saved configuration ──────────────────────────────────────────────
# WHY: The fresh @ has default /etc files from pacstrap. We overwrite them
#      with the saved versions to restore users, passwords, hostname, etc.
echo "Restoring system configuration..."

# _restore: copy from save area into the new /mnt/etc
_restore() {
    local src="$SAVED/etc/$1"
    local dst="/mnt/etc/$1"
    [[ -e "$src" ]] || return 0  # Skip if we didn't save this file
    mkdir -p "$(dirname "$dst")"
    cp -a "$src" "$dst"
}

# Auth files get explicit permission modes with `install` rather than `cp -a`
# because cp -a preserves the original permissions which may be wrong after
# copying across filesystem boundaries.
# passwd/group: world-readable (644), shadow/gshadow: inaccessible to non-root (000)
[[ -f "$SAVED/etc/passwd"  ]] && install -m 644 "$SAVED/etc/passwd"  /mnt/etc/passwd
[[ -f "$SAVED/etc/group"   ]] && install -m 644 "$SAVED/etc/group"   /mnt/etc/group
[[ -f "$SAVED/etc/shadow"  ]] && install -m 000 "$SAVED/etc/shadow"  /mnt/etc/shadow
[[ -f "$SAVED/etc/gshadow" ]] && install -m 000 "$SAVED/etc/gshadow" /mnt/etc/gshadow

# Restore everything else using the _restore helper
_restore sudoers
_restore sudoers.d
_restore pam.d         # FIDO2 PAM rules restored here
_restore hostname
_restore locale.conf
_restore locale.gen
_restore vconsole.conf
_restore fstab         # CRITICAL: without fstab the system won't mount properly on boot
_restore mkinitcpio.conf      # CRITICAL: must have correct encryption hooks
_restore mkinitcpio.conf.d
_restore default/grub  # CRITICAL: must have correct cryptdevice= kernel parameters
_restore NetworkManager

# Restore the timezone symlink using the saved target path
TZ_TARGET=$(cat "$SAVED/timezone")
ln -sf "$TZ_TARGET" /mnt/etc/localtime

# ── Chroot: regenerate initramfs, GRUB menu, services ───────────────────────
# WHY: mkinitcpio and grub-mkconfig must run from inside the new root because
#      they reference files and kernel modules specific to that root.
#      Running them from outside would produce incorrect initramfs/GRUB configs.
echo "Finalizing inside chroot..."
arch-chroot /mnt /bin/bash <<'CHROOT_EOF'
set -euo pipefail

# Regenerate locale files (locale.gen was restored, re-run locale-gen to apply)
locale-gen

# Sync hardware clock from system time
hwclock --systohc

# Enable NetworkManager so the system has network on first boot
systemctl enable NetworkManager

# Regenerate initramfs for ALL installed kernels (-P = all presets)
# WHY: The restored mkinitcpio.conf has the correct hooks (encrypt, btrfs, etc.)
#      but the initramfs itself was in the old @ and is now gone.
mkinitcpio -P

# Regenerate GRUB boot menu
# WHY: GRUB needs a fresh grub.cfg pointing to the new kernel images.
#      The restored /etc/default/grub already has the correct cmdline.
grub-mkconfig -o /boot/grub/grub.cfg

# Fix home directory ownership using the restored UID/GID from /etc/passwd
# WHY: @home was preserved but the new @ has a fresh uid/gid namespace until
#      /etc/passwd is restored. Now that passwd is in place, we re-apply
#      ownership to ensure home dirs are accessible by the right users.
# IFS=: splits /etc/passwd fields by colon; fields: name:pass:uid:gid:gecos:home:shell
while IFS=: read -r _uname _ _uid _gid _ _home _; do
    # Only fix user accounts (UID >= 1000), not system accounts
    (( _uid >= 1000 )) && [[ -d "$_home" ]] && chown -R "${_uid}:${_gid}" "$_home" || true
done < /etc/passwd
CHROOT_EOF

echo ""
echo "======================================="
echo "  Reset complete!"
echo "======================================="
echo "  umount -R /mnt && reboot"