#!/usr/bin/env bash
#
# Synex GRUB ZFS snapshot menu generator
#
# Generates /boot/grub/grub-zfs.cfg from bootable ZFS objects related to
# the canonical Synex root dataset:
#
#   - Synex-managed snapshots
#   - external snapshots
#   - temporary RW clones created by zfs-initramfs snapshot boot
#
# Snapshot entries load kernel and initrd directly from the snapshot.
# RW clone entries load kernel and initrd directly from the clone dataset.
#
# grub-zfs does not rename, promote or otherwise modify these objects.
#

set -e
set -o pipefail

VERSION="0.1.0"

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

[ -f "$CONFIG_FILE" ] && . "$CONFIG_FILE"
[ -f "$GRUB_DEFAULT_FILE" ] && . "$GRUB_DEFAULT_FILE"


print_info()
{
    printf 'grub-zfs: %s\n' "$*" >&2
}


print_warning()
{
    printf 'grub-zfs: Warning: %s\n' "$*" >&2
}


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


refresh_only=false

case "${1:-}" in
    "")
        ;;

    --refresh-only)
        refresh_only=true
        ;;

    -V|--version)
        printf 'grub-zfs %s\n' "$VERSION"
        exit 0
        ;;

    -h|--help)
        cat <<'HELP'
Usage:
  41_snapshots-zfs
  41_snapshots-zfs --refresh-only
  41_snapshots-zfs --version

Without options, refreshes grub-zfs.cfg and emits the GRUB submenu hook.
With --refresh-only, only refreshes grub-zfs.cfg.
HELP
        exit 0
        ;;

    *)
        print_error "Unknown argument: ${1}"
        exit 2
        ;;
esac


case "$(printf '%s' "${GRUB_ZFS_DISABLE:-false}" | tr '[:upper:]' '[:lower:]')" in
    true|yes|1)
        print_info "Disabled by GRUB_ZFS_DISABLE."
        exit 0
        ;;
esac


for command_name in \
    zfs \
    findmnt \
    flock \
    mount \
    umount \
    grub-probe \
    grub-script-check \
    date
do
    if ! command -v "$command_name" >/dev/null 2>&1; then
        print_warning "Required command not found: $command_name"
        exit 0
    fi
done


mkconfig_lib="${GRUB_ZFS_MKCONFIG_LIB:-/usr/share/grub/grub-mkconfig_lib}"

if [ ! -r "$mkconfig_lib" ]; then
    print_warning "GRUB helper library not found: $mkconfig_lib"
    exit 0
fi

. "$mkconfig_lib"


grub_directory="${GRUB_ZFS_GRUB_DIRNAME:-/boot/grub}"
grub_cfg="${grub_directory}/grub-zfs.cfg"
grub_new="${grub_directory}/grub-zfs.new"


submenu_name="${GRUB_ZFS_SUBMENUNAME:-Synex ZFS Snapshots}"
snapshot_limit="${GRUB_ZFS_LIMIT:-50}"


case "$snapshot_limit" in
    ''|*[!0-9]*)
        print_warning \
            "Invalid GRUB_ZFS_LIMIT='$snapshot_limit'; using 50."

        snapshot_limit=50
        ;;
esac


emit_submenu_hook()
{
    [ "$refresh_only" = true ] && return 0

    cat <<EOF
if [ -e "\${prefix}/grub-zfs.cfg" ]; then
submenu '${submenu_name}' --class snapshots --class synex --class gnu-linux --class gnu --class os {
    configfile "\${prefix}/grub-zfs.cfg"
}
fi
EOF
}


if [ ! -d "$grub_directory" ]; then
    print_warning "GRUB directory not found: $grub_directory"
    exit 0
fi


#
# Prevent update-grub, Synex Snapshots and manual refreshes from
# generating grub-zfs.cfg concurrently.
#

mkdir -p /run/lock

exec 9>/run/lock/grub-zfs.lock
flock 9


#
# Detect the currently running ZFS dataset.
#
# Normal boot:
#
#   zpsynex/ROOT/synex/root
#
# Snapshot boot:
#
#   zpsynex/ROOT/synex/root_<snapshot>
#
# In the latter case, origin points back to:
#
#   zpsynex/ROOT/synex/root@<snapshot>
#
# and therefore allows us to recover the canonical root dataset.
#

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

if [ -z "$current_root" ] \
    || ! zfs list -H -o name "$current_root" >/dev/null 2>&1
then
    print_info "Root filesystem is not ZFS; skipping."
    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
    print_warning \
        "Canonical root dataset not found: $root_dataset"

    exit 0
fi


pool="${root_dataset%%/*}"
dataset_relative="${root_dataset#*/}"


if [ "$pool" = "$root_dataset" ] \
    || [ -z "$dataset_relative" ]
then
    print_warning \
        "Could not derive pool and dataset from: $root_dataset"

    exit 0
fi


#
# ZFS Snapshot Boot v1 deliberately requires /boot to live inside
# the root dataset.
#
# /boot/efi may remain a separate EFI System Partition.
#
# If /boot itself is a separate filesystem, the snapshot does not
# contain a guaranteed matching historical kernel/initramfs set.
# In that case grub-zfs does not publish the snapshot as bootable.
#

if findmnt -rn -o TARGET | grep -Fxq '/boot'; then
    rm -f \
        "$grub_new" \
        "$grub_cfg"

    print_warning \
        "Separate /boot detected. ZFS snapshot boot entries are not generated."

    emit_submenu_hook
    exit 0
fi


#
# Locate the canonical dataset so its .zfs/snapshot namespace can be
# inspected.
#
# When the canonical root itself is running, it is already mounted.
#
# When Synex is currently booted from an OpenZFS temporary clone, the
# canonical root normally is not mounted. Mount it read-only temporarily.
#

cleanup_mount=false
inspection_root=""

canonical_target="$(
    findmnt \
        -rn \
        -S "$root_dataset" \
        -o TARGET \
        2>/dev/null \
        | head -n1 \
        || true
)"


if [ -n "$canonical_target" ]; then
    inspection_root="$canonical_target"
else
    inspection_root="$(
        mktemp -d /run/grub-zfs-root.XXXXXX
    )"

    if mount \
        -t zfs \
        -o ro,zfsutil \
        "$root_dataset" \
        "$inspection_root"
    then
        cleanup_mount=true
    else
        rmdir "$inspection_root" 2>/dev/null || true

        print_warning \
            "Could not mount canonical root dataset for snapshot inspection: $root_dataset"

        exit 0
    fi
fi


cleanup()
{
    rm -f "$grub_new"

    if [ "$cleanup_mount" = true ]; then
        umount "$inspection_root" >/dev/null 2>&1 || true
        rmdir "$inspection_root" >/dev/null 2>&1 || true
    fi
}


trap cleanup EXIT


snapshot_control_dir="${inspection_root%/}/.zfs/snapshot"


if [ ! -d "$snapshot_control_dir" ]; then
    print_warning \
        "ZFS snapshot control directory is not accessible: $snapshot_control_dir"

    rm -f "$grub_cfg"

    emit_submenu_hook
    exit 0
fi


#
# Ask GRUB itself how to access the underlying device.
#
# This generates the same kind of insmod/search stanza used by
# Debian's normal GRUB Linux entries.
#

root_device="$(
    "${grub_probe}" \
        --target=device \
        / \
        2>/dev/null \
        | head -n1 \
        || true
)"

prepare_root_cache=""


if [ -n "$root_device" ]; then
    prepare_root_cache="$(
        prepare_grub_to_access_device \
            "$root_device" \
            2>/dev/null \
            || true
    )"
fi


#
# Conservative fallback for the normal Synex GPT + ZFS layout.
#

if [ -z "$prepare_root_cache" ]; then
    root_uuid="$(
        "${grub_probe}" \
            --target=fs_uuid \
            / \
            2>/dev/null \
            || true
    )"

    if [ -z "$root_uuid" ]; then
        print_warning \
            "Could not determine GRUB filesystem UUID for root."

        exit 0
    fi

    prepare_root_cache="$(
        cat <<EOF
insmod part_gpt
insmod zfs
search --no-floppy --fs-uuid --set=root ${root_uuid}
EOF
    )"
fi


#
# Preserve the normal Synex GRUB kernel options, but never inherit a
# root= or zfs= selector. Each generated entry supplies the correct
# root=ZFS= and zfs= selectors for its own snapshot or clone.
#

base_extra_args=""


for arg in \
    ${GRUB_CMDLINE_LINUX:-} \
    ${GRUB_CMDLINE_LINUX_DEFAULT:-}
do
    case "$arg" in
        root=*|zfs=*)
            ;;

        *)
            if [ -n "$base_extra_args" ]; then
                base_extra_args="${base_extra_args} ${arg}"
            else
                base_extra_args="$arg"
            fi
            ;;
    esac
done


format_created_at()
{
    local value="$1"
    local date_part
    local time_part

    case "$value" in
        ????-??-??_??-??-??)
            date_part="${value%%_*}"
            time_part="${value#*_}"
            time_part="${time_part//-/:}"

            printf '%s %s' \
                "$date_part" \
                "$time_part"
            ;;

        *)
            printf '%s' "$value"
            ;;
    esac
}


format_purpose()
{
    case "$1" in
        manual)
            printf 'Manual'
            ;;

        auto)
            printf 'Automatic'
            ;;

        pre-restore)
            printf 'Pre-Restore'
            ;;

        *)
            printf '%s' "$1"
            ;;
    esac
}


sanitize_id()
{
    printf '%s' "$1" \
        | tr -c 'A-Za-z0-9_.-' '_'
}


#
# Helpers for bootable objects.
#

format_zfs_creation()
{
    local object="$1"
    local epoch

    epoch="$(
        zfs get \
            -H \
            -p \
            -o value \
            creation \
            "$object" \
            2>/dev/null \
            || true
    )"

    case "$epoch" in
        ''|*[!0-9]*)
            printf '%s' '-'
            ;;

        *)
            date \
                -d "@${epoch}" \
                '+%Y-%m-%d %H:%M:%S' \
                2>/dev/null \
                || printf '%s' "$epoch"
            ;;
    esac
}


kernel_versions_for_dir()
{
    local content_dir="$1"
    local kernel
    local version

    for kernel in "$content_dir"/boot/vmlinuz-*
    do
        [ -e "$kernel" ] || continue

        version="${kernel##*/}"
        version="${version#vmlinuz-}"

        [ -e "$content_dir/boot/initrd.img-${version}" ] \
            || continue

        [ -d "$content_dir/lib/modules/${version}" ] \
            || continue

        printf '%s\n' "$version"
    done \
        | sort -Vr
}


emit_boot_object()
{
    local object_class="$1"
    local object_title="$2"
    local content_dir="$3"
    local grub_dataset_path="$4"
    local root_spec="$5"
    local zfs_selector="$6"
    local object_id="$7"
    local source_description="$8"
    local target_extra_args
    local version
    local kernel_grub
    local initrd_grub
    local menu_id
    local -a versions

    mapfile -t versions < <(
        kernel_versions_for_dir "$content_dir"
    )

    if [ "${#versions[@]}" -eq 0 ]; then
        print_warning \
            "No complete kernel/initrd/modules set found in ${source_description}"

        return 1
    fi

    if [ -n "$base_extra_args" ]; then
        target_extra_args="zfs=${zfs_selector} ${base_extra_args}"
    else
        target_extra_args="zfs=${zfs_selector}"
    fi

    printf "submenu '%s' {\n" \
        "$object_title" \
        >> "$grub_new"

    for version in "${versions[@]}"
    do
        kernel_grub="${grub_dataset_path}/boot/vmlinuz-${version}"
        initrd_grub="${grub_dataset_path}/boot/initrd.img-${version}"

        menu_id="$(
            sanitize_id \
                "synex-zfs-${object_id}-${version}"
        )"

        printf \
            "    menuentry 'Linux %s' --class snapshots --class synex --class gnu-linux --class gnu --class os \$menuentry_id_option '%s' {\n" \
            "$version" \
            "$menu_id" \
            >> "$grub_new"

        printf '%s\n' \
            '        load_video' \
            >> "$grub_new"

        printf '%s\n' \
            '        insmod gzio' \
            >> "$grub_new"

        printf '%s\n' "$prepare_root_cache" \
            | sed 's/^/        /' \
            >> "$grub_new"

        printf \
            "        echo 'Loading %s: %s'\n" \
            "$object_class" \
            "$source_description" \
            >> "$grub_new"

        printf \
            "        echo 'Loading Linux %s ...'\n" \
            "$version" \
            >> "$grub_new"

        printf \
            '        linux "%s" root=ZFS=%s ro %s\n' \
            "$kernel_grub" \
            "$root_spec" \
            "$target_extra_args" \
            >> "$grub_new"

        printf \
            "        echo 'Loading initial ramdisk ...'\n" \
            >> "$grub_new"

        printf \
            '        initrd "%s"\n' \
            "$initrd_grub" \
            >> "$grub_new"

        printf '%s\n' \
            '    }' \
            >> "$grub_new"

        printf '\n' \
            >> "$grub_new"

        count_entries=$((count_entries + 1))
    done

    printf '%s\n' \
        '}' \
        >> "$grub_new"

    printf '\n' \
        >> "$grub_new"

    return 0
}


#
# Mount an unmounted filesystem dataset read-only so its /boot and
# /lib/modules can be inspected without changing its ZFS properties.
#
# The caller must invoke release_object_mount afterwards.
#

object_mount=""
object_mount_temporary=false


prepare_object_mount()
{
    local dataset="$1"
    local target

    object_mount=""
    object_mount_temporary=false

    target="$(
        findmnt \
            -rn \
            -S "$dataset" \
            -o TARGET \
            2>/dev/null \
            | head -n1 \
            || true
    )"

    if [ -n "$target" ]; then
        object_mount="$target"
        return 0
    fi

    object_mount="$(
        mktemp -d /run/grub-zfs-object.XXXXXX
    )"

    if mount \
        -t zfs \
        -o ro,zfsutil \
        "$dataset" \
        "$object_mount"
    then
        object_mount_temporary=true
        return 0
    fi

    rmdir "$object_mount" 2>/dev/null || true
    object_mount=""

    return 1
}


release_object_mount()
{
    if [ "$object_mount_temporary" = true ] \
        && [ -n "$object_mount" ]
    then
        umount "$object_mount" >/dev/null 2>&1 || true
        rmdir "$object_mount" >/dev/null 2>&1 || true
    fi

    object_mount=""
    object_mount_temporary=false
}


#
# Build the new fragment separately.
#
# The active grub-zfs.cfg is replaced only after grub-script-check
# validates the complete generated file.
#

rm -f "$grub_new"


{
    printf '%s\n' \
        '# Generated by grub-zfs. Do not edit.'

    printf '%s\n' \
        "menuentry '| Type       | Date                | Details |' { echo }"

} > "$grub_new"


count_objects=0
count_entries=0
count_clones=0
count_synex_snapshots=0
count_external_snapshots=0
count_snapshot_objects=0


#
# Collect direct snapshots of the canonical root dataset.
# Child dataset snapshots are intentionally ignored.
#

mapfile -t snapshots < <(
    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 '%s\n' "$snapshot"
                ;;
        esac
    done
)


synex_snapshots=()
external_snapshots=()


for snapshot in "${snapshots[@]}"
do
    managed="$(
        zfs get \
            -H \
            -o value \
            ar.synex.snapshots:managed \
            "$snapshot" \
            2>/dev/null \
            || true
    )"

    member="$(
        zfs get \
            -H \
            -o value \
            ar.synex.snapshots:member \
            "$snapshot" \
            2>/dev/null \
            || true
    )"

    if [ "$managed" = "true" ] \
        && [ "$member" = "root" ]
    then
        synex_snapshots+=("$snapshot")
    else
        external_snapshots+=("$snapshot")
    fi
done


#
# 1. RW clones created by zfs-initramfs Snapshot Boot.
#
# They are filesystem datasets whose origin is a direct snapshot of the
# canonical root and whose properties match the temporary clone created
# by OpenZFS:
#
#   mountpoint=none
#   canmount=noauto
#   org.zol:mountpoint=/
#
# They are deliberately NOT renamed, promoted or modified here.
#

root_parent="${root_dataset%/*}"


mapfile -t filesystem_candidates < <(
    zfs list \
        -H \
        -t filesystem \
        -o name \
        -S creation \
        -r "$root_parent" \
        2>/dev/null \
        || true
)


for clone_dataset in "${filesystem_candidates[@]}"
do
    [ "$clone_dataset" = "$root_dataset" ] && continue

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

    case "$clone_origin" in
        "$root_dataset"@*)
            ;;

        *)
            continue
            ;;
    esac

    clone_mountpoint="$(
        zfs get -H -o value mountpoint "$clone_dataset" 2>/dev/null || true
    )"

    clone_canmount="$(
        zfs get -H -o value canmount "$clone_dataset" 2>/dev/null || true
    )"

    clone_zol_mountpoint="$(
        zfs get \
            -H \
            -o value \
            org.zol:mountpoint \
            "$clone_dataset" \
            2>/dev/null \
            || true
    )"

    [ "$clone_mountpoint" = "none" ] || continue
    [ "$clone_canmount" = "noauto" ] || continue
    [ "$clone_zol_mountpoint" = "/" ] || continue

    if ! prepare_object_mount "$clone_dataset"; then
        print_warning \
            "Could not mount RW clone for inspection: $clone_dataset"

        continue
    fi

    clone_relative="${clone_dataset#*/}"
    clone_name="${clone_dataset##*/}"
    clone_origin_name="${clone_origin#*@}"
    clone_created_at="$(
        format_zfs_creation "$clone_dataset"
    )"

    clone_title="| [RW Clone] | ${clone_created_at} | ${clone_name} <- ${clone_origin_name} |"

    if emit_boot_object \
        "RW clone" \
        "$clone_title" \
        "$object_mount" \
        "/${clone_relative}@" \
        "$clone_dataset" \
        "$clone_dataset" \
        "clone-${clone_name}" \
        "$clone_dataset"
    then
        count_objects=$((count_objects + 1))
        count_clones=$((count_clones + 1))

        if [ "${GRUB_ZFS_SHOW_SNAPSHOTS_FOUND:-true}" = "true" ]; then
            print_info \
                "Found bootable RW clone: $clone_dataset (origin: $clone_origin)"
        fi
    fi

    release_object_mount
done


#
# 2. Synex-managed snapshots.
#

for snapshot in "${synex_snapshots[@]}"
do
    if [ "$snapshot_limit" -gt 0 ] \
        && [ "$count_snapshot_objects" -ge "$snapshot_limit" ]
    then
        break
    fi

    snap_name="${snapshot#*@}"
    snap_dir="${snapshot_control_dir}/${snap_name}"

    if [ ! -d "$snap_dir" ]; then
        print_warning \
            "Synex snapshot is not accessible through .zfs: $snapshot"

        continue
    fi

    purpose="$(
        zfs get \
            -H \
            -o value \
            ar.synex.snapshots:purpose \
            "$snapshot" \
            2>/dev/null \
            || true
    )"

    created_at="$(
        zfs get \
            -H \
            -o value \
            ar.synex.snapshots:created-at \
            "$snapshot" \
            2>/dev/null \
            || true
    )"

    [ -n "$purpose" ] \
        && [ "$purpose" != "-" ] \
        || purpose="unknown"

    [ -n "$created_at" ] \
        && [ "$created_at" != "-" ] \
        || created_at="$(format_zfs_creation "$snapshot")"

    display_date="$(
        format_created_at "$created_at"
    )"

    display_purpose="$(
        format_purpose "$purpose"
    )"

    snapshot_title="| [Synex] | ${display_date} | ${display_purpose} | ${snap_name} |"

    if emit_boot_object \
        "Synex snapshot" \
        "$snapshot_title" \
        "$snap_dir" \
        "/${dataset_relative}@${snap_name}" \
        "$snapshot" \
        "$root_dataset" \
        "snapshot-${snap_name}" \
        "$snapshot"
    then
        count_objects=$((count_objects + 1))
        count_synex_snapshots=$((count_synex_snapshots + 1))
        count_snapshot_objects=$((count_snapshot_objects + 1))

        if [ "${GRUB_ZFS_SHOW_SNAPSHOTS_FOUND:-true}" = "true" ]; then
            print_info \
                "Found bootable Synex snapshot: $snapshot"
        fi
    fi
done


#
# 3. External snapshots of the canonical root dataset.
#
# They are not adopted by Synex and receive no Synex metadata. They are
# merely exposed as bootable when they contain a complete historical
# kernel/initrd/modules set.
#

for snapshot in "${external_snapshots[@]}"
do
    if [ "$snapshot_limit" -gt 0 ] \
        && [ "$count_snapshot_objects" -ge "$snapshot_limit" ]
    then
        break
    fi

    snap_name="${snapshot#*@}"
    snap_dir="${snapshot_control_dir}/${snap_name}"

    if [ ! -d "$snap_dir" ]; then
        print_warning \
            "External snapshot is not accessible through .zfs: $snapshot"

        continue
    fi

    display_date="$(
        format_zfs_creation "$snapshot"
    )"

    snapshot_title="| [External] | ${display_date} | ${snap_name} |"

    if emit_boot_object \
        "external snapshot" \
        "$snapshot_title" \
        "$snap_dir" \
        "/${dataset_relative}@${snap_name}" \
        "$snapshot" \
        "$root_dataset" \
        "external-${snap_name}" \
        "$snapshot"
    then
        count_objects=$((count_objects + 1))
        count_external_snapshots=$((count_external_snapshots + 1))
        count_snapshot_objects=$((count_snapshot_objects + 1))

        if [ "${GRUB_ZFS_SHOW_SNAPSHOTS_FOUND:-true}" = "true" ]; then
            print_info \
                "Found bootable external snapshot: $snapshot"
        fi
    fi
done


#
# No valid bootable objects:
#
# Remove any stale fragment. The main grub.cfg contains only a
# conditional submenu hook, so the menu disappears automatically.
#

if [ "$count_objects" -eq 0 ] \
    || [ "$count_entries" -eq 0 ]
then
    rm -f \
        "$grub_new" \
        "$grub_cfg"

    print_info \
        "No bootable ZFS snapshots or RW clones found."

    emit_submenu_hook
    exit 0
fi


#
# Validate before replacing the active file.
#
# Because grub-zfs.new and grub-zfs.cfg live in the same directory,
# mv provides an atomic replacement after successful validation.
#

if grub-script-check "$grub_new"; then
    chmod 0644 "$grub_new"

    mv -f \
        "$grub_new" \
        "$grub_cfg"

    print_info \
        "Generated ${grub_cfg}: ${count_clones} RW clone(s), ${count_synex_snapshots} Synex snapshot(s), ${count_external_snapshots} external snapshot(s), ${count_entries} boot entry/entries."
else
    print_error \
        "Generated GRUB fragment failed syntax validation; keeping the previous grub-zfs.cfg."

    exit 1
fi

emit_submenu_hook
