#!/bin/sh

# set some defaults
KERNEL_CMDLINE_BASE="modules=loop,squashfs,sd-mod,usb-storage quiet console=tty1"
KERNEL_CMDLINE_EXTRA=""

# shellcheck disable=SC1091
. /usr/lib/libalpine.sh

techo() {
	echo
	echo "# $*"
	echo
}

usage() {
	cat <<-__EOF__
$PROGRAM 0.1.26-r0
usage: $PROGRAM <SRC> <DEV>

Wipes all data from DEV, creates GPT partitioning for RKOS UEFI based
dual boot partition schema and copies the mounted RKOS image from SRC
to DEV.

SRC must be a mount point.
DEV must be a block device.

Options:
 -s OPT Serial kernel console setting:
          off - do not enable serial console
          on  - configure detected serial console
          ttySx,115200 - tty+speed setting
 -h  Show this help.
 -v  Verbose mode. Display whats going on.

__EOF__
	exit 1
}

# detect if run from serial console
if [ -n "$(setserial -b -g "$(tty)" 2> /dev/null)" ]; then
	serial_console=on
else
	serial_console=off
fi

while getopts "hvSs:" opt; do
	case "$opt" in
	h) usage;;
	v) VERBOSE=1; export VERBOSE;;
	\?) usage;;
	s) serial_console="$OPTARG";;
	*) usage;;
	esac
done

shift $((OPTIND - 1))

img="$1"
dev="$2"

# Sanity checks
[ -z "$dev" ] || [ -z "$img" ] && usage

apk add --quiet --virtual .rkos-install-deps e2fsprogs efibootguard jq util-linux util-linux-misc || die "Failed to install dependencies!"

[ ! -f "$img/.alpine-release" ] && die "Source is not a valid image mount!"
[ ! -b "$dev" ] && die "Destination $dev is not a block device!"
disk_size=$(blockdev --getsize64 "$dev")

# check min disk size
if [ -z "$disk_size" ] || [ "$disk_size" -lt 4294967296 ]; then
	cat <<- ERRSIZE

	The disk size of $dev is to small. At least 4GiB storage are required!

	ERRSIZE
	exit 1
fi

# efibootguard uses some custom CPU architecture slugs
case "$(uname -m)" in
	x86_64)
		bg_arch="x64"
		;;
	aarch64)
		bg_arch="a64"
		;;
	*)
		die "Unsupported CPU architecture '$(uname -m)'!"
esac

techo "Umounting $dev"
# find and unmount all mounted partitions of $dev
for part in "$dev"*; do
	if test -b "$part"; then
		# find all mount points of partition device
		while read mountpoint; do
			# find loop mounts from $mountpoint (i.e. .modloop)
			while read lomountpoint; do
				echo "  umounting ${lomountpoint/loop\//loop}"
				umount "$lomountpoint" 2> /dev/null || umount "${lomountpoint/loop\//loop}" || exit 1
			done < <(losetup -J | jq -r '.["loopdevices"][]|select(.["back-file"] | startswith("'"$mountpoint"'/"))["back-file"]')
			echo "  umounting $mountpoint"
			umount "$part" || exit 1
		done < <(findmnt -rn -o target "$part")
	fi
done

# shellcheck disable=SC2046
set -- $(ls "$img/boot/vmlinuz-"*)
kernel_image="$1"
kernel_flavor="${kernel_image##*-}"

# prepare serial console
case "$serial_console" in
	off)
		serial_console=""
		;;
	on)
		_tty="$(tty)"
		serial_console="$(basename "$_tty"),$(stty speed)"
		echo "Serial console: $serial_console"
		serial_console=" console=$serial_console"
		;;
	*)
		serial_console=" console=$serial_console"
		;;
esac
KERNEL_CMDLINE_EXTRA="$KERNEL_CMDLINE_EXTRA$serial_console"

techo "Partitioning $dev"

# remove and discard previous data
wipefs --all "$dev"
blkdiscard "$dev"

# build unique install id
install_id="$(tr -dc a-z0-9 </dev/urandom | head -c 16)"

# prepare partition sizes
part_sys_size=1792MiB
part_cfg_opts=""
part_cache=""
if [ "$disk_size" -ge 21474836480 ]; then
	part_sys_size=6GiB
	part_cfg_opts=size=3GiB,
	part_cache="5: size=4GiB,type=0FC63DAF-8483-4772-8E79-3D69D8477DE4,name=rkos-cache@$install_id"
elif [ "$disk_size" -ge 17179869184 ]; then
	part_sys_size=6GiB
	part_cfg_opts=size=3GiB,
elif [ "$disk_size" -ge 12884901888 ]; then
	part_sys_size=4GiB
	part_cfg_opts=size=3GiB,
elif [ "$disk_size" -ge 8589934592 ]; then
	part_sys_size=3840MiB
fi

# set GPT label and create partitions
sfdisk "$dev" <<__EOP__
label: gpt
device: $dev
unit: sectors

1: size=128MiB,type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B,name=rkos-efi@$install_id
2: size=$part_sys_size,type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7,name=rkos-sys-1@$install_id
3: size=$part_sys_size,type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7,name=rkos-sys-2@$install_id
4: ${part_cfg_opts}type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7,name=rkos-config@$install_id
$part_cache
__EOP__

techo "Restart mdev service"
rc-service -D -q mdev restart

techo "Create file systems"
init_tmpdir DESTDIR
labels="rkos-efi rkos-sys-1 rkos-sys-2 rkos-config"
[ -n "$part_cache" ] && labels="$labels rkos-cache"
for _label in $labels; do
	_partlabel="$_label@$install_id"
	_part="$(findfs "PARTLABEL=$_partlabel")"
	[ -z "$_part" ] && die "Could not find $_label partition block device!"
	echo "$_part => $_label"
	[ "$_label" != "rkos-cache" ] && mkfs.vfat -n "$_label" "$_part" || mkfs.ext4 -L "$_label" "$_part"
	mount "$_part" "$DESTDIR"
	case "$_label" in
		rkos-efi)
			mkdir -p "$DESTDIR/EFI/boot"
			cp "/usr/lib/efibootguard/efibootguard$bg_arch.efi" "$DESTDIR/EFI/boot/boot$bg_arch.efi"
			;;
		rkos-cache)
			mkdir "$DESTDIR/chunks"
			;;
		rkos-config)
			# dump setup configuration
			cat <<- EOD  | jq -n -R '[inputs | select(length>0) | [splits("=")] | {(.[0]): .[1:]|join("=")}] | add' > "$DESTDIR/.rkos_vars.json"
			install_id=$install_id
			install_date=$(date -Iseconds)
			cmdline_base=$KERNEL_CMDLINE_BASE
			cmdline_extra=$KERNEL_CMDLINE_EXTRA

			EOD
			;;
		rkos-sys-*)
			echo "Copying files"
			cp -p -r "$img/"* "$DESTDIR"

			echo "Make partition bootable"

			# build unified initrd if needed
			initrd="$DESTDIR/boot/initramfs-$kernel_flavor"
			case "$(uname -m)" in
				x86_64)
					cat "$DESTDIR/boot/amd-ucode.img" "$DESTDIR/boot/intel-ucode.img" "$initrd" > "$DESTDIR/uinitrd"
					initrd="$DESTDIR/uinitrd"
					;;
				*)
					;;
			esac

			# get device for rkos-config
			configdev="$(findfs PARTLABEL=rkos-config@$install_id)"
			[ -z "$configdev" ] && die "Could not find config partition block device!"

			# build final kernel cmdline
			cmdline="$KERNEL_CMDLINE_BASE $KERNEL_CMDLINE_EXTRA apkovl=$configdev:/rkos-sys-${_label##*-}.apkovl.tar.gz alpine_repo=/media/$(basename "$_part")/apks modloop=PARTLABEL=$_partlabel rkos_id=$install_id rkos_boot=${_label##*-}"

			# keep cmdline
			cat <<- EOD  | jq -S -n -R '[inputs | select(length>0) | [splits("=")] | {(.[0]): .[1:]|join("=")}] | add' > "$DESTDIR/rkos/vars.json"
			install_id=$install_id
			kernel_cmdline=$cmdline
			kernel_initrd=${initrd#$DESTDIR/}

			EOD

			# build unified kernel
			bg_gen_unified_kernel \
				"/usr/lib/efibootguard/kernel-stub$bg_arch.efi" \
				"$DESTDIR/boot/vmlinuz-$kernel_flavor" \
				"$DESTDIR/unified" \
				--cmdline "$cmdline" \
				--initrd "$initrd"

			# remove no more required unified initrd
			rm -f "$DESTDIR/uinitrd"

			# set efibootguard env
			# shellcheck disable=SC3037
			echo -n "$_label" | iconv -f ascii -t UTF-16LE > "$DESTDIR/EFILABEL"
			bg_setenv --filepath="$DESTDIR/BGENV.DAT" \
				--revision="$(( 3 - ${_label##*-}))" \
				--kernel="C:$_label:unified"
			echo
			;;
	esac
	umount "$_part"
done

apk del --quiet .rkos-install-deps
