#!/bin/sh
set -eu

DO_REBOOT="${1:-}"

OS_RELEASE="/etc/os-release"

[ -f "$OS_RELEASE" ] || {
	echo "[OTA] ERROR: $OS_RELEASE missing"
	exit 1
}

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

[ -n "$VERSION" ] || {
	echo "[OTA] ERROR: VERSION missing from $OS_RELEASE"
	exit 1
}

[ -n "$MACHINE" ] || {
	echo "[OTA] ERROR: MACHINE missing from $OS_RELEASE"
	exit 1
}

case "$DO_REBOOT" in
	""|"--reboot")
		;;
	*)
		echo "Usage: $0 [--reboot]"
		exit 1
		;;
esac

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

command -v curl >/dev/null || { echo "[OTA] ERROR: curl missing"; exit 1; }
command -v findmnt >/dev/null || { echo "[OTA] ERROR: findmnt missing"; exit 1; }
command -v mkfs.btrfs >/dev/null || { echo "[OTA] ERROR: mkfs.btrfs missing"; exit 1; }
command -v mountpoint >/dev/null || { echo "[OTA] ERROR: mountpoint missing"; exit 1; }
command -v sha256sum >/dev/null || { echo "[OTA] ERROR: sha256sum missing"; exit 1; }
command -v tar >/dev/null || { echo "[OTA] ERROR: tar missing"; exit 1; }

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

BOOT_PART="/dev/mmcblk0p1"
ROOT_A_PART="/dev/mmcblk0p2"
ROOT_B_PART="/dev/mmcblk0p3"
HOME_PART="/dev/mmcblk0p4"

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

stage_recovery_image() {
	echo "[OTA] 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=$VERSION
MACHINE=$MACHINE
STAGED_AT=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
ROOT=$NEW_ROOT
EOF

	cd "$WORKDIR"

	sync

	echo "[OTA] Recovery rootfs staged"
}

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

trap cleanup EXIT

echo "[OTA] Machine: $MACHINE"
echo "[OTA] Version: $VERSION"
echo "[OTA] Update URL: $BASE_URL"

mkdir -p "$WORKDIR" "$TARGET_MNT"
cd "$WORKDIR"

echo "[OTA] Cleaning old downloads..."
rm -f rootfs.tar.xz boot.tar.xz SHA256SUMS os-release

echo "[OTA] 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"

echo "[OTA] Verifying checksums..."
sha256sum -c SHA256SUMS

ACTIVE_SOURCE="$(findmnt -no SOURCE / || true)"

case "$ACTIVE_SOURCE" in
	"$ROOT_A_PART"|"/dev/root")
		if [ "$ACTIVE_SOURCE" = "/dev/root" ]; then
			ACTIVE_PART="$ROOT_A_PART"
		else
			ACTIVE_PART="$ACTIVE_SOURCE"
		fi
		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"
		;;
	*)
		echo "[OTA] ERROR: active root is not recognised"
		echo "[OTA] Active source: $ACTIVE_SOURCE"
		exit 1
		;;
esac

echo "[OTA] Active root: $ACTIVE_PART"
echo "[OTA] Target partition: $TARGET_PART"
echo "[OTA] Previous root: $PREV_ROOT"
echo "[OTA] New root: $NEW_ROOT"

if findmnt -rn "$TARGET_PART" >/dev/null 2>&1; then
	echo "[OTA] ERROR: $TARGET_PART is currently mounted"
	findmnt "$TARGET_PART" || true
	exit 1
fi

echo "[OTA] Ensuring /boot is mounted..."
mkdir -p /boot

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

[ -f /boot/cmdline.txt ] || {
	echo "[OTA] ERROR: /boot/cmdline.txt missing"
	exit 1
}

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

cp -a /boot/cmdline.txt "$WORKDIR/cmdline.before"

echo "[OTA] Formatting inactive rootfs partition..."
mkfs.btrfs -f -L "$TARGET_LABEL" "$TARGET_PART"

echo "[OTA] Mounting target rootfs..."
mount "$TARGET_PART" "$TARGET_MNT"

echo "[OTA] Extracting new rootfs..."
tar --numeric-owner -xpf rootfs.tar.xz -C "$TARGET_MNT"

echo "[OTA] Writing target fstab..."
cat > "$TARGET_MNT/etc/fstab" <<EOF
/dev/root      /             auto   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_PART     /boot         vfat   defaults                           0 0
$HOME_PART     /home         btrfs  defaults                           0 0
EOF

echo "[OTA] Forcing /home resize on next boot..."
rm -f "$TARGET_MNT/etc/home-resize.done"
rm -f "$TARGET_MNT/home/.home-resize.done" 2>/dev/null || true
echo "[OTA] 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

echo "[OTA] 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

echo "[OTA] Restoring known-good cmdline.txt..."
cp -a "$WORKDIR/cmdline.before" /boot/cmdline.txt

echo "[OTA] 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
echo "[OTA] Final cmdline.txt:"
cat /boot/cmdline.txt

echo
echo "[OTA] /boot/ota:"
ls -la "$OTA_BOOT_DIR"

echo
echo "[OTA] /boot/recovery:"
ls -la "$RECOVERY_DIR"

echo
echo "[OTA] Target /etc/fstab:"
cat "$TARGET_MNT/etc/fstab"

echo
echo "[OTA] Upgrade staged successfully."

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