#!/bin/sh

PATH=/sbin:/bin:/usr/sbin:/usr/bin

mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev 2>/dev/null || true

echo "[rollback] initramfs started"

DISK="/dev/mmcblk0"

find_boot_disk() {
	for d in /dev/mmcblk1 /dev/mmcblk0 /dev/mmcblk2; do
		[ -b "${d}p1" ] || continue

		mkdir -p /boot-test

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

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

	return 1
}

DISK="$(find_boot_disk)" || {
	echo "[rollback] ERROR: could not find boot disk"
	exec sh
}

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

BOOT_MNT="/boot"
OTA_DIR="$BOOT_MNT/ota"
REPORT_DIR="$OTA_DIR/report"

RECOVERY_DIR="$BOOT_MNT/recovery"
RECOVERY_TAR="$RECOVERY_DIR/rootfs.tar.xz"
RECOVERY_SUM="$RECOVERY_DIR/rootfs.tar.xz.sha256"

MAX_BOOT_ATTEMPTS=3

mkdir -p "$BOOT_MNT"
mkdir -p /newroot

i=0
while [ ! -b "$BOOT_PART" ] && [ "$i" -lt 100 ]; do
	sleep 0.1
	i=$((i + 1))
done

fsck.vfat -a "$BOOT_PART" 2>/dev/null || true

mount -t vfat "$BOOT_PART" "$BOOT_MNT" || {
	echo "[rollback] ERROR: cannot mount /boot"
	exec sh
}

mkdir -p "$OTA_DIR"
mkdir -p "$REPORT_DIR"

LOG_FILE="$OTA_DIR/initramfs.log"

log() {
	echo "[rollback] $*" | tee -a "$LOG_FILE"
	sync
}

log "initramfs started from active boot partition"
log "cmdline=$(cat /proc/cmdline 2>/dev/null || true)"

boot_root="$ROOT_A_PART"

ROLLBACK_FLAG="$OTA_DIR/rollback.flag"
PENDING_ROOT="$OTA_DIR/pending-root"
PREVIOUS_ROOT="$OTA_DIR/previous-root"
CURRENT_ROOT="$OTA_DIR/current-root"
ATTEMPTED_ROOT="$OTA_DIR/attempted-root"
BOOT_ATTEMPTS="$OTA_DIR/boot-attempts"
LAST_FAILED_ROOT="$OTA_DIR/last-failed-root"
LAST_BOOTED_ROOT="$OTA_DIR/last-booted-root"

log "checking partitions"
[ -b "$ROOT_A_PART" ] && log "rootA exists" || log "rootA missing"
[ -b "$ROOT_B_PART" ] && log "rootB exists" || log "rootB missing"
[ -b "$HOME_PART" ] && log "home exists" || log "home missing"
[ -f "$RECOVERY_TAR" ] && log "recovery rootfs exists" || log "recovery rootfs missing"
[ -f "$RECOVERY_SUM" ] && log "recovery checksum exists" || log "recovery checksum missing"
command -v sfdisk >/dev/null 2>&1 && log "sfdisk exists" || log "sfdisk missing"
command -v partprobe >/dev/null 2>&1 && log "partprobe exists" || log "partprobe missing"
command -v mkfs.btrfs >/dev/null 2>&1 && log "mkfs.btrfs exists" || log "mkfs.btrfs missing"

root_label_for_part() {
	case "$1" in
		"$ROOT_A_PART") echo "rootA" ;;
		"$ROOT_B_PART") echo "rootB" ;;
		*) echo "rootA" ;;
	esac
}

increment_boot_attempts() {
	attempts=0

	if [ -f "$BOOT_ATTEMPTS" ]; then
		attempts="$(cat "$BOOT_ATTEMPTS" 2>/dev/null || echo 0)"
	fi

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

	attempts=$((attempts + 1))
	echo "$attempts" > "$BOOT_ATTEMPTS"
	echo "$attempts"
}

repair_root_from_recovery() {
	target_part="$1"
	target_label="$(root_label_for_part "$target_part")"

	log "attempting auto-repair of $target_part as $target_label"

	echo "$target_part" > "$REPORT_DIR/autorepair-root"
	echo "$target_label" > "$REPORT_DIR/autorepair-label"
	date -u +"%Y-%m-%dT%H:%M:%SZ" > "$REPORT_DIR/autorepair-started-at"

	if [ ! -f "$RECOVERY_TAR" ]; then
		log "ERROR recovery rootfs missing"
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "recovery_rootfs_missing" > "$REPORT_DIR/autorepair-reason"
		return 1
	fi

	if [ ! -f "$RECOVERY_SUM" ]; then
		log "ERROR recovery checksum missing"
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "recovery_checksum_missing" > "$REPORT_DIR/autorepair-reason"
		return 1
	fi

	cd "$RECOVERY_DIR" || {
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "recovery_dir_missing" > "$REPORT_DIR/autorepair-reason"
		return 1
	}

	expected="$(awk '{print $1}' "$RECOVERY_SUM")"

	echo "$expected  $RECOVERY_TAR" | sha256sum -c - >> "$LOG_FILE" 2>&1 || {
		log "ERROR recovery checksum failed"
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "recovery_checksum_failed" > "$REPORT_DIR/autorepair-reason"
		return 1
	}

	umount /newroot 2>/dev/null || true

	log "formatting $target_part"
	if ! mkfs.btrfs -f -L "$target_label" "$target_part" >> "$LOG_FILE" 2>&1; then
		log "ERROR mkfs.btrfs failed"
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "mkfs_failed" > "$REPORT_DIR/autorepair-reason"
		return 1
	fi

	log "mounting repaired root"
	if ! mount -t btrfs -o rw "$target_part" /newroot; then
		log "ERROR repaired root mount failed"
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "mount_after_mkfs_failed" > "$REPORT_DIR/autorepair-reason"
		return 1
	fi

	log "extracting recovery rootfs"
	if ! tar -xJpf "$RECOVERY_TAR" -C /newroot >> "$LOG_FILE" 2>&1; then
		log "ERROR recovery extract failed"
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "extract_failed" > "$REPORT_DIR/autorepair-reason"
		umount /newroot 2>/dev/null || true
		return 1
	fi

	sync
	umount /newroot 2>/dev/null || true

	echo "success" > "$REPORT_DIR/autorepair-result"
	echo "recovered_from_fat_rootfs" > "$REPORT_DIR/autorepair-reason"
	date -u +"%Y-%m-%dT%H:%M:%SZ" > "$REPORT_DIR/autorepair-finished-at"

	log "auto-repair completed"
	return 0
}

repair_partitions_if_needed() {
	log "partition repair check started"

	if [ -b "$ROOT_A_PART" ] && [ -b "$ROOT_B_PART" ] && [ -b "$HOME_PART" ]; then
		log "expected partitions exist, no partition repair needed"
		return 0
	fi

	log "expected partitions missing, starting automatic partition repair"

	echo "started" > "$REPORT_DIR/partition-repair-result"
	date -u +"%Y-%m-%dT%H:%M:%SZ" > "$REPORT_DIR/partition-repair-started-at"

	[ -b "$DISK" ] || {
		log "ERROR disk missing: $DISK"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "disk_missing" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	[ -f "$RECOVERY_TAR" ] || {
		log "ERROR recovery rootfs missing"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "recovery_rootfs_missing" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	[ -f "$RECOVERY_SUM" ] || {
		log "ERROR recovery checksum missing"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "recovery_checksum_missing" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	cd "$RECOVERY_DIR" || {
		log "ERROR recovery directory missing"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "recovery_dir_missing" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	expected="$(awk '{print $1}' "$RECOVERY_SUM")"

	echo "$expected  $RECOVERY_TAR" | sha256sum -c - >> "$LOG_FILE" 2>&1 || {
		log "ERROR recovery checksum failed"
		echo "failed" > "$REPORT_DIR/autorepair-result"
		echo "recovery_checksum_failed" > "$REPORT_DIR/autorepair-reason"
		return 1
	}

	log "recreating missing p2/p3/p4 only"

	sfdisk --no-reread --append "$DISK" > "$REPORT_DIR/sfdisk.log" 2>&1 <<EOF || {
start=2052MiB, size=4096MiB, type=83
start=6148MiB, size=4096MiB, type=83
start=10244MiB, type=83
EOF
		log "ERROR sfdisk append failed"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "sfdisk_failed" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	partprobe "$DISK" >> "$LOG_FILE" 2>&1 || true
	sleep 2

	i=0
	while { [ ! -b "$ROOT_A_PART" ] || [ ! -b "$ROOT_B_PART" ] || [ ! -b "$HOME_PART" ]; } && [ "$i" -lt 100 ]; do
		sleep 0.1
		i=$((i + 1))
	done

	[ -b "$ROOT_A_PART" ] || {
		log "ERROR rootA missing after partition repair"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "rootA_missing_after_repair" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	[ -b "$ROOT_B_PART" ] || {
		log "ERROR rootB missing after partition repair"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "rootB_missing_after_repair" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	[ -b "$HOME_PART" ] || {
		log "ERROR home missing after partition repair"
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "home_missing_after_repair" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	log "formatting rootB and home"

	mkfs.btrfs -f -L rootB "$ROOT_B_PART" >> "$LOG_FILE" 2>&1 || {
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "rootB_mkfs_failed" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	mkfs.btrfs -f -L home "$HOME_PART" >> "$LOG_FILE" 2>&1 || {
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "home_mkfs_failed" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	log "restoring rootA from recovery"
	repair_root_from_recovery "$ROOT_A_PART" || {
		echo "failed" > "$REPORT_DIR/partition-repair-result"
		echo "rootA_restore_failed" > "$REPORT_DIR/partition-repair-reason"
		return 1
	}

	echo "success" > "$REPORT_DIR/partition-repair-result"
	date -u +"%Y-%m-%dT%H:%M:%SZ" > "$REPORT_DIR/partition-repair-finished-at"

	rm -f "$PENDING_ROOT"
	rm -f "$PREVIOUS_ROOT"
	rm -f "$ROLLBACK_FLAG"
	rm -f "$ATTEMPTED_ROOT"
	rm -f "$BOOT_ATTEMPTS"
	rm -f "$OTA_DIR/pending-label"
	rm -f "$OTA_DIR/pending-partition"
	rm -f "$OTA_DIR/staged-at"

	echo "$ROOT_A_PART" > "$LAST_BOOTED_ROOT"
	echo "$ROOT_A_PART" > "$OTA_DIR/partition-repaired-root"
	date -u +"%Y-%m-%dT%H:%M:%SZ" > "$OTA_DIR/partition-repair-complete"

	sync

	log "partition repair complete, rebooting"
	reboot -f
}

repair_partitions_if_needed || {
	log "ERROR partition repair failed"
	exec sh
}

if [ -f "$PENDING_ROOT" ] && [ ! -f "$ATTEMPTED_ROOT" ]; then
	boot_root="$(cat "$PENDING_ROOT")"
	log "first attempt booting pending root: $boot_root"
	echo "$boot_root" > "$ATTEMPTED_ROOT"

elif [ -f "$ROLLBACK_FLAG" ]; then
	log "rollback.flag exists"

	if [ -f "$PENDING_ROOT" ] && [ -f "$PREVIOUS_ROOT" ]; then
		pending="$(cat "$PENDING_ROOT")"
		previous="$(cat "$PREVIOUS_ROOT")"

		log "pending root failed: $pending"

		attempts="$(increment_boot_attempts)"
		log "boot attempts: $attempts"
		echo "$attempts" > "$REPORT_DIR/boot-attempts"

		echo "$pending" > "$REPORT_DIR/failed-root"
		echo "$previous" > "$REPORT_DIR/rolled-back-to"
		echo "boot_not_confirmed" > "$REPORT_DIR/reason"
		date -u +"%Y-%m-%dT%H:%M:%SZ" > "$REPORT_DIR/rollback-detected-at"

		echo "$pending" > "$LAST_FAILED_ROOT"

		if [ "$attempts" -ge "$MAX_BOOT_ATTEMPTS" ]; then
			log "max boot attempts reached, auto-repairing pending root"

			if repair_root_from_recovery "$pending"; then
				log "auto-repair succeeded, booting repaired root: $pending"

				echo "autorepair_after_repeated_failures" > "$REPORT_DIR/reason"
				echo "$pending" > "$REPORT_DIR/repaired-root"

				rm -f "$ROLLBACK_FLAG"
				rm -f "$ATTEMPTED_ROOT"
				rm -f "$BOOT_ATTEMPTS"

				boot_root="$pending"
			else
				log "auto-repair failed, rolling back to: $previous"

				echo "autorepair_failed_rollback" > "$REPORT_DIR/reason"

				rm -f "$PENDING_ROOT"
				rm -f "$PREVIOUS_ROOT"
				rm -f "$ROLLBACK_FLAG"
				rm -f "$ATTEMPTED_ROOT"
				rm -f "$BOOT_ATTEMPTS"
				rm -f "$OTA_DIR/pending-label"
				rm -f "$OTA_DIR/pending-partition"
				rm -f "$OTA_DIR/staged-at"

				boot_root="$previous"
			fi
		else
			log "rolling back to: $previous"

			rm -f "$PENDING_ROOT"
			rm -f "$PREVIOUS_ROOT"
			rm -f "$ROLLBACK_FLAG"
			rm -f "$ATTEMPTED_ROOT"
			rm -f "$OTA_DIR/pending-label"
			rm -f "$OTA_DIR/pending-partition"
			rm -f "$OTA_DIR/staged-at"

			boot_root="$previous"
		fi

	elif [ -f "$PREVIOUS_ROOT" ]; then
		previous="$(cat "$PREVIOUS_ROOT")"

		log "rollback.flag exists but pending-root missing"
		log "rolling back to: $previous"

		echo "unknown" > "$REPORT_DIR/failed-root"
		echo "$previous" > "$REPORT_DIR/rolled-back-to"
		echo "rollback_flag_without_pending_root" > "$REPORT_DIR/reason"
		date -u +"%Y-%m-%dT%H:%M:%SZ" > "$REPORT_DIR/rollback-detected-at"

		rm -f "$PREVIOUS_ROOT"
		rm -f "$ROLLBACK_FLAG"
		rm -f "$ATTEMPTED_ROOT"
		rm -f "$OTA_DIR/pending-label"
		rm -f "$OTA_DIR/pending-partition"
		rm -f "$OTA_DIR/staged-at"

		boot_root="$previous"
	else
		log "rollback.flag present but previous-root missing"

		echo "unknown" > "$REPORT_DIR/failed-root"
		echo "$ROOT_A_PART" > "$REPORT_DIR/rolled-back-to"
		echo "rollback_flag_without_previous_root" > "$REPORT_DIR/reason"
		date -u +"%Y-%m-%dT%H:%M:%SZ" > "$REPORT_DIR/rollback-detected-at"

		rm -f "$ROLLBACK_FLAG"
		rm -f "$ATTEMPTED_ROOT"

		boot_root="$ROOT_A_PART"
	fi

else
	if [ -f "$CURRENT_ROOT" ]; then
		boot_root="$(cat "$CURRENT_ROOT")"
		log "no pending-root, booting current-root: $boot_root"
	else
		log "no pending-root/current-root, defaulting to rootA"
		boot_root="$ROOT_A_PART"
	fi
fi

case "$boot_root" in
	"$ROOT_A_PART"|"$ROOT_B_PART")
		;;
	*)
		log "invalid root: $boot_root"
		boot_root="$ROOT_A_PART"
		;;
esac

echo "$boot_root" > "$LAST_BOOTED_ROOT"

sync

log "attempting mount of $boot_root"

if ! mount -t btrfs -o rw "$boot_root" /newroot; then
	log "mount failed for $boot_root"
	echo "$boot_root" > "$REPORT_DIR/mount-failed-root"
	echo "mount_failed" > "$REPORT_DIR/reason"

	if repair_root_from_recovery "$boot_root"; then
		log "mounting repaired root: $boot_root"

		mount -t btrfs -o rw "$boot_root" /newroot || {
			log "ERROR repaired root still cannot mount"
			echo "failed" > "$REPORT_DIR/autorepair-result"
			echo "repaired_root_mount_failed" > "$REPORT_DIR/autorepair-reason"
			exec sh
		}
	else
		log "ERROR auto-repair failed"
		exec sh
	fi
fi

log "mount successful"

ls -lah /newroot >> "$LOG_FILE" 2>&1 || true

if [ -x /newroot/sbin/init ]; then
	log "/sbin/init exists"
else
	log "/sbin/init missing"
	echo "$boot_root" > "$REPORT_DIR/invalid-root"
	echo "missing_sbin_init" > "$REPORT_DIR/reason"

	umount /newroot 2>/dev/null || true

	if repair_root_from_recovery "$boot_root"; then
		log "mounting repaired root after invalid-root repair: $boot_root"

		mount -t btrfs -o rw "$boot_root" /newroot || {
			log "ERROR repaired root still cannot mount"
			echo "failed" > "$REPORT_DIR/autorepair-result"
			echo "repaired_root_mount_failed" > "$REPORT_DIR/autorepair-reason"
			exec sh
		}
	else
		log "ERROR auto-repair failed"
		exec sh
	fi
fi

[ -d /newroot/etc ] && log "/etc exists" || log "/etc missing"
[ -f /newroot/etc/fstab ] && log "/etc/fstab exists" || log "/etc/fstab missing"

log "preparing switch_root"

mkdir -p /newroot/proc
mkdir -p /newroot/sys
mkdir -p /newroot/dev
mkdir -p /newroot/run
mkdir -p /newroot/boot

log "moving pseudo filesystems"

mount --move /proc /newroot/proc || log "WARN could not move /proc"
mount --move /sys /newroot/sys || log "WARN could not move /sys"
mount --move /dev /newroot/dev || log "WARN could not move /dev"

log "unmounting boot before switch_root"
umount "$BOOT_MNT" || log "WARN could not unmount /boot"

sync

echo "[rollback] switch_root starting"

exec switch_root /newroot /sbin/init
