#!/bin/sh

# shellcheck shell=ash

# set some defaults
files_to_move="boot efi apks rkos syslinux.cfg .alpine-release unified"
read_only_mounts=
umounts=
arch="${APK_ARCH:-$(apk --print-arch 2>/dev/null || uname -m)}"
kernel_flavor="$(jq -r '.kernel_flavor // "lts"' < /media/rkos-config/.rkos_vars.json || echo lts)"
image_format="$(jq -r '.image_format // "iso"' < /media/rkos-config/.rkos_vars.json || echo iso)"
rkos_release=next
rkos_url=https://mirrors.routerkit.net/rkos

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

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

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

cleanup_tmpdata() {
	if [ -d "$dst" ] && [ -d "$dst/.new" ]; then
		rm -rf "$dst/.new"
	fi
}

cleanup_mounts() {
	local i=
	cd_assert /
	sync
	sleep 1
	for i in $read_only_mounts; do
		mount -o remount,ro "$i" || echo "Warning: Failed to remount as read-only. Is modloop mounted?"
	done
	read_only_mounts=""
	if [ -n "$umounts" ]; then
		# shellcheck disable=SC2086
		umount $umounts
		umounts=""
	fi
}

cleanup_workdir() {
	rm -rf "$WORKDIR"
}

cleanup() {
	rc=$?
	cleanup_tmpdata
	cleanup_mounts
	cleanup_workdir
	exit $rc
}

WORKDIR="$(mktemp -d "/tmp/$PROGRAM-XXXXXXXXXX")" || die "Failed to create tmp dir"
trap cleanup EXIT INT TERM QUIT

# check if given dir is read-only
is_read_only() {
	local tmpfile
	tmpfile="$(mktemp -p "$1" 2>/dev/null)"
	[ -z "$tmpfile" ] && return 0
	rm -f "$tmpfile"
	return 1
}

update_check() {
	local new_manifest="$1" old_manifest="$2"
	local new old key

	if [ -z "$force" ] && [ ! -f "$new_manifest" ]; then
		die "Update image seems to be broken, manifest not available. Aborting."
	fi

	if [ -z "$force" ] && [ ! -f "$old_manifest" ]; then
		die "Destination seems to be broken, manifest not available. Aborting."
	fi

	techo "Available Upgrades"
	new_kernel=
	for key in alpine_release kernel_version image_arch rkos_release; do
		old="$(jq -r ".$key" "$old_manifest")"
		new="$(jq -r ".$key" "$new_manifest")"
		if [ "$old" != "$new" ]; then
			printf '%10s: \033[1m%s => %s\033[0m\n' "${key%_*}" "$old" "$new"
			[ "$key" = "image_arch" ] && die "The CPU architecture does not match, aborting! "
			[ "$key" = "kernel_version" ] && new_kernel=1
		else
			[ "$key" = "image_arch" ] || printf '%10s: %s == %s\n' "${key%_*}" "$old" "$new"
		fi
	done

	new="$(jq -r ".kernel_flavor" "$new_manifest")"
	if [ "$kernel_flavor" != "$new" ]; then
		eecho -e "\033[1mWARNING:\033[0m Changing kernel flavor from $kernel_flavor to $new!"
		kernel_flavor="$new"
	fi

	printf '  packages: \033[5mchecking\033[0m'
	pkg_updates=
	if [ -n "$rkos_current" ]; then
		# check for installed packages, only
		apk query --format json --installed '*' | jq 'reduce .[] as $pkg ({}; . + {($pkg.name): $pkg.version})' > "$WORKDIR/packages.json"
	else
		# check for available packages
		jq .packages "$old_manifest" > "$WORKDIR/packages.json"
	fi
	for pkg in $(jq -r 'keys|.[]' "$WORKDIR/packages.json"); do
		old="$(jq -r ".[\"$pkg\"]" "$WORKDIR/packages.json")"
		new="$(jq -r ".packages[\"$pkg\"] // \"\"" "$new_manifest")"
		if [ -z "$new" ]; then
			eecho -e "\033[1mWARNING:\033[0m Package $pkg is no more available!"
		fi
		if [ "$old" != "$new" ]; then
			[ -n "$VERBOSE" ] && [ -z "$pkg_updates" ] && techo "Package Upgrades"
			[ -n "$VERBOSE" ] && printf "  %s: %s => %s\n" "$pkg" "$old" "$new"
			pkg_updates=$((pkg_updates+1))
		fi
	done
	if [ -z "$VERBOSE" ]; then
		if [ -z "$pkg_updates" ]; then
			printf '\r  packages: no upgrades\n'
		else
			printf '\r  packages: \033[1m%d upgrades\033[0m\n' "$pkg_updates"
		fi
	fi
	echo

	if [ -z "$force" ] && [ -z "$pkg_updates$new_kernel" ]; then
		echo "Nothing to upgrade, exiting."
		exit 0
	fi

	if ! ask_yesno "Continue upgrading $dst? (y/n)" y; then
		exit 0
	fi
	echo
}

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

Upgrades the RKOS image files including the unified kernel image and
the apk repository archive.

When SRC is omitted, the official RKOS
image download will be used. If DST is omitted, the partition from
which the current system has been booted will be used.

Options:
 -u OPT  base URL used for auto download (default: $rkos_url)
 -i OPT  override image format for auto download (default: $image_format)
 -k OPT  override kernel flavor for auto download (default: $kernel_flavor)
 -r OPT  override release for auto download (default: $rkos_release)
 -h  Show this help.
 -v  Verbose mode. Display whats going on.

__EOF__
	exit 1
}


while getopts "hvu:i:k:r:u:" opt; do
	case "$opt" in
	h) usage;;
	u) rkos_url="$OPTARG";;
	v) VERBOSE=1; export VERBOSE;;
	\?) usage;;
	k) kernel_flavor="$OPTARG";;
	v) rkos_release="$OPTARG";;
	*) usage;;
	esac
done

shift $((OPTIND - 1))

# Get directory of currently running RKOS
rkos_current=
# shellcheck disable=SC2013
for opt in $(cat /proc/cmdline); do
	case "$opt" in
		rkos_boot=*)
		rkos_current="/media/rkos-sys-${opt#rkos_boot=}"
		;;
	esac
done

case "$#" in
	0)
		src=""
		[ -n "$rkos_current" ] && dst="$rkos_current" || dst="/media/rkos-sys-1"
		;;
	1)
		src=""
		dst="$1"
		;;
	2)
		src="$1"
		dst="$2"
		;;
	*) usage;;
esac

# Check if running system will be upgraded
[ "$rkos_current" != "$dst" ] && rkos_current=

# Sanity checks
if [ ! -r "$dst/EFILABEL" ]; then
	error "Directory $dst seems not to be a valid RKOS installation mountpoint!" 1>&2
	exit 1
fi

# get kernel cmdline, required for building the unified kernel
cmdline="$(jq -r '.kernel_cmdline // ""' < "$dst/rkos/vars.json")"
[ -z "$cmdline" ] && die "Unable to get kernel cmdline!"

# use download URL if no source was specified
if [ -z "$src" ]; then
	src="$rkos_url/$rkos_release/releases/$arch/rkos-$kernel_flavor-$rkos_release-$arch.$image_format"
fi

vecho "Checking available upgrade from $src"

# copy data from source to .new (if it is not a mounted image)
if [ -f "$src"/rkos/manifest.json ]; then
	src="${src%/}"
	update_check "$src/rkos/manifest.json" "$dst/rkos/manifest.json"

	# remount as rw if needed
	if is_read_only "$dst"; then
		vecho "Remounting $dst as read/write"
		mount -o remount,rw "$dst" || die "Failed to remount $dst as rw"
		read_only_mounts="$read_only_mounts $dst"
	fi

	# remove partial upgrades if any
	rm -rf "$dst/.new" "$dst/.old"
	mkdir -p "$dst/.new" || die "Failed to create $dst/.new"

	(cd_assert "$dst/.new"; exec ${UNISO:-uniso}) < "$src" \
		|| die "Failed to download or extract $src"
	src="$dst/.new"
else
	vecho "Extracting $src to $dst/.new/"

	case "$src" in
	https://*|http://*|ftp://*)
		${WGET:-wget} -q -O "$WORKDIR/manifest.json" "$src.manifest.json" || die "Could not download manifest from $src.manifest.json!"
		update_check "$WORKDIR/manifest.json" "$dst/rkos/manifest.json"

		# remount as rw if needed
		if is_read_only "$dst"; then
			vecho "Remounting $dst as read/write"
			mount -o remount,rw "$dst" || die "Failed to remount $dst as rw"
			read_only_mounts="$read_only_mounts $dst"
		fi

		# remove partial upgrades if any
		rm -rf "$dst/.new" "$dst/.old"
		mkdir -p "$dst/.new" || die "Failed to create $dst/.new"

		# get and extract image
		${WGET:-wget} -O - "$src" | (cd_assert "$dst/.new"; exec ${UNISO:-uniso}) \
			|| die "Failed to download or extract $src"
		echo ""
		src="$dst/.new"
		;;
	*)
		update_check "$src.manifest.json" "$dst/rkos/manifest.json"

		# remount as rw if needed
		if is_read_only "$dst"; then
			vecho "Remounting $dst as read/write"
			mount -o remount,rw "$dst" || die "Failed to remount $dst as rw"
			read_only_mounts="$read_only_mounts $dst"
		fi

		# remove partial upgrades if any
		rm -rf "$dst/.new" "$dst/.old"
		mkdir -p "$dst/.new" || die "Failed to create $dst/.new"

		(cd_assert "$dst/.new"; exec ${UNISO:-uniso}) < "$src" \
			|| die "Failed to download or extract $src"
		src="$dst/.new"
		;;
	esac
fi

# Build unified initrd
initrd="$dst/.new/boot/initramfs-$kernel_flavor"
case "$arch" in
	x86_64)
		vecho "Build unified initrd..."
		cat "$src/boot/amd-ucode.img" "$src/boot/intel-ucode.img" "$src/boot/initramfs-$kernel_flavor" > "$WORKDIR/uinitrd"
		initrd="$WORKDIR/uinitrd"
		;;
	*)
		;;
esac

vecho "Build unified kernel image..."
bg_gen_unified_kernel \
	"/usr/lib/efibootguard/kernel-stub$bg_arch.efi" \
	"$src/boot/vmlinuz-$kernel_flavor" \
	"$dst/.new/unified" \
	--cmdline "$cmdline" \
	--initrd "$initrd"

vecho "Preserve RKOS vars.json file..."
cp -p "$dst/rkos/vars.json" "$dst/.new/rkos/vars.json" || die "Could not preserve 'rkos/vars.json'"

# make sure files are really there before we replace existing
vecho "Flushing cache..."
sync

vecho "Replacing existing files..."
mkdir -p "$dst/.old" || die "Failed to create $dst/.old"

# move current files to .old
for i in $files_to_move; do
	if [ -e "$dst"/"$i" ]; then
		mv "$dst"/"$i" "$dst/.old"/ || die "Failed to move $dst/$i to $dst/.old/"
	fi
done

# move .new to current
for i in $files_to_move; do
	if [ -e "$dst/.new"/"$i" ]; then
		mv "$dst/.new"/"$i" "$dst"/ \
			|| die "Failed to move $dst/.new/ to $dst"
	fi
done

# make sure files are finally there
vecho "Flushing cache..."
sync

# cleanup
rm -rf "$dst/.old" "$dst/.new"

echo
echo -n "$dst has been upgraded."
if [ -n "$rkos_current" ]; then
	[ -n "$pkg_updates" ] && echo -n " Consider to upgrade $pkg_updates packages."
	[ -n "$new_kernel" ] && echo -n " Consider to reboot for the kernel upgrade."
fi
echo
