#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-3.0-only
# eniatos-snap - ZFS snapshots of the running boot environment (bootable from ZFSBootMenu)
set -uo pipefail

PROG=eniatos-snap
CONF=/etc/eniatos/snap.conf
MAX_AGE_DAYS=30
KEEP_MIN=3
DATASETS=""
# shellcheck source=/dev/null
[[ -r $CONF ]] && . "$CONF"

usage() {
    cat <<'EOF'
Usage: eniatos-snap [--permanent] [LABEL]   snapshot the running boot environment
       eniatos-snap --auto LABEL             snapshot only if changed since the last one
       eniatos-snap --hook                   pacman PreTransaction hook (targets on stdin)
       eniatos-snap --list                   list snapshots of the boot environment
       eniatos-snap --prune [--dry-run]      destroy expired snapshots

Snapshot name: <boot-env>@YYYY-MM-DD_HH:MM[_LABEL]
  --permanent  sets eniatos:keep=yes, --prune never destroys it
Config: /etc/eniatos/snap.conf (MAX_AGE_DAYS, KEEP_MIN, DATASETS)
EOF
    exit "${1:-0}"
}

HOOK=0
err() {
    echo "$PROG: error: $*" >&2
    logger -t "$PROG" -p user.err -- "$*" 2>/dev/null || true
}
info() { echo "$PROG: $*"; logger -t "$PROG" -p user.info -- "$*" 2>/dev/null || true; }
die() {
    err "$*"
    # never abort a pacman transaction
    [[ $HOOK -eq 1 ]] && exit 0
    exit 1
}

root_dataset() {
    local src fstype
    read -r src fstype < <(findmnt -n -o SOURCE,FSTYPE /)
    [[ $fstype == zfs ]] || return 1
    echo "$src"
}

sanitize_label() {
    local l=$1
    l=$(printf '%s' "$l" | tr -c 'A-Za-z0-9._:-' '-' | tr -s '-')
    l=${l#-}; l=${l%-}
    printf '%s' "${l:0:64}"
}

# newest snapshot of $1 created by eniatos-snap (short name)
last_snap() {
    zfs list -H -p -t snapshot -d 1 -s creation -o name,eniatos:snap "$1" 2>/dev/null \
        | awk -F'\t' '$2 != "-" { n = $1 } END { if (n) { sub(/^[^@]*@/, "", n); print n } }'
}

do_snapshot() {
    local kind=$1 label=$2 permanent=$3
    local be ds stamp name n=1 last written
    be=$(root_dataset) || die "/ is not on ZFS"
    local -a targets=("$be")
    for ds in $DATASETS; do
        zfs list -H -o name "$ds" >/dev/null 2>&1 || die "dataset not found: $ds"
        [[ $ds == "$be" ]] || targets+=("$ds")
    done

    if [[ $kind == auto ]]; then
        last=$(last_snap "$be")
        if [[ -n $last ]]; then
            written=$(zfs get -H -p -o value "written@$last" "$be" 2>/dev/null)
            if [[ $written == 0 ]]; then
                info "no changes since $be@$last, skipping"
                return 0
            fi
        fi
    fi

    stamp=$(date +%Y-%m-%d_%H:%M)
    label=$(sanitize_label "$label")
    name=$stamp${label:+_$label}
    local base=$name
    while zfs list -H -o name "$be@$name" >/dev/null 2>&1; do
        n=$((n + 1)); name=$base-$n
    done

    local -a props=(-o "eniatos:snap=$kind")
    [[ $permanent -eq 1 ]] && props+=(-o eniatos:keep=yes)
    local -a snaps=()
    for ds in "${targets[@]}"; do snaps+=("$ds@$name"); done

    local out
    if ! out=$(zfs snapshot "${props[@]}" "${snaps[@]}" 2>&1); then
        die "zfs snapshot ${snaps[*]} failed: $out"
    fi
    local suffix=""
    [[ $permanent -eq 1 ]] && suffix=" (permanent)"
    info "created ${snaps[*]}$suffix"
}

do_hook() {
    HOOK=1
    local -a pkgs=()
    local p
    while IFS= read -r p; do [[ -n $p ]] && pkgs+=("$p"); done
    local label=pacman
    if [[ ${#pkgs[@]} -gt 0 && ${#pkgs[@]} -le 3 ]]; then
        label+="_$(IFS=_; echo "${pkgs[*]}")"
    elif [[ ${#pkgs[@]} -gt 3 ]]; then
        label+="-${#pkgs[@]}pkgs"
    fi
    if [[ ${ENIATOS_SNAP:-1} == 0 ]]; then
        info "ENIATOS_SNAP=0, skipping snapshot"
        exit 0
    fi
    do_snapshot auto "$label" 0
    exit 0
}

do_list() {
    local be
    be=$(root_dataset) || die "/ is not on ZFS"
    zfs list -t snapshot -d 1 -s creation \
        -o name,creation,used,refer,eniatos:snap,eniatos:keep "$be"
}

do_prune() {
    local dry=$1 now cutoff
    now=$(date +%s)
    cutoff=$((now - MAX_AGE_DAYS * 86400))
    # all snapshots made by eniatos-snap, any dataset, oldest first
    local list
    list=$(zfs list -H -p -t snapshot -s creation \
        -o name,creation,eniatos:snap,eniatos:keep,clones,userrefs 2>/dev/null \
        | awk -F'\t' '$3 != "-"') || true
    [[ -z $list ]] && { info "nothing to prune"; return 0; }

    local rc=0 name creation kind keep clones refs ds total idx
    declare -A count seen
    while IFS=$'\t' read -r name creation kind keep clones refs; do
        [[ $keep == yes ]] && continue                  # permanent: not counted in KEEP_MIN
        ds=${name%%@*}; count[$ds]=$(( ${count[$ds]:-0} + 1 ))
    done <<< "$list"

    while IFS=$'\t' read -r name creation kind keep clones refs; do
        [[ $keep == yes ]] && continue
        ds=${name%%@*}
        seen[$ds]=$(( ${seen[$ds]:-0} + 1 ))
        idx=${seen[$ds]}; total=${count[$ds]}
        (( creation < cutoff )) || continue
        (( total - idx < KEEP_MIN )) && continue        # keep the newest KEEP_MIN
        [[ $clones != - && -n $clones ]] && continue    # origin of a boot environment
        (( refs > 0 )) && continue                      # zfs hold
        if [[ $dry -eq 1 ]]; then
            echo "would destroy $name"
        elif out=$(zfs destroy "$name" 2>&1); then
            info "destroyed $name"
        else
            err "zfs destroy $name failed: $out"; rc=1
        fi
    done <<< "$list"
    [[ $dry -eq 1 ]] || info "prune done (MAX_AGE_DAYS=$MAX_AGE_DAYS, KEEP_MIN=$KEEP_MIN)"
    return $rc
}

MODE=snap KIND=manual PERMANENT=0 DRY=0 LABEL=""
while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--help) usage 0 ;;
        --permanent) PERMANENT=1; shift ;;
        --auto) KIND=auto; shift ;;
        --hook) MODE=hook; shift ;;
        --list) MODE=list; shift ;;
        --prune) MODE=prune; shift ;;
        --dry-run) DRY=1; shift ;;
        --) shift; break ;;
        -*) echo "unknown option: $1" >&2; usage 1 ;;
        *) break ;;
    esac
done
LABEL="$*"
[[ $KIND == auto && $PERMANENT -eq 1 ]] && { echo "--auto and --permanent are exclusive" >&2; exit 1; }
[[ $MODE == hook ]] && HOOK=1
[[ $(id -u) -eq 0 || $MODE == list ]] || die "must be root"
[[ $MAX_AGE_DAYS =~ ^[0-9]+$ && $KEEP_MIN =~ ^[0-9]+$ ]] || die "invalid MAX_AGE_DAYS/KEEP_MIN in $CONF"

case $MODE in
    snap)  do_snapshot "$KIND" "$LABEL" "$PERMANENT" ;;
    hook)  do_hook ;;
    list)  do_list ;;
    prune) do_prune "$DRY" ;;
esac
