#!/usr/bin/env bash
# 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

TMPDIR=$(mktemp -d /tmp/arch-reset.XXXXXX)
trap 'rm -rf "$TMPDIR"' EXIT

pause() { read -rp "Press ENTER to continue..."; }

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 ──────────────────────────────────────
pacman -Syd --noconfirm cryptsetup btrfs-progs jq libfido2

# ── Drive selection ──────────────────────────────────────────────────────────
lsblk
echo ""
read -rp "Enter drive to reset (e.g., /dev/sda): " DRIVE
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 ""
read -rp "Type YES to continue: " _CONFIRM
[[ "$_CONFIRM" == "YES" ]] || { echo "Aborted."; exit 1; }

# ── LUKS detection and unlock ────────────────────────────────────────────────
MAPPER_DEV="$ROOT_PART"

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}"

    case "$_UNLOCK" in
        1)
            echo "Insert FIDO2 key if using one, then press ENTER..."
            pause
            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)
            cryptsetup open "$ROOT_PART" cryptroot
            ;;
        3)
            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
    MAPPER_DEV="/dev/mapper/cryptroot"
    echo "Partition unlocked."
fi

# ── Detect installed kernel from EFI partition ───────────────────────────────
TMPBOOT=$(mktemp -d /tmp/arch-reset-boot.XXXXXX)
mount "$EFI_PART" "$TMPBOOT"
KERNEL_PKG="linux"
for _img in "$TMPBOOT"/vmlinuz-*; do
    [[ -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 ────────────────────────────────────────────
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"
elif echo "$GPU_INFO" | grep -qi amd;    then GPU_PKGS="xf86-video-amdgpu"
elif echo "$GPU_INFO" | grep -qi intel;  then GPU_PKGS="xf86-video-intel"
fi

# ── Mount btrfs top-level ────────────────────────────────────────────────────
BTRFS_MNT="$TMPDIR/btrfs"
mkdir -p "$BTRFS_MNT"
mount -o subvolid=5 "$MAPPER_DEV" "$BTRFS_MNT"

# ── Save critical configuration from @ ──────────────────────────────────────
echo "Saving system configuration..."
SAVED="$TMPDIR/saved"

_save() {
    local src="$BTRFS_MNT/@/etc/$1"
    local dst="$SAVED/etc/$1"
    [[ -e "$src" ]] || return 0
    mkdir -p "$(dirname "$dst")"
    cp -a "$src" "$dst"
}

_save passwd
_save shadow
_save group
_save gshadow
_save sudoers
_save sudoers.d
_save pam.d
_save hostname
_save locale.conf
_save locale.gen
_save vconsole.conf
_save fstab
_save mkinitcpio.conf
_save mkinitcpio.conf.d
_save default/grub
_save NetworkManager

# Save timezone symlink target as plain text (the symlink itself can't cross roots)
{ readlink "$BTRFS_MNT/@/etc/localtime" 2>/dev/null || echo "/usr/share/zoneinfo/UTC"; } \
    > "$SAVED/timezone"

# ── Clear ~/.config in @home (preserve auth-critical subdirs) ────────────────
echo "Clearing user app configs..."
# Yubico/ holds U2F/FIDO2 PAM keys — deleting these would break FIDO2 login
PRESERVED_CONFIG_DIRS=("Yubico" "pam-u2f")

for _homedir in "$BTRFS_MNT/@home"/*/; do
    [[ -d "$_homedir" ]] || continue
    _user=$(basename "$_homedir")
    _cfgdir="$_homedir/.config"
    [[ -d "$_cfgdir" ]] || continue
    echo "  Clearing ~/.config for: $_user"
    # Top-level files in .config
    find "$_cfgdir" -mindepth 1 -maxdepth 1 ! -type d -delete
    # Subdirectories, skipping preserved ones
    while IFS= read -r -d '' _subdir; do
        _dname=$(basename "$_subdir")
        _skip=false
        for _keep in "${PRESERVED_CONFIG_DIRS[@]}"; do
            [[ "$_dname" == "$_keep" ]] && _skip=true && break
        done
        $_skip || rm -rf "$_subdir"
    done < <(find "$_cfgdir" -mindepth 1 -maxdepth 1 -type d -print0)
done

# ── Delete @ and recreate fresh ──────────────────────────────────────────────
echo "Deleting root subvolume @..."
btrfs subvolume delete "$BTRFS_MNT/@"
echo "Creating fresh root subvolume @..."
btrfs subvolume create "$BTRFS_MNT/@"

# ── Mount for installation ───────────────────────────────────────────────────
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 ─────────────────────────────────────────────────────
echo "Reinstalling base system (this will take a while)..."
# shellcheck disable=SC2086
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 ──────────────────────────────────────────────
echo "Restoring system configuration..."

_restore() {
    local src="$SAVED/etc/$1"
    local dst="/mnt/etc/$1"
    [[ -e "$src" ]] || return 0
    mkdir -p "$(dirname "$dst")"
    cp -a "$src" "$dst"
}

# Auth files — explicit permissions
[[ -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 sudoers
_restore sudoers.d
_restore pam.d
_restore hostname
_restore locale.conf
_restore locale.gen
_restore vconsole.conf
_restore fstab
_restore mkinitcpio.conf
_restore mkinitcpio.conf.d
_restore default/grub
_restore NetworkManager

TZ_TARGET=$(cat "$SAVED/timezone")
ln -sf "$TZ_TARGET" /mnt/etc/localtime

# ── Chroot: regenerate initramfs, GRUB menu, services ───────────────────────
echo "Finalizing inside chroot..."
arch-chroot /mnt /bin/bash <<'CHROOT_EOF'
set -euo pipefail

locale-gen
hwclock --systohc
systemctl enable NetworkManager
mkinitcpio -P
grub-mkconfig -o /boot/grub/grub.cfg

# Re-apply correct ownership for user home directories using restored UIDs
while IFS=: read -r _uname _ _uid _gid _ _home _; do
    (( _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"