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

EFIMOUNT=/boot/efi
GRUBENV=${EFIMOUNT}/EFI/BOOT/grubenv

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
    _partuuid=$(get_partuuid "$1")
    _id=id-${_partuuid}
    printf 'Setting default boot entry to %s\n' "${_id}"
    grub-editenv "${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 "${EFIMOUNT}"; then
        mount -o ro "${_esp}" "${EFIMOUNT}" \
            || warn 'Failed to mount EFI System Partition'
    fi
    if [ -f "${GRUBENV}" ]; then
        info 'Saving current GRUB environment ...'
        cp "${GRUBENV}" grubenv \
            || warn 'Failed to save GRUB environment'
    fi
    if mountpoint -q "${EFIMOUNT}"; then
        umount "${EFIMOUNT}" || 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}" "${EFIMOUNT}" || rc=$?
    fi
    if [ $_rc -eq 0 ]; then
        if [ -f grubenv ]; then
            printf 'Restoring GRUB environment ...\n'
            cp grubenv "${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