UPGRADE_FW_MAJOR="2025-10-20-0-62a2a20f-25.10"
UPGRADE_FW_VERSION="2026-07-21-0-34648a21-26.07-rc"
UPGRADE_FW_REQUIRE="2025-10-20-0-62a2a20f-25.10"

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

BLOCK_SIZE=4096

# Staging area for the Nix payload extracted from the sysupgrade
# tarball. /dev/shm is tmpfs, so it disappears on reboot by itself.
NIX_STAGE_DIR="/dev/shm/nix-payload"

# Pending package installs handed over by bmc right before it
# triggers the upgrade; consumed while staging the profile in this
# same boot, so the tmpfs lifetime is exactly right.
NIX_PENDING_INSTALL="/dev/shm/bmc-nix-pending-install.json"

# One sysupgrade run validates the image twice (/sbin/sysupgrade's
# validate_firmware_image and procd's upgraded re-validation both call
# platform_check_image), and each call would stage the profile. The
# marker records the version whose profile was staged so later passes
# skip instead of building another generation. It is written only
# after staging succeeded, so a match always means a valid profile for
# that version exists in the persistent store and skipping is safe for
# any number of passes or repeated runs in the same boot. The marker
# must not be consumed on match: the re-validation pass runs after
# network teardown and cannot re-stage, so a run whose first pass
# consumed a leftover marker from an interrupted run would fail there.
# /dev/shm is tmpfs, so a reboot clears it.
NIX_PREPARED_MARKER="/dev/shm/bos-nix-profile-prepared"

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
}

stage_nix_payload_file() {
	local name=$1
	shift

	get_image "$@" | tar xf - "$(sysupgrade_dir)/$name" -O \
		> "$NIX_STAGE_DIR/$name" || {
		v "Unable to extract '$name' from the sysupgrade image"
		return 1
	}

	return 0
}

# Required members of the Nix payload. COMMAND itself is extracted
# from the same tarball, so the payload layout always matches this
# script.
NIX_PAYLOAD_MEMBERS="bmc-nix-cli servers.json.default"

nix_cli() {
	"$NIX_STAGE_DIR/bmc-nix-cli" --log-format internal-json "$@"
}

# Side-effect-free payload classification, safe under `sysupgrade -T`.
# Returns 0 when the payload is complete, 1 when the image carries no
# payload at all, 2 when it is partial — a malformed image.
validate_nix_payload() {
	local member length present="" missing=""

	for member in $NIX_PAYLOAD_MEMBERS; do
		length=$(image_length "$(sysupgrade_dir)/$member" "$@")
		if [ "$length" = 0 ]; then
			missing="$missing $member"
		else
			present="$present $member"
		fi
	done

	[ -n "$present" ] || return 1
	[ -z "$missing" ] || {
		v "Malformed Nix payload: missing or empty:$missing"
		return 2
	}

	return 0
}

init_nix_store() {
	local default_servers_config

	default_servers_config="$NIX_STAGE_DIR/servers.json.default"
	[ -n "$BOS_NIX_DEFAULT_SERVERS_CONFIG" ] \
		&& default_servers_config="$BOS_NIX_DEFAULT_SERVERS_CONFIG"

	# No promoted store exists, so there is nothing valid to
	# preserve; --wipe clears partial wreckage left by an earlier
	# aborted attempt. The factory tarball is downloaded to
	# /mnt/data because it does not fit into tmpfs. A valid
	# /etc/nix-upgrade/servers.json wins over the staged default,
	# which is a read-only fallback and is never persisted.
	v "Initializing Nix store..."
	nix_cli init \
		--default-servers-config "$default_servers_config" \
		--firmware "$UPGRADE_FW_VERSION" \
		--download-dir /mnt/data \
		--wipe || {
		v "Unable to initialize the Nix store"
		return 1
	}

	return 0
}

stage_nix_upgrade() {
	local default_servers_config

	default_servers_config="$NIX_STAGE_DIR/servers.json.default"
	[ -n "$BOS_NIX_DEFAULT_SERVERS_CONFIG" ] \
		&& default_servers_config="$BOS_NIX_DEFAULT_SERVERS_CONFIG"

	# Stage the new firmware's application set as a `next` profile
	# generation bound to that firmware's BOS version; its boot-time
	# activator consumes the marker after the reboot. The --firmware
	# scope resolves each feed-linked server to the incoming
	# firmware's package index. A valid runtime
	# /etc/nix-upgrade/servers.json wins over the staged default,
	# which is a read-only fallback and is never persisted.
	v "Staging Nix profile for the new firmware..."
	set --
	if [ -f "$NIX_PENDING_INSTALL" ]; then
		v "Including pending package installs from $NIX_PENDING_INSTALL"
		set -- --install-from "$NIX_PENDING_INSTALL"
	fi
	nix_cli upgrade \
		--default-servers-config "$default_servers_config" \
		--firmware "$UPGRADE_FW_VERSION" \
		--next-boot \
		"$@" || {
		v "Unable to stage the Nix profile for the new firmware"
		return 1
	}

	return 0
}

prepare_nix_profile() {
	if [ "$BOS_NIX_SKIP" = 1 ]; then
		echo "WARNING: BOS_NIX_SKIP=1: skipping Nix store/profile preparation"
		echo "WARNING: the flashed firmware may lack a matching Nix profile"
		return 0
	fi

	if [ "$(cat "$NIX_PREPARED_MARKER" 2>/dev/null)" = "$UPGRADE_FW_VERSION" ]; then
		v "Nix profile already prepared for '$UPGRADE_FW_VERSION', skipping"
		return 0
	fi

	# Stale files from an earlier aborted attempt must not leak into
	# this run.
	rm -rf "$NIX_STAGE_DIR" && mkdir -p "$NIX_STAGE_DIR" || return 1
	stage_nix_payload_file "bmc-nix-cli" "$@" || return 1
	stage_nix_payload_file "servers.json.default" "$@" || return 1
	chmod +x "$NIX_STAGE_DIR/bmc-nix-cli" || return 1

	# nix-store ships in the Nix profile, not as a system binary, and
	# its bin directory reaches PATH only through login shells
	# (/root/.profile). Extend PATH so the realisation spawned by
	# upgrade also works when sysupgrade runs from a non-login shell.
	# The gcroots path is the in-store fallback for a boot whose
	# activation never created /run/current-profile; it resolves once
	# the /nix bind mount is in place.
	export PATH="$PATH:/run/current-profile/bin"
	export PATH="$PATH:/nix/var/nix/gcroots/profiles/bmc/current/bin"

	# fsck/format/mount the data partition as needed; a no-op when
	# the right partition is already mounted. The store-presence
	# check below is meaningful only after this verified mount.
	nix_cli prepare-data-partition || {
		v "Unable to prepare the data partition"
		return 1
	}

	# Exit 0 also vouches for the /nix bind mount, so the staging
	# below can build on the persistent store directly.
	nix_cli is-initialized --data-dir /mnt/data
	case $? in
	0)
		stage_nix_upgrade || return 1
		;;
	1|3)
		# 1 means no promoted store; 3 means a promoted store exists
		# but /nix is unmounted. Since Nix-era systems always mount
		# /nix on boot, both states require initialization here.
		init_nix_store || return 1
		;;
	*)
		# Only the statuses above carry an answer; never route an
		# unexpected status into the wipe branch.
		v "Unable to determine the Nix store state"
		return 1
		;;
	esac

	# The marker is an optimization: if it cannot be written, the
	# second validation pass just re-stages.
	printf '%s\n' "$UPGRADE_FW_VERSION" > "$NIX_PREPARED_MARKER" \
		|| v "Unable to record the Nix staging marker"

	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

	validate_nix_payload "$@"
	case $? in
	0) ;;
	1)
		v "No Nix payload in the image, skipping Nix staging"
		return 0
		;;
	*)
		return 1
		;;
	esac

	# `sysupgrade -T` must stay side-effect free; the Nix staging
	# downloads packages and writes to the data partition.
	[ "$TEST" = 1 ] && return 0

	prepare_nix_profile "$@" || 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
}
