#!/bin/sh
set -eu

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
}

DO_REBOOT=""
LOCAL_DIR=""

while [ "$#" -gt 0 ]; do
	case "$1" in
		--reboot)
			DO_REBOOT="--reboot"
			;;
		--local-dir)
			shift
			[ "$#" -gt 0 ] || {
				echo "Usage: $0 [--local-dir DIR] [--reboot]"
				exit 1
			}
			LOCAL_DIR="$1"
			;;
		*)
			echo "Usage: $0 [--local-dir DIR] [--reboot]"
			exit 1
			;;
	esac
	shift
done

[ "$(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}"

preferred_mmc_disks() {
	case "$MACHINE" in
		khadas-vim1s)
			echo "/dev/mmcblk2 /dev/mmcblk1 /dev/mmcblk0"
			;;
		*)
			echo "/dev/mmcblk1 /dev/mmcblk0 /dev/mmcblk2"
			;;
	esac
}

partition_disk() {
	case "$1" in
		/dev/mmcblk*p[0-9]*|/dev/nvme*n*p[0-9]*)
			echo "${1%p*}"
			;;
		/dev/*[0-9])
			echo "${1%[0-9]}"
			;;
		*)
			return 1
			;;
	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
}

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 "${disk}p1" ]; 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 "${disk}p1" ]; then
					echo "$disk"
					return 0
				fi
			fi
			;;
	esac

	for d in $(preferred_mmc_disks); do
		[ -b "${d}p1" ] || continue

		mkdir -p /tmp/ota-boot-test

		if mount -t vfat "${d}p1" /tmp/ota-boot-test 2>/dev/null; then
			if [ -f /tmp/ota-boot-test/extlinux/extlinux.conf ] || [ -d /tmp/ota-boot-test/ota ] || [ -d /tmp/ota-boot-test/recovery ]; then
				umount /tmp/ota-boot-test 2>/dev/null || true
				echo "$d"
				return 0
			fi

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

	return 1
}

DISK="$(find_boot_disk || true)"

[ -n "$DISK" ] || {
	echo "[OTA] ERROR: could not detect boot disk"
	exit 1
}

BOOT_PART="${DISK}p1"
ROOT_A_PART="${DISK}p2"
ROOT_B_PART="${DISK}p3"
HOME_PART="${DISK}p4"

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

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

	[ -d "$LOCAL_DIR" ] || {
		echo "[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" ] || {
			echo "[OTA] ERROR: missing local bundle file: $LOCAL_DIR/$f"
			exit 1
		}
		cp -f "$LOCAL_DIR/$f" "./$f"
	done
else
	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"
fi

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

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

BOOT_CONFIG_REL="$(find_boot_config_relpath || true)"

[ -n "$BOOT_CONFIG_REL" ] || {
	echo "[OTA] ERROR: could not detect boot config in /boot"
	exit 1
}

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

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] 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='./extlinux/extlinux.conf' \
	--exclude='extlinux/extlinux.conf' \
	--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 boot config..."
mkdir -p "/boot/$(dirname "$BOOT_CONFIG_REL")"
cp -a "$WORKDIR/$BOOT_CONFIG_REL.before" "/boot/$BOOT_CONFIG_REL"

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 boot config (/boot/$BOOT_CONFIG_REL):"
cat "/boot/$BOOT_CONFIG_REL"

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
