#!/bin/bash -ex
# usage:
# tar2fs chroot.tar image.raw [size_in_bytes [fstype]]

. shell-error
export LANG=C

if [ $# -lt 2 ]; then
	fatal "error: tar2fs needs at least two arguments"
fi

# this needs env_keep sudo setup to actually work
if [ -n "$GLOBAL_BUILDDIR" ]; then
	WORKDIR="$GLOBAL_BUILDDIR/vmroot"
else
	WORKDIR="$(mktemp --tmpdir -d vmroot-XXXXX)"
fi

[ -n "$WORKDIR" ] || fatal "couldn't come up with suitable WORKDIR"

[ -n "$GLOBAL_DEBUG" ] || message "WORKDIR: $WORKDIR"

MB=1048576		# a parted's "megabyte" in bytes is *broken*

SIZE_FACTOR=2		# multiply the sizes found by this value
BOOT_SIZE_FACTOR=2	# multiply /boot size by this value additionally
BOOTLOADERPARTSIZEM=0	# PReP partition size (ppc*)

CUR_BOUNDARY=0		# align first partition at 1MB for performance (+1)

BOOTFSTYPE=
BOOTPART=
EFIPARTFSTYPE=
EFIPART=

BOOTLOADER="$5"

if [ -n "$6" ]; then
	ARCH="$6"
else
	ARCH="$(arch)"
fi

BOOTTYPE="$8"

case "$ARCH" in
e2k*)
	BOOTFSTYPE="ext2"       # firmware knows it
	BLOCKDEV="/dev/sda"     # ...hopefully...
	BOOTPART="1"
	ROOTPART="2"
	;;
ppc*)
	BOOTFSTYPE="ext4"
	BLOCKDEV="/dev/sda"
	BOOTLOADERPART="1"
	BOOTLOADERPARTSIZEM="8"
	ROOTPART="2"
	;;
arm*|aarch64)
	ROOTPART="1"
	BLOCKDEV="/dev/sda"
	CUR_BOUNDARY=15 # offset 16 MiB for singleboard's
	;;
riscv64)
	ROOTPART="1"
	BLOCKDEV="/dev/sda"
	CUR_BOUNDARY=33 # offset 34 MiB for singleboard's
	;;
*)
	ROOTPART="1"
	BLOCKDEV="/dev/sda"
	;;
esac

if [ "$BOOTLOADER" == grub-efi -o "$BOOTTYPE" == EFI ]; then
	EFIPART="1"
	EFIPARTSIZEM="256"
	EFIPARTFSTYPE="fat"
	if [ "$ARCH" = x86_64 ]; then
		BIOSPART="2"
		ROOTPART="3"
	else
		ROOTPART="2"
	fi
fi

# tested to work: ext[234], jfs
# NB: xfs doesn't have a spare sector for the bootloader
ROOTFSTYPE="${4:-ext4}"

if [ -f "$ROOTFS/boot/extlinux/extlinux.conf" ] && [ "$ROOTFSTYPE" != ext4 ]; then
	if [ -n "$EFIPART" ]; then
		BOOTPART="2"
		ROOTPART="3"
	else
		BOOTPART="1"
		ROOTPART="2"
	fi
	BOOTFSTYPE="ext4"
fi

PARTTABLE="$7"
if [ -z "$PARTTABLE" ]; then
	if [ "$BOOTLOADER" == grub-efi ]; then
		PARTTABLE=gpt
	else
		PARTTABLE=msdos
	fi
fi

# figure out the part taken by /boot in the given tarball
boot_size() {
	if [ -n "$BOOTPART" ]; then
		tar tvf "$1" \
		| awk ' \
			BEGIN { sum=0 }
			/^-.* \.\/boot\// { sum=sum+$3 }
			END { print sum }'
	else
		echo "0"
	fi
}

# parted wrapper for convenience
parting() { parted "$LOOPDEV" --align optimal --script -- "$@"; }

# unfortunately parted is insane enough to lump alignment controls
# into unit controls so creating adjacent partitions sized in MiB
# is not as straightforward as it might be... thus "+1" padding;
# see also http://www.gnu.org/software/parted/manual/parted.html#unit
mkpart() {
	# a bit different start/end semantics to handle end of device too
	local start="$(($CUR_BOUNDARY + 1))"	# yes, we lose a megabyte
	if [ -n "$1" ]; then
		CUR_BOUNDARY="$(($start + $1))"
		local end="$CUR_BOUNDARY"MiB
	else
		local end="$OFFSET"MiB
	fi
	if [ -n "$2" ]; then
		CUR_FS="$2"
	else
		CUR_FS=ext2
	fi
	parting mkpart primary "$CUR_FS" "$start"MiB "$end"
}

create_inner_grub_cfg() {
	local boot_path boot_uuid cfg_path
	cfg_path=$1
	[ -n "$cfg_path" ] || return 1
	if [ -n "$BOOTUUID" ]; then
		boot_uuid="$BOOTUUID"
		boot_path=grub
	else
		boot_uuid="$ROOTUUID"
		boot_path=boot/grub
	fi
	[ -n "$boot_uuid" ] || return 1
	cat >"$cfg_path" <<-GRUB_EOF
	search.fs_uuid $boot_uuid root
	set prefix=(\$root)/$boot_path
	configfile \$prefix/grub.cfg
	GRUB_EOF
}

# a tarball containing chroot with a kernel
TAR="$1"
[ -s "$TAR" ] || fatal "source tarball doesn't really exist"

# a path to the image to be generated
IMG="$2"
[ -d "$(dirname "$IMG")" ] || fatal "target directory doesn't exist"

# 0 means auto; if a value is specified it holds (no /boot subtracted)
ROOTSIZE="$3"
[ -n "$ROOTSIZE" -a "$ROOTSIZE" != 0 ] || unset ROOTSIZE

BOOTSIZE="$9"
[ -n "$BOOTSIZE" -a "$BOOTSIZE" != 0 ] || unset BOOTSIZE

# image size in bytes
TARSIZE="$(stat -Lc %s "$TAR")"
# /boot size in that tarball
BOOTDEFSIZE="$(boot_size "$TAR")"
DEFSIZE="$(($SIZE_FACTOR * ($TARSIZE - ${BOOTSIZE:-$BOOTDEFSIZE})))"	# (exact sizes)
ROOTSIZE="$((${ROOTSIZE:-$DEFSIZE} + $MB - 1))"	# for ceil rounding to MB
# image and /boot sizes in megabytes
ROOTSIZEM="$(($ROOTSIZE / $MB))"
BOOTDEFSIZE="$(($SIZE_FACTOR * $BOOT_SIZE_FACTOR * $BOOTDEFSIZE))"
BOOTSIZE="$((${BOOTSIZE:-$BOOTDEFSIZE} + $MB - 1))"	# for ceil rounding to MB
BOOTSIZEM="$(($BOOTSIZE / $MB))"

# single root partition hardwired so far,
# add another image for home/data/swap if needed
ROOTDEV="$BLOCKDEV$ROOTPART"

# last preparations...
MKFS="mkfs.$ROOTFSTYPE ${BOOTFSTYPE:+mkfs.$BOOTFSTYPE} \
  ${EFIPARTFSTYPE:+mkfs.$EFIPARTFSTYPE}"
for i in losetup sfdisk parted kpartx $MKFS; do
	if ! type -t "$i" >&/dev/null; then
		fatal "$i required but not found in host system"
	fi
done

ROOTFS="$WORKDIR/chroot"

BOOTFS=
EFIPARTFS=
if [ -n "$BOOTPART" ]; then
	BOOTFS="$ROOTFS/boot"
fi
if [ -n "$EFIPART" ]; then
	EFIPARTFS="$ROOTFS/boot/efi"
fi

exit_handler() {
	rc=$?
	cd /
	if [ -n "$ROOTFS" ]; then
		umount ${EFIPARTFS:+"$EFIPARTFS"} ${BOOTFS:+"$BOOTFS"} \
		  "$ROOTFS"{/dev,/proc,/sys,}

		if [ -n "$LOOPDEV" ]; then
			kpartx -d -s "$LOOPDEV" || {
				sleep 10
				kpartx -d -s -v "$LOOPDEV"
			}
			losetup --detach "$LOOPDEV"
		fi
		rm -r -- "$ROOTFS"
		rmdir -- "$WORKDIR"
	fi
	exit $rc
}

# handle -e in shebang as well
trap exit_handler EXIT ERR

# prepare disk image and a filesystem inside it
rm -f -- "$IMG"
OFFSET="$(($CUR_BOUNDARY + $EFIPARTSIZEM + $BOOTLOADERPARTSIZEM + $BOOTSIZEM + ${BIOSPART:+1} + $ROOTSIZEM - 1))"
dd if=/dev/zero of="$IMG" conv=notrunc bs=$MB count=1 seek="$OFFSET"
losetup -f "$IMG"
LOOPDEV=$(losetup -j "$IMG" | cut -f 1 -d ':')

parting mklabel "$PARTTABLE"

if [ -n "$BOOTLOADERPART" ] && [ -n "$BOOTLOADERPARTSIZEM" ]; then
	case "$ARCH" in
		ppc*)
			parting mkpart primary ext2 $((CUR_BOUNDARY+1))MiB $((BOOTLOADERPARTSIZEM + 1))MiB
			CUR_BOUNDARY="$BOOTLOADERPARTSIZEM"
			parting set 1 prep on
			parting set 1 boot on
			;;
	esac
fi

if [ -n "$EFIPART" ]; then
	EFIDEV="$EFIDEV$EFIPART"
	if [ "$PARTTABLE" == gpt ]; then
		parting mkpart fat32 $((CUR_BOUNDARY+1))MiB $(($EFIPARTSIZEM + 1))MiB
	else
		parting mkpart primary fat32 $((CUR_BOUNDARY+1))MiB $(($EFIPARTSIZEM + 1))MiB
	fi
	CUR_BOUNDARY="$EFIPARTSIZEM"
	parting set 1 boot on
	if [ "$PARTTABLE" == gpt ]; then
		parting set 1 esp on
	fi
fi

if [ -n "$BIOSPART" ]; then
	parting mkpart bios $((CUR_BOUNDARY+1))MiB $(($CUR_BOUNDARY + 2))MiB
	CUR_BOUNDARY="$(($CUR_BOUNDARY + 1))"
	parting set "$BIOSPART" bios on
fi

if [ -n "$BOOTPART" ]; then
	BOOTDEV="$BLOCKDEV$BOOTPART"
	mkpart "$BOOTSIZEM"
fi

# not ROOTSIZEM but "the rest"; somewhat non-trivial arithmetics lurk in parted
mkpart

kpartx -a -s "$LOOPDEV"
LOOPROOT="/dev/mapper/$(basename "$LOOPDEV")p$ROOTPART"

mkfs."$ROOTFSTYPE" "$LOOPROOT"

if [ -n "$BOOTPART" ]; then
	LOOPBOOT="/dev/mapper/$(basename "$LOOPDEV")p$BOOTPART"
	mkfs."$BOOTFSTYPE" "$LOOPBOOT"
fi

if [ -n "$BOOTLOADERPART" ] && [ -n "$BOOTLOADERPARTSIZEM" ]; then
	LOOPBOOTLOADER="/dev/mapper/$(basename "$LOOPDEV")p$BOOTLOADERPART"
fi

if [ -n "$EFIPART" ]; then
	LOOPEFI="/dev/mapper/$(basename "$LOOPDEV")p$EFIPART"
	mkfs.fat -F32 "$LOOPEFI"
fi

ROOTUUID="$(blkid -s UUID -o value -c /dev/null "$LOOPROOT")"
if [ -n "$ROOTUUID" ]; then
       ROOTDEV="UUID=$ROOTUUID"
else
       ROOTDEV="$LOOPROOT"
fi

if [ -n "$BOOTPART" ]; then
	BOOTUUID="$(blkid -s UUID -o value -c /dev/null "$LOOPBOOT")"
	if [ -n "$BOOTUUID" ]; then
		BOOTDEV="UUID=$BOOTUUID"
	fi
fi

if [ -n "$EFIPART" ]; then
	EFIUUID="$(blkid -s UUID -o value -c /dev/null "$LOOPEFI")"
	if [ -n "$EFIUUID" ]; then
		EFIDEV="UUID=$EFIUUID"
	fi
fi

# mount and populate it
mkdir -pm755 "$ROOTFS"
mount "$LOOPROOT" "$ROOTFS"

if [ -n "$BOOTPART" ]; then
	mkdir -pm700 "$BOOTFS"
	mount "$LOOPBOOT" "$BOOTFS"
fi

if [ -n "$EFIPART" ]; then
	mkdir -pm751 "$EFIPARTFS"
	mount "$LOOPEFI" "$EFIPARTFS"
fi

tar -C "$ROOTFS" --numeric-owner -xf "$TAR"
for i in /dev /proc /sys; do mount --bind "$i" "$ROOTFS$i"; done

if grep -qe "[[:space:]]/[[:space:]]" "$ROOTFS/etc/fstab"; then \
	sed -i "s/LABEL=ROOT/$ROOTDEV/" "$ROOTFS/etc/fstab"
else
	echo "$ROOTDEV / $ROOTFSTYPE relatime 1 1" >> "$ROOTFS/etc/fstab"
fi

# target device at once
if [ -n "$BOOTPART" ]; then
	echo "$BOOTDEV /boot $BOOTFSTYPE defaults 1 2" >> "$ROOTFS/etc/fstab"
fi
if [ -n "$EFIPART" ]; then
	echo "$EFIDEV /boot/efi vfat umask=0,quiet,showexec,iocharset=utf8,codepage=866 1 2" >> "$ROOTFS/etc/fstab"
fi

# clean fstab
sed -i "/LABEL=ROOT/d" "$ROOTFS/etc/fstab"
# ...target device too
sed -i "s,$LOOPROOT,$ROOTDEV," "$ROOTFS/etc/fstab"

# Update cmdline.txt for Raspberry Pi
[ -f "$ROOTFS/boot/efi/cmdline.txt" ] &&
    sed -i "s/LABEL=ROOT/$ROOTDEV/" "$ROOTFS/boot/efi/cmdline.txt"

# Update extlinux.conf
if [ -f "$ROOTFS/boot/extlinux/extlinux.conf" ]; then
	sed -i "s/LABEL=ROOT/$ROOTDEV/g" "$ROOTFS/boot/extlinux/extlinux.conf"
	if [ "$PARTTABLE" == gpt ]; then
		if [ -n "$BOOTPART" ]; then
			parting set "$BOOTPART" legacy_boot on
		else
			parting set "$ROOTPART" legacy_boot on
		fi
	fi
	if [ "$PARTTABLE" == msdos ]; then
		if [ -n "$BOOTPART" ]; then
			parting set "$BOOTPART" boot on
		else
			parting set "$ROOTPART" boot on
		fi
	fi
fi

# e2k
if [ -f "$ROOTFS/boot/boot.conf" ]; then
	sed -i "s/LABEL=ROOT/$ROOTDEV/g" "$ROOTFS/boot/boot.conf"
fi

# Query ARCH in chroot and redefine arch-dependent variable
ARCH="$(chroot "$ROOTFS" rpm --eval '%_host_cpu')"

# Setup bootloader
case "$BOOTLOADER" in
grub-efi)
	echo 'GRUB_DISABLE_OS_PROBER=true' >> "$ROOTFS"/etc/sysconfig/grub2
	chroot "$ROOTFS" grub-mkconfig -o /boot/grub/grub.cfg
	mkdir -p "$ROOTFS"/boot/efi/EFI/BOOT
	case "$ARCH" in
		x86_64)
			chroot "$ROOTFS" grub-install --target=i386-efi --recheck \
			  --removable --uefi-secure-boot
			chroot "$ROOTFS" grub-install --target=x86_64-efi --recheck \
			  --removable --uefi-secure-boot
			sed -i 's/initrd16/initrd/g' "$ROOTFS/boot/grub/grub.cfg"
			sed -i 's/linux16/linux/g' "$ROOTFS/boot/grub/grub.cfg"
			[ -n "$BIOSPART" ] &&
				chroot "$ROOTFS" grub-install --target=i386-pc "$LOOPDEV"
		;;
		aarch64)
			cp "$ROOTFS"/usr/lib64/efi/grubaa64.efi \
			   "$ROOTFS"/boot/efi/EFI/BOOT/bootaa64.efi
		;;
		riscv64)
			cp "$ROOTFS"/usr/lib64/efi/grubriscv64.efi \
			   "$ROOTFS"/boot/efi/EFI/BOOT/bootriscv64.efi
		;;
		loongarch64)
			cp "$ROOTFS"/usr/lib64/efi/grubloongarch64.efi \
			   "$ROOTFS"/boot/efi/EFI/BOOT/bootloongarch64.efi
		;;
	esac
	sed -i '/GRUB_DISABLE_OS_PROBER=true/d' "$ROOTFS/etc/sysconfig/grub2"
	[ -s "$ROOTFS"/boot/efi/EFI/BOOT/grub.cfg ] ||
		create_inner_grub_cfg "$ROOTFS"/boot/efi/EFI/BOOT/grub.cfg
	;;
grub)
	echo 'GRUB_DISABLE_OS_PROBER=true' >> "$ROOTFS"/etc/sysconfig/grub2
	chroot "$ROOTFS" grub-mkconfig -o /boot/grub/grub.cfg
	case "$ARCH" in
		*86*)
			chroot "$ROOTFS" grub-install --target=i386-pc "$LOOPDEV"
			sed -i 's/initrdefi/initrd/g' "$ROOTFS/boot/grub/grub.cfg"
			sed -i 's/linuxefi/linux/g' "$ROOTFS/boot/grub/grub.cfg"
		;;
		ppc*)
			[ -z "$LOOPBOOTLOADER" ] ||
				chroot "$ROOTFS" grub-install --target=powerpc-ieee1275 \
				--no-nvram "$LOOPBOOTLOADER"
		;;
	esac
	sed -i '/GRUB_DISABLE_OS_PROBER=true/d' "$ROOTFS/etc/sysconfig/grub2"
	;;
esac

if [ -n "$SUDO_USER" ]; then
	chown "$SUDO_USER:$(id -g "$SUDO_USER")" "$IMG" ||:
fi
# maybe qemu interpreter was copied to chroot;
# this is no longer necessary, remove
rm -rf "$ROOTFS"/.host ||: