#!/bin/sh
# arctic-snapshot - take and manage btrfs snapshots of the root subvolume.
#
#   arctic-snapshot create [description]   take one now
#   arctic-snapshot list                   show what exists
#   arctic-snapshot delete <name>          remove one
#   arctic-snapshot rollback <name>        put the system back to that snapshot
#   arctic-snapshot prune [n]              keep only the newest n (default 20)
#
# alpm calls "create" before every install and every removal, so there is always
# a way back from a package change. That is configured in /etc/alpm/alpm.conf.
#
# If the real timeshift or snapper is installed, this hands over to it so the
# snapshots show up where those tools expect. Otherwise it drives btrfs
# directly - which is all either of them does underneath - so snapshots work on
# a fresh install with nothing extra on it.
# shellcheck shell=sh

set -u

CONF=/etc/arctic/snapshot.conf
TOOL=timeshift
SUBVOL_DIR=/.snapshots
KEEP=20
[ -f "$CONF" ] && . "$CONF"

if [ -t 1 ]; then
	R=$(printf '\033[0m'); B=$(printf '\033[1m')
	TEAL=$(printf '\033[38;5;44m'); GREY=$(printf '\033[38;5;245m')
	AMB=$(printf '\033[38;5;215m'); RED=$(printf '\033[38;5;203m')
else
	R= ; B= ; TEAL= ; GREY= ; AMB= ; RED=
fi

die() { printf '%serror:%s %s\n' "$RED" "$R" "$*" >&2; exit 1; }
note() { printf '%s::%s %s\n' "$TEAL" "$R" "$*"; }

root_is_btrfs() {
	[ "$(stat -f -c %T / 2>/dev/null)" = "btrfs" ] && return 0
	awk '$2=="/" && $3=="btrfs"{found=1} END{exit !found}' /proc/mounts
}

need_btrfs() {
	root_is_btrfs || {
		printf '%s/ is not btrfs - snapshots are a btrfs feature.%s\n' "$GREY" "$R"
		exit 0
	}
}

# Hand over to the real tool when it is installed.
delegate() {
	case "$TOOL" in
	timeshift)
		command -v timeshift >/dev/null 2>&1 || return 1
		case "$1" in
		create) timeshift --create --btrfs --comments "${2:-arctic}" --scripted ;;
		list)   timeshift --list ;;
		*)      return 1 ;;
		esac
		return 0 ;;
	snapper)
		command -v snapper >/dev/null 2>&1 || return 1
		case "$1" in
		create) snapper -c root create --description "${2:-arctic}" ;;
		list)   snapper -c root list ;;
		*)      return 1 ;;
		esac
		return 0 ;;
	esac
	return 1
}

snap_name() { printf 'arctic-%s' "$(date '+%Y%m%d-%H%M%S')"; }

cmd_create() {
	need_btrfs
	desc=${1:-manual}
	delegate create "$desc" && return 0

	[ -d "$SUBVOL_DIR" ] || mkdir -p "$SUBVOL_DIR"
	name=$(snap_name)
	target="$SUBVOL_DIR/$name"
	# Read-only, because a snapshot you can edit is not a snapshot.
	if btrfs subvolume snapshot -r / "$target" >/dev/null 2>&1; then
		printf '%s\n' "$desc" >"$target.description" 2>/dev/null || :
		note "snapshot $name  ($desc)"
		cmd_prune "$KEEP" quiet
		return 0
	fi
	printf '%swarning:%s could not create a snapshot; continuing anyway\n' "$AMB" "$R" >&2
	return 0
}

cmd_list() {
	need_btrfs
	delegate list && return 0
	[ -d "$SUBVOL_DIR" ] || { printf '  no snapshots yet\n'; return 0; }
	printf '\n  %s%-26s %-20s %s%s\n' "$B" "SNAPSHOT" "TAKEN" "DESCRIPTION" "$R"
	for d in "$SUBVOL_DIR"/arctic-*; do
		[ -d "$d" ] || continue
		n=$(basename "$d")
		when=$(printf '%s' "$n" | sed 's/arctic-//; s/\(....\)\(..\)\(..\)-\(..\)\(..\)\(..\)/\1-\2-\3 \4:\5:\6/')
		desc=$(cat "$d.description" 2>/dev/null || echo '-')
		printf '  %-26s %-20s %s\n' "$n" "$when" "$desc"
	done
	printf '\n'
}

cmd_delete() {
	need_btrfs
	n=${1:?usage: arctic-snapshot delete <name>}
	[ -d "$SUBVOL_DIR/$n" ] || die "no such snapshot: $n"
	btrfs subvolume delete "$SUBVOL_DIR/$n" >/dev/null 2>&1 || die "could not delete $n"
	rm -f "$SUBVOL_DIR/$n.description"
	note "deleted $n"
}

cmd_prune() {
	need_btrfs
	keep=${1:-$KEEP}
	quiet=${2:-}
	n=0
	# Newest first; anything past the limit goes.
	for d in $(ls -1d "$SUBVOL_DIR"/arctic-* 2>/dev/null | sort -r); do
		n=$((n+1))
		[ "$n" -le "$keep" ] && continue
		btrfs subvolume delete "$d" >/dev/null 2>&1 && rm -f "$d.description"
		[ -n "$quiet" ] || note "pruned $(basename "$d")"
	done
}

cmd_rollback() {
	need_btrfs
	n=${1:?usage: arctic-snapshot rollback <name>}
	[ -d "$SUBVOL_DIR/$n" ] || die "no such snapshot: $n"
	printf '\n%sThis replaces the running root subvolume with %s.%s\n' "$AMB$B" "$n" "$R"
	printf 'The current root is kept as a snapshot first, so this is reversible.\n'
	printf 'Reboot is required afterwards.\n\n'
	printf 'Continue? [y/N] '
	read -r a || a=n
	case "$a" in y|Y|yes) ;; *) echo "nothing was changed."; exit 0 ;; esac

	# Snapshot the current state before replacing it, so a bad rollback is
	# itself recoverable.
	pre="$SUBVOL_DIR/pre-rollback-$(date '+%Y%m%d-%H%M%S')"
	btrfs subvolume snapshot -r / "$pre" >/dev/null 2>&1 || :

	# The root subvolume has to be swapped from outside itself: mount the
	# top-level subvolume, move @ aside and put the snapshot in its place.
	dev=$(awk '$2=="/"{print $1; exit}' /proc/mounts)
	tmp=$(mktemp -d)
	mount -o subvolid=5 "$dev" "$tmp" || die "cannot mount the btrfs top level"
    if [ -d "$tmp/@" ]; then
		mv "$tmp/@" "$tmp/@.rollback-$(date '+%s')" || die "cannot move the old root aside"
		btrfs subvolume snapshot "$SUBVOL_DIR/$n" "$tmp/@" >/dev/null 2>&1 \
			|| die "cannot put the snapshot in place"
	else
		umount "$tmp"; rmdir "$tmp"
		die "no @ subvolume - this system does not use the Arctic btrfs layout"
	fi
	umount "$tmp"; rmdir "$tmp"
	note "rolled back to $n - reboot to use it"
	note "the previous root was kept as $(basename "$pre")"
}

case "${1:-list}" in
create)   shift; cmd_create "$@" ;;
list|ls)  cmd_list ;;
delete|rm) shift; cmd_delete "$@" ;;
prune)    shift; cmd_prune "$@" ;;
rollback) shift; cmd_rollback "$@" ;;
*)        sed -n '2,16p' "$0" | sed 's/^# \?//'; exit 1 ;;
esac
