#!/usr/bin/env bash
#
# Copyright 2022 RnD Center "ELVEES", JSC
#
# SPDX-License-Identifier: GPLv3
#

set -euo pipefail

error() {
    echo "Error: $*" > /dev/stderr
    exit 1
}

hideifok() {
    local out
    out=$(eval "$*" 2>&1) || (echo "$out" && exit 1)
}

# NOTE: There's false positive here: dm-xxx, but who cares...
iswholedisk() {
    test -e "/sys/block/${1##*/}"
}

getpart() {
    local partprefix=
    [[ "$1" =~ [0-9]$ ]] && partprefix=p
    echo "$1$partprefix$2"
}

cmdexists() {
    hash "$1" 2>/dev/null
}

unmount() {
    if cmdexists udisksctl; then
        udisksctl unmount -b "$1"
    elif cmdexists systemd-umount; then
        systemd-umount "$1"
    else
        umount "$1"
    fi
}

unmountall() {
    # shellcheck disable=SC2013
    for i in $(grep -F "$1" /etc/mtab | cut -f1 -d' '); do
        unmount "$i"
    done
}

help() {
    echo 'tar2dev - Tool to unpack a tarball onto block device'
    echo '          Partition, format and unpack tarball in one command'
    echo
    echo 'Usage: tar2dev [options] <tarball> <device>'
    echo
    echo 'Options:'
    echo '  -h   Print this help'
    echo '  -n   Set hostname'
    echo '  -s   Comma-separated systemd services to enable'
    echo
    echo 'After extraction /usr/bin/extlinux-gen is called in target root filesystem.'
    echo 'It gets its parameters using environment variables. For more information see'
    echo 'extlinux-gen source code.'
    echo
    echo 'Examples:'
    echo '  ROOT=/dev/mmcblk0p1 FDT=board.dtb tar2dev -n myhost rootfs.tar /dev/sdX'
    echo '  ssh root@myboard zcat \| [vars] tar2dev [options] - /dev/mmcblkX < rootfs.tar.gz'
    echo '  ROOT=/dev/mmcblk0p1 FDT=board.dtb tar2dev -s stress-test rootfs.tar /dev/sdX'
    exit 1
}

unset HOSTNAME

while getopts 'hn:s:' opt; do
    case $opt in
        h) help;;
        n) HOSTNAME="$OPTARG";;
        s) SYSTEMD_SERVICES="$OPTARG";;
        *) help;;
    esac
done
shift $((OPTIND-1))

set +u
[[ -z "$1" ]] && help
[[ -z "$2" ]] && error "Not enough arguments";
set -u

TAR="$1"
DEV="$2"

[[ "$EUID" != "0" ]] && error "This script must be run as root"
[[ "$TAR" == "-" || -r "$TAR" ]] || error "'$TAR' is not readable"
iswholedisk "$DEV" || error "'$DEV' is not a disk"
[[ -w "$DEV" ]] || error "'$DEV' is not writable"

# Real work starts here!

echo "Unmount all partitions on '$DEV'..."
unmountall "$DEV"

echo "Partition '$DEV'..."
hideifok "echo ';' | sfdisk -f '$DEV'"

PART=$(getpart "$DEV" 1)

echo "Format '$PART'..."
hideifok "mkfs.ext4 -F -L root '$PART'"

echo "Unpack '$TAR' on '$PART'..."
DIR=$(mktemp -d)
# shellcheck disable=SC2064
trap "echo 'Clean...'; rmdir '$DIR'" EXIT

mount "$PART" "$DIR"
# shellcheck disable=SC2064
trap "echo 'Unmount...'; unmount $PART; echo 'Clean...'; rmdir '$DIR'" EXIT

tar -C "$DIR" -xf "$TAR"

if [[ -x "$DIR"/usr/bin/extlinux-gen ]]; then
    echo "Generate extlinux.conf..."
    TARGET_DIR="$DIR" "$DIR"/usr/bin/extlinux-gen
fi

if [[ -v HOSTNAME ]]; then
    echo "Set hostname to '$HOSTNAME'..."
    echo "$HOSTNAME" > "$DIR"/etc/hostname
    sed -i -e "s|buildroot|$HOSTNAME|" "$DIR"/etc/hosts
fi

if [[ -v SYSTEMD_SERVICES ]]; then
    for SERVICE in ${SYSTEMD_SERVICES//,/ }; do
        echo "Enable '$SERVICE' systemd service ..."
        ln -s /usr/lib/systemd/system/"$SERVICE".service \
            "$DIR"/etc/systemd/system/multi-user.target.wants/"$SERVICE".service
    done
fi