#!/bin/sh

# ota-upgrade 1.0.0
# Copyright (c)2026 John Lawson & Sons
# All Rights Reserved

set -eu

OS_RELEASE="/etc/os-release"
BOOT_MARKER="isignage-boot.marker"
OTA_CHANNEL="${OTA_CHANNEL:-latest}"
ANSI_BOLD="$(printf '\033[1m')"
ANSI_RED="$(printf '\033[1;91m')"
ANSI_GREEN="$(printf '\033[1;92m')"
ANSI_YELLOW="$(printf '\033[1;93m')"
ANSI_CYAN="$(printf '\033[1;96m')"
ANSI_RESET="$(printf '\033[0m')"

ota_info() {
	printf '%s[OTA]%s %s\n' "$ANSI_GREEN" "$ANSI_RESET" "$*"
}

ota_warn() {
	printf '%s[OTA]%s %s\n' "$ANSI_YELLOW" "$ANSI_RESET" "$*"
}

ota_error() {
	printf '%s[OTA ERROR]%s %s\n' "$ANSI_RED" "$ANSI_RESET" "$*" >&2
}

ota_kv() {
	printf '  %s%-18s%s %s\n' "$ANSI_CYAN" "$1" "$ANSI_RESET" "$2"
}

[ -f "$OS_RELEASE" ] || {
	ota_error "$OS_RELEASE missing"
	exit 1
}

CURRENT_VERSION="$(awk -F= '$1=="VERSION"{gsub(/"/,"",$2); print $2}' "$OS_RELEASE")"
MACHINE="$(awk -F= '$1=="MACHINE"{gsub(/"/,"",$2); print $2}' "$OS_RELEASE")"

[ -n "$CURRENT_VERSION" ] || {
	ota_error "VERSION missing from $OS_RELEASE"
	exit 1
}

[ -n "$MACHINE" ] || {
	ota_error "MACHINE missing from $OS_RELEASE"
	exit 1
}

DO_REBOOT=""
LOCAL_DIR=""
SECURE_PREP_MODE=""
SECURE_STATUS_ONLY=""

usage() {
	printf '%sUsage:%s %s\n' "$ANSI_YELLOW" "$ANSI_RESET" "$0 [--local-dir DIR] [--reboot] [--secure-prepare] [--secure-prepare-pi4] [--secure-prepare-pi5] [--secure-status] [--help]"
}

help() {
	printf '%sota-upgrade%s\n' "$ANSI_BOLD" "$ANSI_RESET"
	printf 'Upgrade inactive rootfs and refresh boot content for %simage-signage%s devices.\n' "$ANSI_CYAN" "$ANSI_RESET"
	printf '\n'
	printf '%sUsage%s\n' "$ANSI_YELLOW" "$ANSI_RESET"
	printf '  %s%s%s\n' "$ANSI_GREEN" "$0 [--local-dir DIR] [--reboot] [--secure-prepare] [--secure-prepare-pi4] [--secure-prepare-pi5] [--secure-status] [--help]" "$ANSI_RESET"
	printf '\n'
	printf '%sOptions%s\n' "$ANSI_YELLOW" "$ANSI_RESET"
	printf '  %s--local-dir DIR%s         Load OTA bundle files from a local directory instead of downloading them.\n' "$ANSI_GREEN" "$ANSI_RESET"
	printf '  %s--reboot%s                Reboot automatically after staging the update.\n' "$ANSI_GREEN" "$ANSI_RESET"
	printf '  %s--secure-prepare%s        Stage secure-preparation content and accept the bundle mode from secure.env.\n' "$ANSI_GREEN" "$ANSI_RESET"
	printf '  %s--secure-prepare-pi4%s    Stage secure-preparation content and require a Pi 4 secure bundle.\n' "$ANSI_GREEN" "$ANSI_RESET"
	printf '  %s--secure-prepare-pi5%s    Stage secure-preparation content and require a Pi 5 secure bundle.\n' "$ANSI_GREEN" "$ANSI_RESET"
	printf '  %s--secure-status%s         Print the current secure-preparation state from /boot/ota and exit.\n' "$ANSI_GREEN" "$ANSI_RESET"
	printf '  %s--help%s                  Show this help text.\n' "$ANSI_GREEN" "$ANSI_RESET"
	printf '\n'
	printf '%sNotes%s\n' "$ANSI_YELLOW" "$ANSI_RESET"
	printf '  - Secure prepare stages files only. It does %snot%s enrol or lock the device.\n' "$ANSI_BOLD" "$ANSI_RESET"
	printf '  - %s--secure-status%s is read-only.\n' "$ANSI_GREEN" "$ANSI_RESET"
}

env_file_value() {
	file="$1"
	key="$2"
	awk -F= -v key="$key" '$1==key {print substr($0, index($0, "=")+1)}' "$file" 2>/dev/null || true
}

print_secure_status() {
	mount_boot=""
	if ! mountpoint -q /boot; then
		mount_boot="1"
		mount "$BOOT_PART" /boot
	fi

	status="$(cat "$OTA_BOOT_DIR/secure-state" 2>/dev/null || echo "not-prepared")"
	mode="$(cat "$OTA_BOOT_DIR/secure-mode" 2>/dev/null || true)"
	machine="$(cat "$OTA_BOOT_DIR/secure-machine" 2>/dev/null || true)"
	hash="$(cat "$OTA_BOOT_DIR/secure-public-key-hash" 2>/dev/null || true)"
	prepared_at="$(cat "$OTA_BOOT_DIR/secure-prepared-at" 2>/dev/null || true)"

	ota_info "Secure boot status"
	ota_kv "state" "$status"
	[ -n "$mode" ] && ota_kv "mode" "$mode"
	[ -n "$machine" ] && ota_kv "machine" "$machine"
	[ -n "$hash" ] && ota_kv "public-key hash" "$hash"
	[ -n "$prepared_at" ] && ota_kv "prepared at" "$prepared_at"

	if [ -d /boot/secure-boot ]; then
		ota_kv "staged dir" "/boot/secure-boot"
	else
		ota_kv "staged dir" "missing"
	fi

	if [ "$mount_boot" = "1" ]; then
		umount /boot 2>/dev/null || true
	fi
}

while [ "$#" -gt 0 ]; do
	case "$1" in
		--help|-h)
			help
			exit 0
			;;
		--reboot)
			DO_REBOOT="--reboot"
			;;
		--secure-prepare)
			SECURE_PREP_MODE="auto"
			;;
		--secure-prepare-pi4)
			SECURE_PREP_MODE="pi4"
			;;
		--secure-prepare-pi5)
			SECURE_PREP_MODE="pi5"
			;;
		--secure-status)
			SECURE_STATUS_ONLY="1"
			;;
		--local-dir)
			shift
			[ "$#" -gt 0 ] || {
				usage
				exit 1
			}
			LOCAL_DIR="$1"
			;;
		*)
			usage
			exit 1
			;;
	esac
	shift
done

[ "$(id -u)" = "0" ] || {
	ota_error "must run as root"
	exit 1
}

command -v curl >/dev/null || { ota_error "curl missing"; exit 1; }
command -v findmnt >/dev/null || { ota_error "findmnt missing"; exit 1; }
command -v mkfs.btrfs >/dev/null || { ota_error "mkfs.btrfs missing"; exit 1; }
command -v mountpoint >/dev/null || { ota_error "mountpoint missing"; exit 1; }
command -v sha256sum >/dev/null || { ota_error "sha256sum missing"; exit 1; }
command -v tar >/dev/null || { ota_error "tar missing"; exit 1; }

BASE_URL="https://upgrades.isignage.app/${MACHINE}/${OTA_CHANNEL}"

partition_path() {
	disk="$1"
	index="$2"

	case "$disk" in
		/dev/mmcblk*|/dev/nvme*)
			echo "${disk}p${index}"
			;;
		*)
			echo "${disk}${index}"
			;;
	esac
}

preferred_boot_disks() {
	seen=""

	add_disk() {
		disk="$1"

		[ -n "$disk" ] || return 0
		[ -b "$disk" ] || return 0

		case " $seen " in
			*" $disk "*)
				return 0
				;;
		esac

		printf '%s\n' "$disk"
		seen="$seen $disk"
	}

	for disk in /dev/nvme0n1 /dev/mmcblk1 /dev/mmcblk0 /dev/mmcblk2; do
		add_disk "$disk"
	done

	for sysdisk in /sys/class/block/*; do
		name="${sysdisk##*/}"

		[ -e "$sysdisk" ] || continue
		[ ! -e "$sysdisk/partition" ] || continue

		case "$name" in
			loop*|ram*|zram*|dm-*|md*|sr*|mtdblock*|nbd*)
				continue
				;;
		esac

		add_disk "/dev/$name"
	done
}

partition_disk() {
	part="$1"

	if command -v readlink >/dev/null 2>&1; then
		part="$(readlink -f "$part" 2>/dev/null || echo "$part")"
	fi

	case "$part" in
		/dev/mmcblk*p[0-9]*|/dev/nvme*n*p[0-9]*)
			echo "${part%p*}"
			;;
		/dev/*[0-9])
			echo "${part%[0-9]}"
			;;
		*)
			return 1
			;;
	esac
}

partition_mount_source() {
	part="$1"
	echo "$part"
}

disk_removable() {
	disk="$1"
	name="${disk##*/}"
	value="$(cat "/sys/class/block/$name/removable" 2>/dev/null || true)"

	case "$value" in
		0)
			return 1
			;;
		*)
			return 0
			;;
	esac
}

find_boot_config_relpath() {
	for rel in cmdline.txt extlinux/extlinux.conf; do
		if [ -f "/boot/$rel" ]; then
			echo "$rel"
			return 0
		fi
	done

	return 1
}

boot_disk_score() {
	disk="$1"
	score=0
	part1="$(partition_path "$disk" 1)"
	part2="$(partition_path "$disk" 2)"
	part3="$(partition_path "$disk" 3)"
	part4="$(partition_path "$disk" 4)"

	[ -b "$part1" ] || {
		echo 0
		return 0
	}

	mkdir -p /tmp/ota-boot-test

	if ! mount -t vfat "$part1" /tmp/ota-boot-test 2>/dev/null; then
		echo 0
		return 0
	fi

	if [ ! -f "/tmp/ota-boot-test/$BOOT_MARKER" ]; then
		umount /tmp/ota-boot-test 2>/dev/null || true
		echo 0
		return 0
	fi

	if ! disk_removable "$disk"; then
		score=$((score + 1000))
	fi

	if [ -f /tmp/ota-boot-test/cmdline.txt ] || [ -f /tmp/ota-boot-test/extlinux/extlinux.conf ]; then
		score=$((score + 50))
	fi

	if [ -f /tmp/ota-boot-test/config.txt ]; then
		score=$((score + 10))
	fi

	if [ -f /tmp/ota-boot-test/fitImage ] || [ -f /tmp/ota-boot-test/uImage ] || [ -f /tmp/ota-boot-test/Image ]; then
		score=$((score + 25))
	fi

	if [ -f /tmp/ota-boot-test/initramfs-rollback.cpio ] || [ -f /tmp/ota-boot-test/initramfs-rollback.cpio.gz ]; then
		score=$((score + 20))
	fi

	if [ -d /tmp/ota-boot-test/ota ]; then
		score=$((score + 10))
	fi

	if [ -d /tmp/ota-boot-test/recovery ]; then
		score=$((score + 10))
	fi

	umount /tmp/ota-boot-test 2>/dev/null || true

	[ -b "$part2" ] && score=$((score + 15))
	[ -b "$part3" ] && score=$((score + 15))
	[ -b "$part4" ] && score=$((score + 10))

	echo "$score"
}

find_boot_disk() {
	boot_source="$(findmnt -no SOURCE /boot 2>/dev/null || true)"
	if [ -n "$boot_source" ]; then
		disk="$(partition_disk "$boot_source" 2>/dev/null || true)"
		if [ -n "$disk" ] && [ -b "$(partition_path "$disk" 1)" ]; then
			echo "$disk"
			return 0
		fi
	fi

	root_source="$(findmnt -no SOURCE / 2>/dev/null || true)"
	case "$root_source" in
		/dev/root)
			;;
		*)
			if [ -n "$root_source" ]; then
				disk="$(partition_disk "$root_source" 2>/dev/null || true)"
				if [ -n "$disk" ] && [ -b "$(partition_path "$disk" 1)" ]; then
					echo "$disk"
					return 0
				fi
			fi
			;;
	esac

	best_disk=""
	best_score=0

	for d in $(preferred_boot_disks); do
		score="$(boot_disk_score "$d")"

		case "$score" in
			''|*[!0-9]*)
				score=0
				;;
		esac

		if [ "$score" -gt "$best_score" ]; then
			best_disk="$d"
			best_score="$score"
		fi
	done

	if [ -n "$best_disk" ]; then
		echo "$best_disk"
		return 0
	fi

	return 1
}

DISK="$(find_boot_disk || true)"

[ -n "$DISK" ] || {
	ota_error "could not detect boot disk"
	exit 1
}

BOOT_PART="$(partition_path "$DISK" 1)"
ROOT_A_PART="$(partition_path "$DISK" 2)"
ROOT_B_PART="$(partition_path "$DISK" 3)"
HOME_PART="$(partition_path "$DISK" 4)"

WORKDIR="/home/update"
TARGET_MNT="/mnt/ota-target"
OTA_BOOT_DIR="/boot/ota"
RECOVERY_DIR="/boot/recovery"

if [ -n "$SECURE_STATUS_ONLY" ]; then
	print_secure_status
	exit 0
fi

current_root_from_metadata() {
	for file in "$OTA_BOOT_DIR/current-root" "$OTA_BOOT_DIR/last-booted-root"; do
		if [ -f "$file" ]; then
			root="$(cat "$file" 2>/dev/null || true)"
			case "$root" in
				"$ROOT_A_PART"|"$ROOT_B_PART")
					echo "$root"
					return 0
					;;
			esac
		fi
	done

	root="$(awk -F= '$1=="INSTALLED_ROOT"{gsub(/"/,"",$2); print $2}' "$OS_RELEASE" 2>/dev/null || true)"
	case "$root" in
		"$ROOT_A_PART"|"$ROOT_B_PART")
			echo "$root"
			return 0
			;;
	esac

	return 1
}

stage_recovery_image() {
	ota_info "Staging recovery rootfs..."

	mkdir -p "$RECOVERY_DIR"

	rm -f "$RECOVERY_DIR/rootfs.tar.xz.new"
	rm -f "$RECOVERY_DIR/rootfs.tar.xz.sha256.new"
	rm -f "$RECOVERY_DIR/version.new"

	cp "$WORKDIR/rootfs.tar.xz" "$RECOVERY_DIR/rootfs.tar.xz.new"

	cd "$RECOVERY_DIR"

	sha256sum rootfs.tar.xz.new > rootfs.tar.xz.sha256.new
	sha256sum -c rootfs.tar.xz.sha256.new

	cat > version.new <<EOF
VERSION=$BUNDLE_VERSION
MACHINE=$BUNDLE_MACHINE
STAGED_AT=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
ROOT=$NEW_ROOT
EOF

	cd "$WORKDIR"

	sync

	ota_info "Recovery rootfs staged"
}

cleanup() {
	if mountpoint -q "$TARGET_MNT"; then
		ota_warn "Cleanup: unmounting $TARGET_MNT"
		umount "$TARGET_MNT" || true
	fi
}

trap cleanup EXIT

ota_info "Starting OTA upgrade"
ota_kv "machine" "$MACHINE"
ota_kv "current version" "$CURRENT_VERSION"
ota_kv "update channel" "$OTA_CHANNEL"
ota_kv "update url" "$BASE_URL"

mkdir -p "$WORKDIR" "$TARGET_MNT"

ota_info "Ensuring /boot is mounted..."
mkdir -p /boot

if ! mountpoint -q /boot; then
	mount "$BOOT_PART" /boot
fi

mkdir -p "$OTA_BOOT_DIR"
mkdir -p "$RECOVERY_DIR"

cd "$WORKDIR"

ota_info "Cleaning old downloads..."
rm -f rootfs.tar.xz boot.tar.xz SHA256SUMS os-release secure.env secure.tar.xz

if [ -n "$LOCAL_DIR" ]; then
	ota_info "Loading local update bundle from $LOCAL_DIR"

	[ -d "$LOCAL_DIR" ] || {
		ota_error "local bundle directory missing: $LOCAL_DIR"
		exit 1
	}

	for f in rootfs.tar.xz boot.tar.xz SHA256SUMS os-release; do
		[ -f "$LOCAL_DIR/$f" ] || {
			ota_error "missing local bundle file: $LOCAL_DIR/$f"
			exit 1
		}
		cp -f "$LOCAL_DIR/$f" "./$f"
	done
	if [ -n "$SECURE_PREP_MODE" ]; then
		for f in secure.env secure.tar.xz; do
			[ -f "$LOCAL_DIR/$f" ] || {
				ota_error "missing local secure bundle file: $LOCAL_DIR/$f"
				exit 1
			}
			cp -f "$LOCAL_DIR/$f" "./$f"
		done
	fi
else
	ota_info "Downloading update bundle..."
	curl -fL -o rootfs.tar.xz "$BASE_URL/rootfs.tar.xz"
	curl -fL -o boot.tar.xz "$BASE_URL/boot.tar.xz"
	curl -fL -o SHA256SUMS "$BASE_URL/SHA256SUMS"
	curl -fL -o os-release "$BASE_URL/os-release"
	if [ -n "$SECURE_PREP_MODE" ]; then
		curl -fL -o secure.env "$BASE_URL/secure.env"
		curl -fL -o secure.tar.xz "$BASE_URL/secure.tar.xz"
	fi
fi

BUNDLE_VERSION="$(awk -F= '$1=="VERSION"{gsub(/"/,"",$2); print $2}' os-release)"
BUNDLE_MACHINE="$(awk -F= '$1=="MACHINE"{gsub(/"/,"",$2); print $2}' os-release)"

[ -n "$BUNDLE_VERSION" ] || {
	ota_error "VERSION missing from downloaded os-release"
	exit 1
}

[ -n "$BUNDLE_MACHINE" ] || {
	ota_error "MACHINE missing from downloaded os-release"
	exit 1
}

[ "$BUNDLE_MACHINE" = "$MACHINE" ] || {
	ota_error "bundle machine mismatch: expected $MACHINE got $BUNDLE_MACHINE"
	exit 1
}

ota_kv "bundle version" "$BUNDLE_VERSION"

ota_info "Verifying checksums..."
CHECKSUM_FILE="SHA256SUMS"
if [ -z "$SECURE_PREP_MODE" ]; then
	CHECKSUM_FILE="$WORKDIR/SHA256SUMS.standard"
	grep -vE '  (secure\.env|secure\.tar\.xz)$' SHA256SUMS > "$CHECKSUM_FILE"
fi
sha256sum -c "$CHECKSUM_FILE"

SECURE_BUNDLE_MODE=""
SECURE_BUNDLE_MACHINE=""
SECURE_PUBLIC_KEY_HASH=""
if [ -n "$SECURE_PREP_MODE" ]; then
	[ -f secure.env ] || {
		ota_error "secure.env missing"
		exit 1
	}
	[ -f secure.tar.xz ] || {
		ota_error "secure.tar.xz missing"
		exit 1
	}

	SECURE_ENABLED="$(env_file_value secure.env SECURE_PREP)"
	SECURE_BUNDLE_MODE="$(env_file_value secure.env SECURE_PREP_MODE)"
	SECURE_BUNDLE_MACHINE="$(env_file_value secure.env MACHINE)"
	SECURE_PUBLIC_KEY_HASH="$(env_file_value secure.env PUBLIC_KEY_HASH)"

	[ "$SECURE_ENABLED" = "1" ] || {
		ota_error "secure.env does not describe a secure preparation bundle"
		exit 1
	}

	[ "$SECURE_BUNDLE_MACHINE" = "$MACHINE" ] || {
		ota_error "secure bundle machine mismatch: expected $MACHINE got $SECURE_BUNDLE_MACHINE"
		exit 1
	}

	if [ "$SECURE_PREP_MODE" != "auto" ] && [ "$SECURE_BUNDLE_MODE" != "$SECURE_PREP_MODE" ]; then
		ota_error "secure bundle mode mismatch: requested $SECURE_PREP_MODE got $SECURE_BUNDLE_MODE"
		exit 1
	fi

	ota_kv "secure bundle" "$SECURE_BUNDLE_MODE"
	[ -n "$SECURE_PUBLIC_KEY_HASH" ] && ota_kv "public-key hash" "$SECURE_PUBLIC_KEY_HASH"
fi

ACTIVE_SOURCE="$(findmnt -no SOURCE / || true)"
if [ -n "$ACTIVE_SOURCE" ] && [ "$ACTIVE_SOURCE" != "/dev/root" ] && command -v readlink >/dev/null 2>&1; then
	ACTIVE_SOURCE="$(readlink -f "$ACTIVE_SOURCE" 2>/dev/null || echo "$ACTIVE_SOURCE")"
fi

case "$ACTIVE_SOURCE" in
	"$ROOT_A_PART")
		ACTIVE_PART="$ACTIVE_SOURCE"
		TARGET_LABEL="rootB"
		TARGET_PART="$ROOT_B_PART"
		PREV_ROOT="$ROOT_A_PART"
		NEW_ROOT="$ROOT_B_PART"
		;;
	"$ROOT_B_PART")
		ACTIVE_PART="$ROOT_B_PART"
		TARGET_LABEL="rootA"
		TARGET_PART="$ROOT_A_PART"
		PREV_ROOT="$ROOT_B_PART"
		NEW_ROOT="$ROOT_A_PART"
		;;
	"/dev/root")
		ACTIVE_PART="$(current_root_from_metadata || true)"

		case "$ACTIVE_PART" in
			"$ROOT_B_PART")
				TARGET_LABEL="rootA"
				TARGET_PART="$ROOT_A_PART"
				PREV_ROOT="$ROOT_B_PART"
				NEW_ROOT="$ROOT_A_PART"
				;;
			"$ROOT_A_PART"|"")
				ACTIVE_PART="$ROOT_A_PART"
				TARGET_LABEL="rootB"
				TARGET_PART="$ROOT_B_PART"
				PREV_ROOT="$ROOT_A_PART"
				NEW_ROOT="$ROOT_B_PART"
				;;
			*)
				ota_error "active root metadata is not recognised"
				ota_kv "active metadata" "$ACTIVE_PART"
				exit 1
				;;
		esac
		;;
	*)
		ota_error "active root is not recognised"
		ota_kv "active source" "$ACTIVE_SOURCE"
		exit 1
		;;
esac

ota_kv "active root" "$ACTIVE_PART"
ota_kv "target partition" "$TARGET_PART"
ota_kv "previous root" "$PREV_ROOT"
ota_kv "new root" "$NEW_ROOT"

if findmnt -rn "$TARGET_PART" >/dev/null 2>&1; then
	ota_error "$TARGET_PART is currently mounted"
	findmnt "$TARGET_PART" || true
	exit 1
fi

BOOT_CONFIG_REL="$(find_boot_config_relpath || true)"

[ -n "$BOOT_CONFIG_REL" ] || {
	ota_error "could not detect boot config in /boot"
	exit 1
}

ota_info "Preserving boot config: /boot/$BOOT_CONFIG_REL"
mkdir -p "$WORKDIR/$(dirname "$BOOT_CONFIG_REL")"
cp -a "/boot/$BOOT_CONFIG_REL" "$WORKDIR/$BOOT_CONFIG_REL.before"

ota_info "Formatting inactive rootfs partition..."
mkfs.btrfs -f -L "$TARGET_LABEL" "$TARGET_PART"

ota_info "Mounting target rootfs..."
mount "$TARGET_PART" "$TARGET_MNT"

ota_info "Extracting new rootfs..."
tar --numeric-owner -xpf rootfs.tar.xz -C "$TARGET_MNT"

ota_info "Writing target fstab..."
ROOT_MOUNT_SOURCE="$(partition_mount_source "$NEW_ROOT")"
BOOT_MOUNT_SOURCE="$(partition_mount_source "$BOOT_PART")"
HOME_MOUNT_SOURCE="$(partition_mount_source "$HOME_PART")"
cat > "$TARGET_MNT/etc/fstab" <<EOF
$ROOT_MOUNT_SOURCE /          btrfs  defaults                           1 1
proc           /proc         proc   defaults                           0 0
devpts         /dev/pts      devpts mode=0620,ptmxmode=0666,gid=5      0 0
tmpfs          /run          tmpfs  mode=0755,nodev,nosuid,strictatime 0 0
tmpfs          /var/volatile tmpfs  defaults                           0 0
$BOOT_MOUNT_SOURCE /boot     vfat   defaults                           0 0
$HOME_MOUNT_SOURCE /home     btrfs  defaults                           0 0
EOF

ota_info "Installing version metadata..."
mkdir -p "$TARGET_MNT/etc"

cat os-release > "$TARGET_MNT/etc/os-release"

cat >> "$TARGET_MNT/etc/os-release" <<EOF
INSTALLED_AT=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
INSTALLED_ROOT=$NEW_ROOT
PREVIOUS_ROOT=$PREV_ROOT
EOF

ota_info "Installing safe /boot files..."
tar --no-same-owner \
	--exclude='./cmdline.txt' \
	--exclude='cmdline.txt' \
	--exclude='./start*.elf' \
	--exclude='start*.elf' \
	--exclude='./fixup*.dat' \
	--exclude='fixup*.dat' \
	--exclude='./bootcode.bin' \
	--exclude='bootcode.bin' \
	-xpf boot.tar.xz -C /boot

if [ -n "$SECURE_PREP_MODE" ]; then
	ota_info "Staging secure-boot preparation files..."
	rm -rf /boot/secure-boot
	tar --no-same-owner -xpf secure.tar.xz -C /boot
	cp -f secure.env "$OTA_BOOT_DIR/secure.env"
	echo "prepared" > "$OTA_BOOT_DIR/secure-state"
	echo "$SECURE_BUNDLE_MODE" > "$OTA_BOOT_DIR/secure-mode"
	echo "$SECURE_BUNDLE_MACHINE" > "$OTA_BOOT_DIR/secure-machine"
	if [ -n "$SECURE_PUBLIC_KEY_HASH" ]; then
		echo "$SECURE_PUBLIC_KEY_HASH" > "$OTA_BOOT_DIR/secure-public-key-hash"
	fi
	echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ")" > "$OTA_BOOT_DIR/secure-prepared-at"
fi

if [ "$BOOT_CONFIG_REL" = "cmdline.txt" ]; then
	ota_info "Restoring known-good boot config..."
	mkdir -p "/boot/$(dirname "$BOOT_CONFIG_REL")"
	cp -a "$WORKDIR/$BOOT_CONFIG_REL.before" "/boot/$BOOT_CONFIG_REL"
else
	ota_info "Using boot config from OTA payload..."
fi

ota_info "Writing boot status flags..."
echo "$NEW_ROOT" > "$OTA_BOOT_DIR/pending-root"
echo "$PREV_ROOT" > "$OTA_BOOT_DIR/previous-root"
echo "$TARGET_LABEL" > "$OTA_BOOT_DIR/pending-label"
echo "$TARGET_PART" > "$OTA_BOOT_DIR/pending-partition"
echo "$(date -u +"%Y-%m-%dT%H:%M:%SZ")" > "$OTA_BOOT_DIR/staged-at"

touch "$OTA_BOOT_DIR/rollback.flag"

stage_recovery_image

sync

echo
ota_info "Final boot config (/boot/$BOOT_CONFIG_REL):"
cat "/boot/$BOOT_CONFIG_REL"

echo
ota_info "/boot/ota:"
ls -la "$OTA_BOOT_DIR"

echo
ota_info "/boot/recovery:"
ls -la "$RECOVERY_DIR"

if [ -n "$SECURE_PREP_MODE" ]; then
	echo
	ota_info "/boot/secure-boot:"
	ls -la /boot/secure-boot

	echo
	ota_info "Secure state:"
	ls -la "$OTA_BOOT_DIR"/secure* 2>/dev/null || true
fi

echo
ota_info "Target /etc/fstab:"
cat "$TARGET_MNT/etc/fstab"

echo
ota_info "Upgrade staged successfully."

if [ "$DO_REBOOT" = "--reboot" ]; then
	ota_warn "Rebooting..."
	reboot
else
	ota_warn "Not rebooting automatically."
	ota_warn "Reboot manually when ready."
fi
