UPGRADE_FW_MAJOR="2024-03-26-0-5aa91fcb-24.03-plus"
UPGRADE_FW_VERSION="2024-09-12-0-9dbf0134-24.04.3-plus"
UPGRADE_FW_REQUIRE="2023-09-07-0-022c6ed2-23.09-plus"

FS_STATE_READY=2
OVERLAY_PATH="/tmp/overlay"
ENV_SCRIPT_PATH="/tmp/uboot.env"
ENV_MESSAGE_PATH="/tmp/uboot.msg"

BLOCK_SIZE=4096

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
}

file_size() {
	local file_path=$1

	du -b "$file_path" | cut -f1
}

result_to_bool() {
	if [ $? -eq 0 ]; then
		echo "true"
	else
		echo "false"
	fi
}

mmc_block_offset() {
	local offset=$1

	echo $(( offset / BLOCK_SIZE ))
}

mmc_block_count() {
	local length=$1

	echo $(( (length + BLOCK_SIZE - 1) / BLOCK_SIZE ))
}

mmc_dump() {
	local dev=$1
	local offset=$2
	local length=$3

	local dump_offset dump_count
	dump_offset=$(mmc_block_offset $offset)
	dump_count=$(mmc_block_count $length)

	dd if=$dev bs=$BLOCK_SIZE skip=$dump_offset count=$dump_count
}

mmc_unlock() {
	local dev=$1

	local force_path="/sys/class/block/$(basename $dev)/force_ro"
	[ -f $force_path ] && echo 0 > "$force_path"
}

mmc_erase() {
	local dev=$1
	local offset=$2
	local length=$3

	local erase_offset erase_count
	erase_offset=$(mmc_block_offset $offset)
	erase_count=$(mmc_block_count $length)

	v "Erasing device $dev[$offset] (size=$length)..."
	# Do not erase the whole device to speed it up
	dd if=/dev/zero of=$dev bs=$BLOCK_SIZE seek=$erase_offset \
		count=$erase_count conv=fsync
	sync
}

mmc_flash() {
	local name=$1
	local dev=$2
	local offset=${3:-0}

	local flash_offset
	flash_offset=$(mmc_block_offset $offset)

	v "Flashing device $dev[$offset] with $name..."
	dd of=$dev bs=$BLOCK_SIZE seek=$flash_offset conv=fsync
	sync
}

mmc_image_cmp() {
	local dump_path=$1
	shift
	local image_path=$1
	shift

	local dump_count
	dump_length=$(file_size "$dump_path")

	(get_image "$@" | \
		tar xf - "$image_path" -O | \
		cat - /dev/zero | \
		head -c $dump_length | \
		cmp -s "$dump_path") 2>/dev/null
	return $?
}

image_length() {
	local image_path=$1
	shift

	(get_image "$@" | tar xf - "$image_path" -O | wc -c) 2>/dev/null
	return $?
}

get_bootldrs_revert() {
	# NOTE: echo is used here to create single string without newlines
	echo $(cat <<-'EOF'
	emmc_blk_size=0x200;
	fsbl_addr_r=0xc4000000;
	fip_addr_r=0xc4100000;
	part number mmc $dev_emmc fip fip_part &&
	part start mmc $dev_emmc $fip_part fip_blk &&
	setexpr image_blk $fip_blk + 0x1800 &&
	mmc read $loadaddr $image_blk 0x800 &&
	unzip $loadaddr $fsbl_addr_r &&
	setexpr fsbl_cnt $filesize / $emmc_blk_size &&
	setexpr image_blk $image_blk + 0x800 &&
	mmc read $loadaddr $image_blk 0x1000 &&
	unzip $loadaddr $fip_addr_r &&
	setexpr fip_cnt $filesize / $emmc_blk_size &&
	mmc dev $dev_emmc 2 &&
	mmc write $fsbl_addr_r 0 $fsbl_cnt &&
	mmc dev $dev_emmc &&
	mmc write $fip_addr_r $fip_blk $fip_cnt &&
	setenv bootcmd "env default downgrade_env && run downgrade_env && boot" &&
	saveenv &&
	reset
	EOF
	)
}

bootloaders_upgrade() {
	local sysupgrade_dir board_iface
	local fsbl_path fsbl_length
	local fip_path fip_length
	local board_name board_revision

	sysupgrade_dir=$(sysupgrade_dir)
	board_iface=$(board_iface)
	board_name=$(board_name)
	bos_revision=$(bos_revision)
	bos_revision=${bos_revision:-"default"}

	fsbl_path="$sysupgrade_dir/tf-a_${board_iface}_${bos_revision}.stm32"
	fip_path="$sysupgrade_dir/fip_${board_iface}_${bos_revision}.bin"

	v "FIP Image: '$fip_path'"
	v "FSBL Image: '$fsbl_path'"

	fsbl_length=$(image_length "$fsbl_path" "$@")
	[ $? -ne 0 ] && return 1

	fip_length=$(image_length "$fip_path" "$@")
	[ $? -ne 0 ] && return 1

	local fip_dev="$(find_part_fip_dev)"
	local fip_dev_prev="$fip_dev"
	local fsbl_dev="/dev/mmcblk0boot0"
	local fsbl_dev_prev="$fip_dev"

	local fsbl_dev_offset=0
	local fsbl_dev_length=0x100000
	local fip_dev_offset=0
	local fip_dev_length=0x300000

	local fsbl_dev_prev_offset=0x300000
	local fsbl_dev_prev_length=0x100000
	local fip_dev_prev_offset=0x400000
	local fip_dev_prev_length=0x200000

	local fsbl_dump="/tmp/fsbl.dump"
	local fip_dump="/tmp/fip.dump"

	mmc_dump $fsbl_dev $fsbl_dev_offset $fsbl_dev_length > "$fsbl_dump"
	mmc_dump $fip_dev $fip_dev_offset $fip_dev_length > "$fip_dump"

	# Backup previous FSBL (TF-A)
	mmc_erase $fsbl_dev_prev $fsbl_dev_prev_offset $fsbl_dev_prev_length
	gzip -c $fsbl_dump | \
		mmc_flash "TF-A backup" $fsbl_dev_prev $fsbl_dev_prev_offset

	# Backup previous FIP (U-Boot)
	mmc_erase $fip_dev_prev $fip_dev_prev_offset $fip_dev_prev_length
	gzip -c $fip_dump | \
		mmc_flash "U-Boot backup" $fip_dev_prev $fip_dev_prev_offset

	! mmc_image_cmp "$fsbl_dump" "$fsbl_path" "$@"
	local o_update_fsbl=$(result_to_bool)
	! mmc_image_cmp "$fip_dump" "$fip_path" "$@"
	local o_update_fip=$(result_to_bool)

	$o_update_fsbl || $o_update_fip
	local o_update=$(result_to_bool)

	$o_update && {
		# Set up a boot command to trigger a recovery of original bootloaders
		# if something goes wrong during this most critical part of the update
		fw_setenv bootcmd "$(get_bootldrs_revert)"
	}
	# Flash new FSBL (TF-A)
	$o_update_fsbl && {
		mmc_unlock $fsbl_dev
		mmc_erase $fsbl_dev $fsbl_dev_offset $fsbl_dev_length
		get_image "$@" \
			| tar xf - "$fsbl_path" -O \
			| mmc_flash "TF-A" $fsbl_dev
	}
	# Flash new FIP (U-Boot)
	$o_update_fip && {
		mmc_unlock $fip_dev
		mmc_erase $fip_dev $fip_dev_offset $fip_dev_length
		get_image "$@" \
			| tar xf - "$fip_path" -O \
			| mmc_flash "U-Boot" $fip_dev
	}
	$o_update && {
		# Update U-Boot env to the latest version and boot the system
		echo "bootcmd=ii_resetenv && ii_saveenv && boot" >> "$ENV_SCRIPT_PATH"
	}

	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() {
	# Do not truncate kernel messages
	echo on > /proc/sys/kernel/printk_devkmsg
	# Redirect STDOUT and STDERR to /dev/kmsg
	exec 1<&- 2<&- 1>/dev/kmsg 2>&1
	echo "upgrade: - start -"

	local sysupgrade_dir
	local next_rootfs_index rootfs_dev rootfs_path rootfs_length
	local overlay_offset

	# Ensure that the script does not exist
	rm -f "$ENV_SCRIPT_PATH" "$ENV_MESSAGE_PATH"

	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=$(image_length "$rootfs_path" "$@")
	[ $? -ne 0 ] && return 1

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

	# Do not erase the whole device to speed it up
	mmc_erase $rootfs_dev $rootfs_length 0x800000
	get_image "$@" \
	| tar xf - "$rootfs_path" -O \
	| mmc_flash "rootfs" $rootfs_dev

	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
	[ -n "$next_rootfs_index" ] && {
		echo "Switching rootfs to $next_rootfs_index" >> "$ENV_MESSAGE_PATH"
		echo "rootfs_index=$next_rootfs_index" >> "$ENV_SCRIPT_PATH"
		# Set auto-recovery for two failed attempts to start the system
		echo "recovery_counter=2" >> "$ENV_SCRIPT_PATH"
	}

	# The most critical part of upgrade is done at the end. When it is somehow
	# interrupted it can lead to factory reset or brick
	[ "$(bos_mode)" = "emmc" ] && bootloaders_upgrade "$@"

	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

	# Print all messages and modify the U-Env script at once
	[ -f "$ENV_MESSAGE_PATH" ] && {
		while IFS= read -r line; do
			v "$line"
		done < "$ENV_MESSAGE_PATH"
	}
	[ -f "$ENV_SCRIPT_PATH" ] && {
		fw_setenv --script "$ENV_SCRIPT_PATH"
	}

	return 0
}
