#!/bin/sh
# rikkiti-rollback <snapshot-number> [--apply]
#
# Roll the SYSTEM (@) back to snapper snapshot N. On our systemd-boot + `subvol=@` layout,
# plain `snapper rollback` is a no-op: the kernel cmdline hardcodes `rootflags=subvol=@`, so
# boot ignores the btrfs default subvolume that snapper flips. Instead we swap @ itself:
# from the running system, mount the btrfs top-level, move the live @ aside as a timestamped
# backup, and create a fresh writable @ from snapshot N's read-only snapshot — so `subvol=@`
# resolves to the snapshot's content on the next boot.
#
# /home (@home), /var/log (@var_log) and the snapshot history (@snapshots) are SEPARATE
# subvolumes and are left untouched — that's the whole reason @snapshots isn't nested in @.
# Reversible: the old system is kept as the backup subvolume (recover by swapping it back).
#
# Without --apply it prints the plan and changes nothing. Must run as root (Settings ▸
# Snapshots calls it via pkexec). Reboot afterwards to apply. docs/44.
set -u

N="${1:-}"; APPLY="${2:-}"
case "$N" in ''|*[!0-9]*) echo "usage: rikkiti-rollback <snapshot-number> [--apply]" >&2; exit 2;; esac
[ "$N" -ge 1 ] 2>/dev/null || { echo "snapshot number must be >= 1" >&2; exit 2; }
[ "$(id -u)" = 0 ] || { echo "must run as root" >&2; exit 1; }

# Root must be btrfs, mounted from subvol @ (never touch a non-@ / — refuse loudly).
[ "$(findmnt -no FSTYPE / 2>/dev/null)" = btrfs ] || { echo "root is not btrfs" >&2; exit 1; }
rootsrc=$(findmnt -no SOURCE / 2>/dev/null)
case "$rootsrc" in *'[/@]'*) ;; *) echo "root is not on subvol @ ($rootsrc) — refusing" >&2; exit 1;; esac
dev=$(printf '%s' "$rootsrc" | sed 's/\[.*//')
uuid=$(findmnt -no UUID / 2>/dev/null)

top=$(mktemp -d /run/rikkiti-rollback.XXXXXX) || exit 1
cleanup() { umount "$top" 2>/dev/null; rmdir "$top" 2>/dev/null; }
trap cleanup EXIT INT TERM

mount -o subvolid=5 "$dev" "$top" 2>/dev/null || { echo "cannot mount btrfs top-level" >&2; exit 1; }

snap="$top/@snapshots/$N/snapshot"
[ -d "$snap" ] || { echo "snapshot $N not found (expected @snapshots/$N/snapshot)" >&2; exit 1; }
[ -d "$top/@" ]  || { echo "@ subvolume missing?!" >&2; exit 1; }

ts=$(date +%Y%m%d-%H%M%S 2>/dev/null); [ -n "$ts" ] || ts="manual"
backup="@rollback-backup-$ts"
newtmp="@rollback-new-$$"

echo "Rollback plan:"
echo "  device         : $dev (UUID=$uuid)"
echo "  roll @ back to : snapshot $N  (@snapshots/$N/snapshot)"
echo "  keep old @ as  : $backup   (recover by swapping it back)"
echo "  untouched      : /home (@home), /var/log (@var_log), history (@snapshots)"
if [ "$APPLY" != "--apply" ]; then
	echo "(dry run — nothing changed; pass --apply to perform, then reboot)"
	exit 0
fi

# 1. Fresh writable @ from the read-only snapshot. Does NOT touch the live @.
if ! btrfs subvolume snapshot "$snap" "$top/$newtmp" >/dev/null 2>&1; then
	echo "creating the new @ failed" >&2
	btrfs subvolume delete "$top/$newtmp" >/dev/null 2>&1
	exit 1
fi

# 2. Swap: old @ -> backup, then new @ -> @. The running system stays on the old subvolume
#    (the mount tracks the subvol root, not its path) until reboot. Brief window where @ is
#    absent — a deliberate, fast rename with the old system preserved as $backup; if the
#    second rename fails, undo.
if ! mv "$top/@" "$top/$backup"; then
	echo "could not move the current @ aside" >&2
	btrfs subvolume delete "$top/$newtmp" >/dev/null 2>&1
	exit 1
fi
if ! mv "$top/$newtmp" "$top/@"; then
	mv "$top/$backup" "$top/@"   # restore
	echo "swap failed — restored the original @; nothing changed" >&2
	exit 1
fi

# CRITICAL: sync the ESP kernel to the rolled-back @. Our kernel + initrd live on the ESP
# (for the Plymouth FDE unlock), NOT inside @ — so after swapping @ the ESP would still boot
# the CURRENT kernel while the rolled-back @ has the OLDER /lib/modules → drivers fail to
# load → broken boot. The snapshot captured its own /boot/{vmlinuz,initrd.img}; copy those
# onto the ESP entry so the booted kernel matches the rolled-back modules. (/boot/efi is the
# shared ESP, mounted in the running system; a later kernel update re-syncs via esp-sync.)
esp=$(findmnt -no TARGET /boot/efi 2>/dev/null); [ -n "$esp" ] || esp=/boot/efi
newroot="$top/@"
kver=$(ls "$newroot/lib/modules" 2>/dev/null | sort -V | tail -1)
if [ -n "$kver" ] && [ -f "$newroot/boot/vmlinuz-$kver" ] && [ -d "$esp/rikkiti" ]; then
	cp -f "$newroot/boot/vmlinuz-$kver" "$esp/rikkiti/vmlinuz" 2>/dev/null || true
	[ -f "$newroot/boot/initrd.img-$kver" ] && cp -f "$newroot/boot/initrd.img-$kver" "$esp/rikkiti/initrd" 2>/dev/null || true
	sync
else
	echo "warn: could not sync the ESP kernel for $kver — rollback across a kernel change may not boot" >&2
fi

# Breadcrumb (support + a future 'roll forward'): written into the NEW @.
mkdir -p "$top/@/var/lib/rikkiti" 2>/dev/null || true
printf 'rolled_back_to=%s\nfrom_backup=%s\nkernel=%s\nwhen=%s\n' "$N" "$backup" "$kver" "$ts" \
	> "$top/@/var/lib/rikkiti/last-rollback" 2>/dev/null || true

echo "ROLLED_BACK to snapshot $N (kernel $kver). Old system kept as top-level:/$backup. Reboot to apply."
exit 0
