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

. /lib/gentoo/functions.sh

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

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

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}
    ebegin 'Setting default boot entry'
    veindent
    veinfo "Menu entry ID: ${_id}"
    veoutdent
    grub-editenv /boot/efi/EFI/gentoo/grubenv set "default=${_id}" || _rc=$?
    eend $_rc
    umount newroot
    return $_rc
}

write_firmware() {
    _rc=0
    _esp=$(findfs PARTLABEL='EFI System Partition')
    if [ -z "${_esp}" ]; then
        eerror 'Could not identify EFI System Partition'
        return 1
    fi
    if ! mountpoint -q /boot/efi; then
        mount -o ro "${_esp}" /boot/efi \
            || ewarn 'Failed to mount EFI System Partition'
    fi
    if [ -f /boot/efi/EFI/gentoo/grubenv ]; then
        ebegin 'Saving current GRUB environment'
        cp /boot/efi/EFI/gentoo/grubenv .
        ewend $? 'Failed to save GRUB environment'
    fi
    if mountpoint -q /boot/efi; then
        umount /boot/efi || return
    fi
    ebegin "Writing firmware image to EFI System Partition (${_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
            ebegin 'Restoring GRUB environment'
            cp grubenv /boot/efi/EFI/gentoo/grubenv || _rc=$?
            eend $_rc
        fi
    fi
    return $_rc
}

write_rootfs() {
    _rc=0
    ebegin "Writing rootfs image to $1"
    dd if=rootfs.squashfs of="$1" bs=1M || _rc=$?
    eend $_rc
}

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=$?
    eerror 'Failed to set default boot option'
fi

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