#!/usr/bin/env bash
#
# grub-zfsd
#
# Synex ZFS snapshot/clone monitor for GRUB.
#
# Watches the boot-relevant object set of the canonical ZFS root dataset
# and regenerates /boot/grub/grub-zfs.cfg when direct root snapshots or
# zfs-initramfs RW clones are created or removed.
#
# The visible behaviour intentionally follows grub-btrfsd:
#
#   change detected
#       -> wait 5 seconds
#       -> regenerate snapshot menu
#
# ZFS state is monitored through the ZFS command interface rather than
# inotify on .zfs/snapshot, because the ZFS snapshot namespace does not
# provide the filesystem events required for reliable inotify watching.
#

set -u

CONFIG_FILE="/etc/default/grub-zfs/config"

if [ -r "$CONFIG_FILE" ]; then
    . "$CONFIG_FILE"
fi

POLL_INTERVAL="${GRUB_ZFS_DAEMON_POLL_INTERVAL:-1}"
REFRESH_DELAY="${GRUB_ZFS_DAEMON_REFRESH_DELAY:-5}"


log()
{
    printf 'grub-zfsd: %s\n' "$*"
}


err()
{
    printf 'grub-zfsd: Error: %s\n' "$*" >&2
}


validate_uint()
{
    case "$2" in
        ''|*[!0-9]*)
            err "Invalid $1='$2'"
            return 1
            ;;
    esac

    return 0
}


validate_uint \
    "GRUB_ZFS_DAEMON_POLL_INTERVAL" \
    "$POLL_INTERVAL" \
    || POLL_INTERVAL=1

validate_uint \
    "GRUB_ZFS_DAEMON_REFRESH_DELAY" \
    "$REFRESH_DELAY" \
    || REFRESH_DELAY=5


for command_name in \
    zfs \
    findmnt \
    grub-zfs-refresh
do
    if ! command -v "$command_name" >/dev/null 2>&1; then
        err "Required command not found: $command_name"
        exit 1
    fi
done


current_root="$(
    findmnt -n -o SOURCE / 2>/dev/null || true
)"


if [ -z "$current_root" ]; then
    log "Could not determine the root filesystem; exiting."
    exit 0
fi


if ! zfs list \
    -H \
    -o name \
    "$current_root" \
    >/dev/null 2>&1
then
    log "Root filesystem is not ZFS; exiting."
    exit 0
fi


origin="$(
    zfs get \
        -H \
        -o value \
        origin \
        "$current_root" \
        2>/dev/null \
        || printf '%s' '-'
)"


if [ "$origin" != "-" ] \
    && [[ "$origin" == *@* ]]
then
    root_dataset="${origin%%@*}"
else
    root_dataset="$current_root"
fi


if ! zfs list \
    -H \
    -o name \
    "$root_dataset" \
    >/dev/null 2>&1
then
    err "Canonical root dataset not found: $root_dataset"
    exit 1
fi


root_parent="${root_dataset%/*}"

if [ "$root_parent" = "$root_dataset" ]; then
    err "Unable to determine parent dataset for: $root_dataset"
    exit 1
fi


#
# grub-zfs.cfg currently lives in /boot of the canonical root dataset.
#
# While booted from an OpenZFS temporary snapshot clone, /boot belongs
# to that temporary clone as well. Do not pretend that a refresh there
# would persist.
#
# On the next canonical boot the daemon performs an initial refresh.
#

if [ "$current_root" != "$root_dataset" ]; then
    log "Snapshot clone boot detected: $current_root"
    log "Canonical root dataset: $root_dataset"
    log "Automatic persistent GRUB refresh is suspended for this boot."
    exit 0
fi


boot_object_state()
{
    local dataset
    local dataset_origin
    local dataset_mountpoint
    local dataset_canmount
    local dataset_zol_mountpoint
    local snapshot

    #
    # Direct snapshots of the canonical root dataset.
    #
    # This deliberately includes both Synex-managed and external
    # snapshots. The generator decides how each snapshot is classified
    # and whether it is actually bootable.
    #

    zfs list \
        -H \
        -t snapshot \
        -o name \
        -S creation \
        -r "$root_dataset" \
        2>/dev/null \
    | while IFS= read -r snapshot
    do
        case "$snapshot" in
            "$root_dataset"@*)
                printf 'snapshot\t%s\n' "$snapshot"
                ;;
        esac
    done

    #
    # zfs-initramfs RW clones are sibling filesystems of the canonical
    # root. They are intentionally identified by the properties used by
    # zfs-initramfs, so unrelated ZFSBootMenu/manual Boot Environments
    # are not treated as Synex RW clones.
    #

    while IFS=$'\t' read -r \
        dataset \
        dataset_origin \
        dataset_mountpoint \
        dataset_canmount
    do
        [ -n "$dataset" ] || continue
        [ "$dataset" != "$root_dataset" ] || continue

        case "$dataset_origin" in
            "$root_dataset"@*)
                ;;
            *)
                continue
                ;;
        esac

        [ "$dataset_mountpoint" = "none" ] || continue
        [ "$dataset_canmount" = "noauto" ] || continue

        dataset_zol_mountpoint="$(
            zfs get \
                -H \
                -o value \
                org.zol:mountpoint \
                "$dataset" \
                2>/dev/null \
                || printf '%s' '-'
        )"

        [ "$dataset_zol_mountpoint" = "/" ] || continue

        printf 'clone\t%s\t%s\n' \
            "$dataset" \
            "$dataset_origin"

    done < <(
        zfs list \
            -H \
            -t filesystem \
            -o name,origin,mountpoint,canmount \
            -r "$root_parent" \
            2>/dev/null
    )
}


refresh_menu()
{
    if grub-zfs-refresh; then
        log "GRUB ZFS snapshot menu refreshed."
        return 0
    fi

    err "Could not refresh GRUB ZFS snapshot menu."
    return 1
}


trap '
    log "Received termination signal; exiting."
    exit 0
' INT TERM


log "Starting."
log "Canonical root dataset: $root_dataset"
log "Polling interval: ${POLL_INTERVAL}s"
log "Refresh delay: ${REFRESH_DELAY}s"


#
# Synchronize once at daemon startup.
#
# This also covers a snapshot or clone change that may have happened
# while the daemon was stopped, including a RW clone created during a
# previous snapshot boot.
#

refresh_menu || true


last_state="$(
    boot_object_state
)"


while true
do
    sleep "$POLL_INTERVAL"

    current_state="$(
        boot_object_state
    )"

    if [ "$current_state" = "$last_state" ]; then
        continue
    fi

    log "Detected ZFS snapshot/clone change."
    log "Waiting ${REFRESH_DELAY}s before rebuilding the GRUB snapshot menu."

    sleep "$REFRESH_DELAY"

    #
    # Read state again after the debounce delay. More than one ZFS
    # operation may have happened while Synex Snapshots or zfs-initramfs
    # was working.
    #

    current_state="$(
        boot_object_state
    )"

    refresh_menu || true

    last_state="$current_state"
done
