#!/usr/bin/env bash
#
# Copyright 2022-2023 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
}

# Up to here the code is sourceable
if [[ "${BASH_SOURCE[0]}" != "${0}" ]]; then
    return 0
fi

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 '  -v   Verbosely list files processed'
    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 'If ROOT variable is not set or empty then will be used automatic PARTUUID detect.'
    echo
    echo 'Examples:'
    echo '  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
TAR_VERBOSE=""

while getopts 'hvn:s:' opt; do
    case $opt in
        h) help;;
        v) TAR_VERBOSE="-v";;
        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'"

# Use automatic PARTUUID detection if ROOT is not specified or empty.
# Universal way is to get part UUID via "lsblk -po NAME,PARTUUID", but in lsblk
# PARTUUID will available in ~1 second after previous sfdisk call. This way is
# difficult to do because delay may depend on PC speed. More reliable way is
# to use sfdisk to read part UUID from device.
# For GPT part UUID can be read via sfdisk --part-uuid.
# For MBR part UUID can be got by concatenation of MBR label-id and partition number.
if [[ ! -v ROOT || -z "$ROOT" ]]; then
    LABEL=$(sfdisk -d "$DEV" | grep "label:" | awk '{ print $2 }')
    if [[ "$LABEL" == "dos" ]]; then
        ID=$(sfdisk -d "$DEV" | grep "label-id:" | awk '{ print $2 }')
        ROOT="PARTUUID=${ID:2}-01"
    else
        GPT_PARTUUID=$(sfdisk --part-uuid "$DEV" 1)
        # sfdisk return uppercase PARTUUID for GPT partitions, this is why some
        # distributions won't boot, instead lowercase are always booted.
        ROOT="PARTUUID=${GPT_PARTUUID,,*}"
    fi
    export ROOT
fi

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" $TAR_VERBOSE -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