#!/bin/sh
# vim: set sw=4 ts=4 sts=4 et :

die() {
    rc=$?
    if [ $rc -eq 0 ]; then
        rc=1
    fi
    error "$@"
    exit $rc
}

error() {
    printf 'ERROR: '
    info "$@"
}

get_partuuid() {
    blkid -o value -s PARTUUID "$1"
}

info() {
    if [ $# -eq 1 ]; then
        echo "$1" >&2
    elif [ $# -gt 1 ]; then
        printf "$@" >&2
    fi
}

set_default_boot() {
    _rc=0
    mkdir -p newroot || return
    mount -oro "$1" newroot || return
    _kernel=$(find newroot/boot -name 'vmlinuz-*' -printf '%P\n' \
        | sort -V \
        | tail -n1
    )
    _kver=${_kernel#vmlinuz-}
    _partuuid=$(get_partuuid "$1")
    _id=id-${_partuuid}-${_kver}
    printf 'Setting default boot entry to %s\n' "${_id}"
    grub-editenv /boot/efi/EFI/gentoo/grubenv set "default=${_id}" || rc=$?
    umount newroot
    return $rc
}

warn() {
    printf 'WARNING: '
    info "$@"
}

write_firmware() {
    _rc=0
    _esp=$(findfs PARTLABEL='EFI System Partition')
    if [ -z "${_esp}" ]; then
        error 'Could not identify EFI System Partition'
        return 1
    fi
    if ! mountpoint -q /boot/efi; then
        mount -o ro "${_esp}" /boot/efi \
            || warn 'Failed to mount EFI System Partition'
    fi
    if [ -f /boot/efi/EFI/gentoo/grubenv ]; then
        info 'Saving current GRUB environment ...'
        cp /boot/efi/EFI/gentoo/grubenv . \
            || warn 'Failed to save GRUB environment'
    fi
    if mountpoint -q /boot/efi; then
        umount /boot/efi || return
    fi
    info 'Writing firmware image to EFI System Partition (%s) ...\n' "${_esp}"
    dd if=firmware.img of="${_esp}" bs=1M || _rc=$?
    if [ $_rc -eq 0 ]; then
        mount -orw "${_esp}" /boot/efi || rc=$?
    fi
    if [ $_rc -eq 0 ]; then
        if [ -f grubenv ]; then
            printf 'Restoring GRUB environment ...\n'
            cp grubenv /boot/efi/EFI/gentoo/grubenv || _rc=$?
        fi
    fi
    return $_rc
}

write_rootfs() {
    printf 'Writing rootfs image to %s ...\n' "$1"
    dd if=rootfs.squashfs of="$1" bs=1M
}

rc=0
newroot="$1"

write_rootfs "${newroot}" || die 'Failed to write new rootfs image to disk'
write_firmware || die 'Failed to write new firmware image to disk'
if ! set_default_boot "${newroot}"; then
    rc=$?
    error 'Failed to set default boot option'
fi

if [ $rc -eq 0 ]; then
    info 'Successfully installed update'
fi
exit $rc