UPGRADE_FW_MAJOR="2023-09-07-0-022c6ed2-23.09-plus"
UPGRADE_FW_VERSION="2023-09-07-0-022c6ed2-23.09-plus"
UPGRADE_FW_REQUIRE=""

FS_STATE_READY=2
OVERLAY_PATH="/tmp/overlay"

check_fw_version() {
	echo "Gathering current firmware information..."

	local board_name=$(board_name)
	local fw_major=$(bos_major)
	local fw_version=$(bos_version)

	echo "board  : $board_name"
	echo "version: $fw_version"
	echo "major  : $fw_major"

	echo "Checking compatibility..."

	if [ "$UPGRADE_FW_VERSION" ">" "$fw_version" ]; then
		# Firmware upgrade
		if [ "$UPGRADE_FW_REQUIRE" ">" "$fw_version" ]; then
			echo "Firmware upgrade to '$UPGRADE_FW_VERSION' is not possible!"
			echo "Firmware version '$UPGRADE_FW_REQUIRE' is required before upgrading to this version."
			return 1
		fi
	elif [ "$UPGRADE_FW_VERSION" "<" "$fw_version" ]; then
		# Firmware downgrade
		if [ "$UPGRADE_FW_MAJOR" != "$fw_major" ]; then
			echo "Firmware downgrade to '$UPGRADE_FW_VERSION' is not possible!"
			echo "Downgrade is only possible among firmwares with major version '$UPGRADE_FW_MAJOR'."
			echo "Do the factory reset and try to upgrade to this version."
			return 1
		fi
	fi

	return 0
}

get_next_rootfs_index() {
	local curr_index

	curr_index=$(fw_printenv -n rootfs_index 2>/dev/null)
	case "$curr_index" in
		"") ;;
		1) echo 2;;
		2) echo 1;;
		*)
			v "Unexpected rootfs index '$curr_index'"
			return 1
		  ;;
	esac

	return 0
}

find_rootfs_dev() {
	local next_rootfs_index=$1
	local rootfs_label rootfs_dev

	export_bootdevice && export_partdevice diskdev 0 || {
		v "Unable to determine upgrade device"
		return 1
	}

	rootfs_label="rootfs${next_rootfs_index}"
	rootfs_dev=$(blkid \
		--match-token PARTLABEL="$rootfs_label" \
		--output device /dev/${diskdev}p*
	)
	[ -n "$rootfs_dev" ] || {
		v "Unable to find partition '$rootfs_label'"
		return 1
	}

	echo "$rootfs_dev"

	return 0
}

package_check_image() {
	local next_rootfs_index

	check_fw_version || return 1

	next_rootfs_index=$(get_next_rootfs_index)
	[ $? -ne 0 ] && return 1
	find_rootfs_dev $next_rootfs_index >/dev/null || return 1

	return 0
}

package_pre_upgrade() {
	export RAMFS_COPY_BIN="$RAMFS_COPY_BIN blkid"
	export RAMFS_COPY_BIN="$RAMFS_COPY_BIN dmesg"
	export RAMFS_COPY_BIN="$RAMFS_COPY_BIN ln"
	export RAMFS_COPY_BIN="$RAMFS_COPY_BIN mkfs.f2fs"

	return 0
}

package_switch_to_ramfs_required() {
	local next_rootfs_index

	next_rootfs_index=$(get_next_rootfs_index)

	if [ -n "$next_rootfs_index" ]; then
		# Flashing another rootfs device does not require switch to RAMFS
		# The device is not mounted and it is save to overwrite it
		echo "no"
	else
		echo "yes"
	fi

	return 0
}

package_do_upgrade() {
	# Redirect STDOUT and STDERR to /dev/kmsg
	exec 1<&- 2<&- 1>/dev/kmsg 2>&1
	echo "upgrade: - start -"

	local block_size sysupgrade_dir
	local next_rootfs_index rootfs_dev rootfs_path rootfs_length
	local erase_offset overlay_offset

	block_size=4096
	sysupgrade_dir=$(sysupgrade_dir)

	next_rootfs_index=$(get_next_rootfs_index)
	[ $? -ne 0 ] && return 1
	rootfs_dev=$(find_rootfs_dev $next_rootfs_index)
	[ $? -ne 0 ] && return 1

	rootfs_path="$(sysupgrade_dir)/rootfs.img"
	rootfs_length=$(
		(get_image "$@" | tar xf - "$rootfs_path" -O | wc -c) 2>/dev/null
	)
	[ $? -ne 0 ] && return 1

	erase_offset=$(( rootfs_length / block_size ))
	erase_count=$(( 8388608 / block_size ))

	overlay_offset=$(( (rootfs_length + 65535) / 65536 * 65536 ))

	v "Erasing device '$rootfs_dev'..."
	# Do not erase the whole device to speed it up
	dd if=/dev/zero of=$rootfs_dev bs=$block_size seek=$erase_offset \
		count=$erase_count conv=fsync
	sync

	v "Flashing device '$rootfs_dev'..."
	get_image "$@" \
	| tar xf - "$rootfs_path" -O \
	| dd of=$rootfs_dev bs=$block_size conv=fsync
	sync

	v "Creating read-write overlay (offset=$overlay_offset)..."
	export OVERLAY_DEV=$(
		losetup --find --show --offset $overlay_offset $rootfs_dev
	)
	mkfs.f2fs -q -l rootfs_data $OVERLAY_DEV

	mkdir -p "$OVERLAY_PATH"
	mount $OVERLAY_DEV "$OVERLAY_PATH"

	# The index is switched at the end of the whole upgrade process
	export ROOTFS_INDEX="$next_rootfs_index"

	return 0
}

package_copy_config() {
	v "Restoring config files..."

	mv "$UPGRADE_BACKUP" "$OVERLAY_PATH/$BACKUP_FILE"

	return 0
}

package_finish_upgrade() {
	# Mark it ready to avoid "mount_root" wiping it again
	ln -s $FS_STATE_READY "$OVERLAY_PATH/.fs_state"

	v "Saving log..."
	sync
	echo "upgrade: - end -"
	dmesg > "$OVERLAY_PATH/.sysupgrade.log"

	v "Freeing up resources..."
	sync
	umount "$OVERLAY_PATH"
	losetup -d $OVERLAY_DEV

	[ -n "$ROOTFS_INDEX" ] && {
		v "Switching rootfs to $ROOTFS_INDEX"
		fw_setenv rootfs_index $ROOTFS_INDEX
	}

	return 0
}
