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

cleanup() {
    cd /
    if [ -n "${workdir}" ] && [ "${workdir}" != / ]; then
        rm -rf "${workdir}"
    fi
    unset workdir
}

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

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

extract_update() {
    zstd -dc update.tar.zstd | tar -x \
        || die 'Could not extract update source'
    sha256sum -c digests \
        || die 'Invalid update source: checksum mismatch'
}

fetch_update() {
    wget -O update.tar.zstd "$1"
}

get_root() {
    set -- $(cat /proc/cmdline)
    while [ $# -gt 0 ]; do
        case "$1" in
            root=*)
                _root=${1#root=}
                ;;
        esac
        shift
    done
    echo $(findfs "${_root}")
}

get_partlabel() {
    blkid -o value -s PARTLABEL "$1"
}

help() {
    usage
}

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

usage() {
    printf 'usage: %s source_url\n' "${0##*/}"
}

while [ $# -gt 0 ]; do
    case "$1" in
    --help)
        help
        exit 0
        ;;
    *)
        if [ -z "${source_url}" ]; then
            source_url="$1"
        else
            usage >&2
            exit 2
        fi
        ;;
    esac
    shift
done

if [ -z "${source_url}" ]; then
    usage >&2
    exit 2
fi

root=$(get_root)
partlabel=$(get_partlabel "${root}")

case "${partlabel}" in
rootfs-a)
    newpartlabel=rootfs-b
    ;;
rootfs-b)
    newpartlabel=rootfs-a
    ;;
*)
    die \
        'Unsupported system configuration: invalid rootfs partition label: %s\n' \
        "${partlabel}" >&2
esac
newroot=$(findfs PARTLABEL="${newpartlabel}")
if [ -z "${newroot}" ]; then
    die 'Could not find partition with label %s\n' "${partlabel}"
fi
info 'Current rootfs: %s (%s)\n' "${partlabel}" "${root}"
info 'New rootfs: %s (%s)\n' "${newpartlabel}" "${newroot}"

trap cleanup INT TERM QUIT EXIT
workdir=$(mktemp -d)
cd "${workdir}"

systemctl daemon-reload

fetch_update "${source_url}" || die 'Failed to fetch update source'
extract_update || die 'Failed to extact update source'
./install "${newroot}" || die 'Error installing system update'

printf 'Do you want to reboot now? [y/N] '
read confirm
case "${confirm}" in
[yY]|[yY][eE][sS])
    systemctl reboot
    ;;
*)
    info 'A reboot is required to complete the update'
    ;;
esac