#!/bin/busybox sh

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

NEW_ROOT="/new_root"
ROOT_DEV=""
ROOT_FSTYPE=""
FOUND_INIT=""

log() { echo "[initramfs] $*"; }
warn() { echo "[initramfs][WARN] $*"; }
err() { echo "[initramfs][ERROR] $*"; }

rescue_shell() {
	echo "------------------------------------------"
	echo "RootFS handoff failed"
	echo "------------------------------------------"
	echo "[cmdline]"
	cat /proc/cmdline 2>/dev/null || true
	echo "[mounts]"
	mount 2>/dev/null || true
	echo "[partitions]"
	cat /proc/partitions 2>/dev/null || true
	echo "[block devices]"
	ls /dev/vd* /dev/sd* /dev/mmcblk* 2>/dev/null || true
	exec /bin/sh
}

setup_console() {
	if (exec 0</dev/console) 2>/dev/null; then
		exec 0</dev/console
		exec 1>/dev/console
		exec 2>/dev/console
	fi
}

parse_cmdline() {
	for x in $(cat /proc/cmdline 2>/dev/null); do
		case "$x" in
		root=*) ROOT_DEV="${x#root=}" ;;
		rootfstype=*) ROOT_FSTYPE="${x#rootfstype=}" ;;
		esac
	done

	if [ "$ROOT_DEV" = "/dev/ram0" ]; then
		warn "Ignoring initramfs root device: $ROOT_DEV"
		ROOT_DEV=""
	fi

	if [ "$ROOT_FSTYPE" = "cpio" ]; then
		warn "Ignoring initramfs rootfstype: $ROOT_FSTYPE"
		ROOT_FSTYPE=""
	fi
}

load_modules() {
	KVER=$(uname -r)
	MOD_PATH=""

	for mod in $(find "/lib/modules/$KVER" -name zshm.ko 2>/dev/null); do
		MOD_PATH="$mod"
		break
	done

	if [ -n "$MOD_PATH" ] && [ -f "$MOD_PATH" ]; then
		log "Loading $MOD_PATH"
		insmod "$MOD_PATH" || warn "zshm.ko load failed"
	else
		warn "zshm.ko not found for kernel $KVER"
	fi

	mdev -s
}

is_rootfs_candidate() {
	for target in /sbin/init /etc/init /bin/init /init /bin/sh; do
		[ -x "$NEW_ROOT$target" ] && return 0
	done
	return 1
}

cleanup_new_root() {
	umount "$NEW_ROOT" 2>/dev/null || true
}

try_mount_type() {
	DEV="$1"
	FSTYPE="$2"

	if [ "$FSTYPE" = "auto" ]; then
		mount -t auto -o rw "$DEV" "$NEW_ROOT" 2>/dev/null
	else
		mount -t "$FSTYPE" -o rw "$DEV" "$NEW_ROOT" 2>/dev/null
	fi
}

try_mount_one() {
	DEV="$1"

	[ -b "$DEV" ] || return 1
	log "Trying rootfs $DEV"
	cleanup_new_root

	if [ -n "$ROOT_FSTYPE" ]; then
		if try_mount_type "$DEV" "$ROOT_FSTYPE"; then
			is_rootfs_candidate && ROOT_DEV="$DEV" && return 0
			cleanup_new_root
		fi
	fi

	for t in ext4 squashfs auto; do
		if try_mount_type "$DEV" "$t"; then
			if is_rootfs_candidate; then
				ROOT_DEV="$DEV"
				[ -z "$ROOT_FSTYPE" ] && ROOT_FSTYPE="$t"
				return 0
			fi
			cleanup_new_root
		fi
	done

	return 1
}

try_explicit_root() {
	[ -n "$ROOT_DEV" ] || return 1
	case "$ROOT_DEV" in
	/dev/*) try_mount_one "$ROOT_DEV" && return 0 ;;
	esac
	return 1
}

auto_scan_rootfs() {
	round=1

	while [ "$round" -le 20 ]; do
		log "Scanning rootfs round $round/20"
		mdev -s
		cat /proc/partitions 2>/dev/null || true

		for dev in /dev/vda1 /dev/vda2 /dev/vda3 /dev/vda \
			   /dev/sda1 /dev/sda2 /dev/sda3 /dev/sda \
			   /dev/mmcblk0p1 /dev/mmcblk0p2 /dev/mmcblk0p3 /dev/mmcblk0; do
			try_mount_one "$dev" && return 0
		done

		sleep 1
		round=$((round + 1))
	done

	return 1
}

find_init_in_root() {
	for target in /sbin/init /etc/init /bin/init /init /bin/sh; do
		if [ -x "$NEW_ROOT$target" ]; then
			FOUND_INIT="$target"
			return 0
		fi
	done
	return 1
}

copy_one() {
	src="$1"
	dst="$NEW_ROOT$src"

	[ -e "$src" ] || [ -L "$src" ] || return 0
	mkdir -p "${dst%/*}"
	cp -a "$src" "$dst" || warn "copy failed: $src"
}

copy_missing_one() {
	src="$1"
	dst="$NEW_ROOT$src"

	[ -e "$src" ] || [ -L "$src" ] || return 0
	[ -e "$dst" ] || [ -L "$dst" ] && return 0
	mkdir -p "${dst%/*}"
	cp -a "$src" "$dst" || warn "copy failed: $src"
}

copy_missing_tree() {
	src_dir="$1"

	[ -d "$src_dir" ] || return 0
	find "$src_dir" -type d -print 2>/dev/null | while read dir; do
		[ "$dir" = "$src_dir" ] && continue
		mkdir -p "$NEW_ROOT$dir"
	done
	find "$src_dir" -type f -print 2>/dev/null | while read file; do
		copy_missing_one "$file"
	done
	find "$src_dir" -type l -print 2>/dev/null | while read link; do
		copy_missing_one "$link"
	done
}

copy_test_tools_to_rootfs() {
	log "Copying initramfs test tools into $ROOT_DEV"

	for path in \
		/bin/busybox \
		/bin/context1 /bin/dhry2 /bin/fio /bin/fstime /bin/iperf3 \
		/bin/pipe /bin/spawn /bin/syscall /bin/whetstone-double \
		/bin/dmesg /bin/findmnt /bin/lsblk /bin/mount /bin/ping /bin/umount \
		/sbin/blkid /sbin/ip /sbin/ss /sbin/tc \
		/usr/bin/complete_benchmark.sh /usr/bin/cyclictest /usr/bin/fio \
		/usr/bin/iperf3 /usr/bin/zshm \
		/usr/sbin/ethtool /usr/sbin/i2cdetect /usr/sbin/i2cdump \
		/usr/sbin/i2cget /usr/sbin/i2cset /usr/sbin/i2ctransfer \
		/irq_test /virtio_blk_test.sh /virtio_net_test.sh /virtio_gpio_test.sh; do
		copy_one "$path"
	done

	copy_missing_tree /lib
	copy_missing_tree /usr/lib
	copy_missing_tree /lib/modules
	sync
}

prepare_switch_root() {
	mkdir -p "$NEW_ROOT/dev" "$NEW_ROOT/proc" "$NEW_ROOT/sys"

	mount --move /dev "$NEW_ROOT/dev" || return 1
	mount --move /proc "$NEW_ROOT/proc" || return 1
	mount --move /sys "$NEW_ROOT/sys" || return 1
	return 0
}

mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev >/dev/null 2>&1 || mount -t tmpfs tmpfs /dev
setup_console
echo /sbin/mdev > /proc/sys/kernel/hotplug 2>/dev/null || true
mdev -s
mkdir -p "$NEW_ROOT"

parse_cmdline
log "Kernel cmdline: $(cat /proc/cmdline 2>/dev/null)"
load_modules

if ! try_explicit_root; then
	warn "Falling back to auto rootfs scan"
	auto_scan_rootfs || rescue_shell
fi

copy_test_tools_to_rootfs

if ! find_init_in_root; then
	err "No init found in $NEW_ROOT"
	rescue_shell
fi

log "Using root device: $ROOT_DEV"
log "Using init: $FOUND_INIT"

if ! prepare_switch_root; then
	err "Failed to prepare switch_root"
	rescue_shell
fi

log "Switching root to $ROOT_DEV"
exec switch_root "$NEW_ROOT" "$FOUND_INIT"

err "switch_root failed"
rescue_shell
