UPGRADE_FW_MAJOR="2024-11-20-0-8e3abf6a-24.10-plus"
UPGRADE_FW_VERSION="2026-06-09-0-0e8e52af-26.07-plus-nightly"
UPGRADE_FW_REQUIRE="2024-11-20-0-8e3abf6a-24.10-plus"

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

BLOCK_SIZE=4096

# Hard-coded URL of the Nix initialization tarball. The tarball is
# expected to carry an appended usign signature trailer (fwtool-style,
# same pattern used for signed sysupgrade images on OpenWrt). The
# default is overridable via the environment for manual testing /
# staging builds.
NIX_INIT_TARBALL_URL="${NIX_INIT_TARBALL_URL:-https://files.ditigal.xyz/nix-init.tar.gz}"
NIX_INIT_TARBALL_SIG_URL="${NIX_INIT_TARBALL_SIG_URL:-${NIX_INIT_TARBALL_URL}.sig}"

# Minimum total size (in 1K blocks) required on /mnt/data to host the
# Nix store.
NIX_DATA_MIN_KB=${NIX_DATA_MIN_KB:-2097152}

# usign public key that authenticates the Nix initialization tarball.
# Packing it inline is safe: the COMMAND itself travels inside a
# signed sysupgrade image, so this key inherits that outer trust
# chain and bootstraps the Nix-init trust root without any separate
# distribution channel.
NIX_INIT_TARBALL_PUBKEY='untrusted comment: BOS Nix init signing key public key
RWR0LSX/RMCJy8JH9S9sUwxXWgNyIL+HgR58inqRD76YOkMoOJNF2psJ'

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
}

# Return 0 if $1 is a current mount point. Uses /proc/mounts rather
# than busybox mountpoint(1), which is not guaranteed to be present
# in this target's busybox config.
is_mountpoint() {
	awk -v m="$1" '$2==m {found=1} END{exit !found}' /proc/mounts
}

# Mount /mnt/data if no-one else already did. This mirrors
# ensure_data_partition_mounted from
# bos-packages/bos-hooks/files/pre-exec-mount.init so sysupgrade from
# the first shipped firmware (which never mounted /mnt/data) still
# works.
ensure_data_mounted() {
	local partition="/dev/mmcblk0p4"
	local mount_point="/mnt/data"
	local fs_type

	is_mountpoint "$mount_point" && return 0

	[ -b "$partition" ] || {
		v "Data partition $partition is missing"
		return 1
	}

	fs_type=$(blkid -o value -s TYPE "$partition" 2>/dev/null)
	if [ -z "$fs_type" ]; then
		v "Data partition not formatted — creating ext4"
		mkfs.ext4 -F "$partition" || {
			v "Failed to format data partition"
			return 1
		}
	fi

	fsck.ext4 -f -y -C0 "$partition"
	case $? in
		0 | 1 | 2) ;;
		*)
			v "fsck.ext4 failed on $partition"
			return 1
			;;
	esac

	mkdir -p "$mount_point" || {
		v "Failed to create $mount_point"
		return 1
	}
	mount -t ext4 "$partition" "$mount_point" || {
		v "Failed to mount $partition on $mount_point"
		return 1
	}

	return 0
}

prepare_nix_store() {
  local nix_marker="/tmp/.nix-init-staged"
	local nix_tarball="/mnt/data/nix-init.tar.gz"
	local nix_tarball_sig="/tmp/nix-init.sig"
	local nix_tarball_pubkey="/tmp/nix-init.pub"
	local data_total_kb

  # Skip if a prior phase of the same sysupgrade already staged the store.
  # /tmp is tmpfs so this marker dies at reboot on its own. Before we skip,
  # sanity-check that the expected artifacts are actually present — if the
  # marker somehow outlived a failed stage, we must not silently let the
  # upgrade proceed with an empty store.
  if [ -f "$nix_marker" ]; then
      if [ ! -d "/mnt/data/nix/store" ]; then
          v "Nix init marker present but /mnt/data/nix/store is missing"
          return 1
      fi
      if [ ! -e "/mnt/data/nix/var/nix/gcroots/profiles/bmc/1-link" ]; then
          v "Nix init marker present but /mnt/data/nix/var/nix/gcroots/profiles/bmc/1-link is missing"
          return 1
      fi
      v "Nix store already staged this sysupgrade, skipping"
      return 0
  fi

	# Make sure /mnt/data is a real mounted partition; mount it
	# ourselves if the running firmware never did.
	ensure_data_mounted || return 1

	# Sanity check: the partition must be large enough to hold the
	# tarball and the extracted Nix store.
	data_total_kb=$(df -Pk "/mnt/data" | awk 'NR==2 { print $2 }')
	[ -n "$data_total_kb" ] && \
	[ "$data_total_kb" -ge "$NIX_DATA_MIN_KB" ] || {
		v "/mnt/data is too small (total ${data_total_kb:-?} KiB, need $NIX_DATA_MIN_KB KiB)"
		return 1
	}

	# On a Nix-enabled running system /nix is a bind mount of
	# /mnt/data/nix; unwind it first so we can remove the backing
	# directory and rebuild the mount point from scratch. A failed
	# umount is fatal — we cannot safely wipe the store out from under
	# a live mount.
	if is_mountpoint "/nix"; then
		umount "/nix" || {
			v "Unable to unmount /nix"
			return 1
		}
	fi

	# Remove any previous persistent Nix store. This covers the case
	# where the user factory-reset to a firmware version that did not
	# yet support Nix: stale content may still live on /mnt/data and
	# would clash with the fresh bootstrap we are about to populate.
	rm -rf "/mnt/data/nix" || {
		v "Unable to remove /mnt/data/nix"
		return 1
	}

	# Remove /nix so we can rebuild the mount point from scratch.
	rm -rf "/nix" || {
		v "Unable to remove /nix"
		return 1
	}

	# Fetch the initialization tarball onto the persistent partition so
	# it survives the later switch to RAM-FS in package_do_upgrade.
	# Drop any stale copy from a previous (possibly interrupted) run
	# before downloading, so a failed wget cannot leave a truncated
	# tarball sitting on /mnt/data.
	v "Downloading Nix initialization tarball..."
	rm -f "$nix_tarball"
	local attempt
	for attempt in 1 2 3 4 5; do
		wget --timeout=60 -O "$nix_tarball" \
			"$NIX_INIT_TARBALL_URL" && break
		[ "$attempt" = 5 ] && {
			v "Unable to download Nix initialization tarball"
			return 1
		}
		v "Download attempt $attempt failed, retrying..."
		sleep 2
	done

	# Fetch the detached signify signature alongside the tarball.
	v "Downloading Nix initialization tarball signature..."
	rm -f "$nix_tarball_sig"
	for attempt in 1 2 3 4 5; do
		wget --timeout=60 -O "$nix_tarball_sig" \
			"$NIX_INIT_TARBALL_SIG_URL" && break
		[ "$attempt" = 5 ] && {
			v "Unable to download Nix initialization tarball signature"
			return 1
		}
		v "Signature download attempt $attempt failed, retrying..."
		sleep 2
	done

	# Verify the downloaded tarball against the embedded public key.
	# usign consumes the detached signature and the tarball directly.
	# This ties the downloaded artifact back to the signed COMMAND
	# (and thus the signed sysupgrade image) that is executing now.
	# printf '%s' "$NIX_INIT_TARBALL_PUBKEY" > "$nix_tarball_pubkey" || {
	# 	v "Unable to stage Nix init public key"
	# 	return 1
	# }
	# usign -V -m "$nix_tarball" -x "$nix_tarball_sig" \
	# 	-p "$nix_tarball_pubkey" || {
	# 	v "Nix initialization tarball signature verification failed"
	# 	return 1
	# }

	# Prepare the bind mount so that the tarball extraction below lands
	# on persistent storage via /nix -> /mnt/data/nix.
	mkdir -p "/mnt/data/nix" "/nix" || {
		v "Unable to create Nix mount points"
		return 1
	}
	mount --bind "/mnt/data/nix" "/nix" || {
		v "Unable to bind-mount /mnt/data/nix to /nix"
		return 1
	}

	# Extract the tarball at / so /nix/store, /nix/var, ... land inside
	# the bind-mounted persistent directory. Activation is intentionally
	# NOT performed here — it is postponed to the next boot.
	tar -xzf "$nix_tarball" -C / || {
		v "Unable to extract Nix initialization tarball"
		return 1
	}

	# Reclaim the persistent disk space now that the store has been
	# materialised under /mnt/data/nix.
	rm -f "$nix_tarball" "$nix_tarball_sig"

  # Ensure all is on the disk
  sync

  # Set marker for nix-factory-reset. If not set, it would remove
  # the store on next boot.
  fw_setenv nix_init 1

  touch "$nix_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

	prepare_nix_store || 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
}
