#!/bin/bash
# vim:sw=8:noet

# Except TEGRA
prepare_image() {
	print_message "Creating image $IMAGE_OUT..."
	dd if=/dev/zero of="$IMAGE_OUT" bs="$MB" count="$IMAGE_SIZE" ||
	  fatal_error "failed to create image $IMAGE_OUT"
	sync
	print_done
	losetup "$MEDIA" "$IMAGE_OUT"
}

prepare_media() {
	umount ${MEDIA}* 
}

clean_disklabel() {
	print_message "Clean old partitions table at $MEDIA..."
	# Clear the first block (two 512-byte sectors for
	# MBR and header, and 16KiB primary GPT Table):
	dd if=/dev/zero of=$MEDIA bs=512 count=128 &&
	# Clear secondary GPT Table
	dd if=/dev/zero of=$MEDIA bs=512 count=34 seek=$(($(blockdev --getsz $MEDIA) - 34)) ||
	fatal_error "failed clean old partitions table at $MEDIA!!!"
	print_done
}

clear_riscv64_bootloader_partition() {
	print_message "Find and clear old partitions with BBL and FSBL..."
	[ -n "$BBL" ] && [ -n "$FSBL" ] ||
		fatal_error "BBL and FSBL variables is not defined"
	let 'START_SECTOR_RAW = START_SECTOR * 2048' # 1 sector = 512 Byte
	local PART_NUMS=$(ls "$MEDIA"* |sed "s:$MEDIA::" |sed 's/^p//')
	[ -n "$PART_NUMS" ] || fatal_error "$MEDIA not contain partitions"
	local PART_NUM=
	for PART_NUM in $PART_NUMS; do
		FIND_PART_CODE=$(sgdisk -i "$PART_NUM" "$MEDIA" |
			grep 'Partition GUID code:' | cut -d ' ' -f 4)
		case "$FIND_PART_CODE" in
		"$BBL"|"$FSBL")
			sgdisk -d "$PART_NUM" "$MEDIA"
			continue
		;;
		esac
		FIND_PART_START=$(sgdisk -i "$PART_NUM" "$MEDIA" |
			grep 'First sector: ' | cut -d ' ' -f 3)
		[ "$FIND_PART_START" -ge "$START_SECTOR_RAW" ] ||
			fatal_error "Partition $PART_NUM is in the bootloader zone"
	done
}

create_disklabel() {
	clean_disklabel
	print_message "Creating disklabel $MEDIALABEL at $MEDIA..."
	parted -s "$MEDIA" mklabel "$MEDIATABLE" ||
		fatal_error "failed to create partitions table $MEDIATABLE"
	print_done
}

format_partition() {
	if [ -n "$FIRMPART" ]; then
		print_message "Formating firmware partition to FAT32..."
		echo 'y' | mkfs.fat -F 32 $FIRMPART ||
		  fatal_error "failed to format firmware partition"
		sync
		print_done
	fi
	if [ -n "$BOOTPART" ]; then
		print_message "Formating "$BOOTPART" boot partition to ext4..."
		echo 'y' | mkfs.ext4 "$BOOTPART" ||
			fatal_error "failed to format BOOT partition"
		sync
		print_done
	fi
	if [ -n "${LUKS-}" ]; then
		cryptsetup_luks_format
		cryptsetup_luks_open
	fi
	print_message "Formating "$ROOTPART" root partition to $ROOT_FSTYPE..."
	echo 'y' | mkfs.$ROOT_FSTYPE $ROOT_FS_OPT "$ROOTPART" ||
		fatal_error "failed to format ROOT partition"
	sync
	print_done
}

set_bootflag() {
	if [ "$MEDIATABLE" = msdos ] && [ -n "$FIRMPART_NUM" ]; then
		if [ -z "$EFI" ]; then
			if [ -z "$BOOTPART_NUM" ]; then
				print_message "Setting boot flag at root partition..."
				parted -s "$MEDIA" set "$ROOTPART_NUM" boot on ||
					fatal_error "failed to set bootflag on root partition"
				print_done
			else
				print_message "Setting boot flag at boot partition..."
				parted -s "$MEDIA" set "$BOOTPART_NUM" boot on ||
					fatal_error "failed to set bootflag on boot partition"
				print_done
			fi
		else
			print_message "Unsetting boot flag at root partition..."
			parted -s "$MEDIA" set "$ROOTPART_NUM" boot off ||
				fatal_error "failed to unset boot flag on root partition"
			print_done
		fi
	fi
	
	if [ "$MEDIATABLE" = gpt ] && [ -n "$FIRMPART_NUM" ]; then
		if [ -n "$EFI" ]; then
			print_message "Setting esp flag on EFI partition..."
			parted -s "$MEDIA" set "$FIRMPART_NUM" esp on ||
				fatal_error "failed to set esp flag on EFI partition"
			print_done
			print_message "Unsetting bootflag on root partition..."
			parted -s "$MEDIA" set "$ROOTPART_NUM" legacy_boot off ||
				fatal_error "failed to unset boot flag on root partition"
			print_done
		else
			print_message "Unsetting esp flag on EFI partition..."
			parted -s "$MEDIA" set "$FIRMPART_NUM" esp off ||
				fatal_error "failed to unset esp flag on EFI partition"
			print_done
			print_message "Setting bootflag on root partition..."
			parted -s "$MEDIA" set "$ROOTPART_NUM" legacy_boot on ||
				fatal_error "failed to set boot flag on root partition"
			print_done
		fi
	fi
}

create_partition() {
	[ -n "${START_SECTOR}" ] || START_SECTOR=2
	if [ -n "$FIRMPART" ]; then
		print_message "Creating firmware partition..."
		let NEXT_START="$START_SECTOR+$FIRMPART_SIZE"
		echo 'y' | parted -s -a optimal "$MEDIA" \
			mkpart primary fat32 ${START_SECTOR}MiB ${NEXT_START}MiB ||
			fatal_error "failed to create firmware partition"
		sync
		print_done
	fi
	if [ -n "$BOOTPART" ]; then
		print_message "Creating boot partition..."
		[ -n "$NEXT_START" ] && START_BOOT_SECTOR=$NEXT_START ||
			START_BOOT_SECTOR="$START_SECTOR"
		let NEXT_START="$START_BOOT_SECTOR+$BOOTPART_SIZE"
		echo 'y' | parted -s -a optimal "$MEDIA" \
			mkpart primary ext4 ${START_BOOT_SECTOR}MiB ${NEXT_START}MiB ||
			fatal_error "failed to create BOOT partition"
		sync
		print_done
	fi
	print_message "Creating root partition..."
	[ -n "$END_SECTOR" ] || END_SECTOR=100%
	[ -n "$NEXT_START" ] || NEXT_START="$START_SECTOR"
	echo 'y' | parted -s -a optimal "$MEDIA" \
		mkpart primary "$NEXT_START"MiB "$END_SECTOR" ||
		fatal_error "failed to create ROOT partition"
	sync
	print_done
	[ -z "${IMAGE_OUT-}" ] || finish_partitioning
}

cryptsetup_luks_format() {
	print_message "Formatting luks..."
	local PREV_LOG_MODE=$LOG_MODE
	LOG_MODE="con"
	local LUKS_DEFAULT_OPTIONS="--batch-mode --verify-passphrase"
	cryptsetup luksFormat $LUKS_DEFAULT_OPTIONS $LUKS_OPT $ROOTPART ||
		fatal_error "failed to encrypt root partition"
	LOG_MODE=PREV_LOG_MODE
	print_done
}

cryptsetup_luks_open() {
	print_message "Try to decrypt root partition..."
	local UNIQUE_MAPPING_NAME="luks_$(basename $ROOTPART)"
	cryptsetup open $ROOTPART $UNIQUE_MAPPING_NAME
	ROOTPART="/dev/mapper/$UNIQUE_MAPPING_NAME"
	LUKS_DEVICE_IS_OPEN="1"
}

cryptsetup_luks_close() {
	cryptsetup close $ROOTPART
	LUKS_DEVICE_IS_OPEN=
}

finish_partitioning() {
	# Create device maps from partition tables
	if [ -n "${IMAGE_OUT-}" ]; then
		kpartx -a -s "$MEDIA"
		ROOTPART="/dev/mapper/$(basename "$MEDIA")p$ROOTPART_NUM"
		[ -n "$FIRMPART_NUM" ] &&
			FIRMPART="/dev/mapper/$(basename "$MEDIA")p$FIRMPART_NUM"
		[ -n "$BOOTPART_NUM" ] &&
			BOOTPART="/dev/mapper/$(basename "$MEDIA")p$BOOTPART_NUM"
		[ -n "$BBLPART_NUM" ] &&
			BBLPART="/dev/mapper/$(basename "$MEDIA")p$BBLPART_NUM"
	fi
}

find_firmware_partition() {
	FIRMPART_NUM=$(parted -s "${MEDIA}" print |grep "fat" |
		head -n 1 |cut -d ' ' -f 2)
	[ "$FIRMPART_NUM" = "$ROOTPART_NUM" ] &&
		let ROOTPART_NUM=${FIRMPART_NUM}+1
}

change_partition() {
	[ -n "$FIRMPART_NUM" ] && FIRMPART="${MEDIA}${partsuffix}${FIRMPART_NUM}" ||
		FIRMPART=
	[ -n "$BOOTPART_NUM" ] && BOOTPART="${MEDIA}${partsuffix}${BOOTPART_NUM}" ||
		BOOTPART=
	[ -n "$ROOTPART_NUM" ] && ROOTPART="${MEDIA}${partsuffix}${ROOTPART_NUM}"
}

mount_partition() {
	[ -n "${ROOTFS-}" ] || [ -z "${LUKS-}" ] || cryptsetup_luks_open
	# Get UUID
	ROOTPART_UID="$(blkid $ROOTPART | sed -E 's;.*\sUUID="([^"]*)".*;\1;')"
	TMPROOT="$(mktemp -d --tmpdir 'rootpart.XXXXXXXX')"
	[ -d "${TMPROOT-}" ] ||
		fatal_error "Can't create a temporary directory for rootfs"
	print_message "Mounting root partition to the temporary directory..."
	mount "$ROOTPART" "$TMPROOT" ||
		fatal_error "failed to mount $ROOTPART to $TMPROOT"
	print_done
	if [ -n "$BOOTPART" ]; then
		BOOTPART_UID="$(blkid $BOOTPART |sed -E 's;.*\sUUID="([^"]*)".*;\1;')"
		print_message "Mounting boot partition to the temporary directory..."
		TMPBOOT="$TMPROOT/boot"
		mkdir -p "$TMPBOOT"
		[ -d "${TMPBOOT-}" ] ||
			fatal_error "Can't create a temporary directory for boot"
		mount "$BOOTPART" "$TMPBOOT" ||
			fatal_error "failed to mount $BOOTPART to $TMPBOOT"
		print_done
	fi
	if [ -n "$FIRMPART" ]; then
		FIRMPART_UID="$(blkid $FIRMPART |sed -E 's;.*\sUUID="([^"]*)".*;\1;')"
		print_message "Mounting firmware partition to the temporary directory..."
		TMPFIRM="$TMPROOT/boot/efi"
		mkdir -p "$TMPFIRM"
		[ -d "${TMPFIRM-}" ] ||
			fatal_error "Can't create a temporary directory for firmware"
		mount "$FIRMPART" "$TMPFIRM" ||
			fatal_error "failed to mount $FIRMPART to $TMPFIRM"
		print_done
	fi
}

unmount_partition() {
	if [ -n "$TMPFIRM" ] && [ -d "$TMPFIRM" ]; then
		umount "$TMPFIRM"
	fi

	if [ -n "$TMPBOOT" ] && [ -d "$TMPBOOT" ]; then
		umount "$TMPBOOT"
		rmdir "$TMPBOOT"
	fi

	if [ -n "$TMPROOT" ] && [ -d "$TMPROOT" ]; then
		umount "$TMPROOT"
		rmdir "$TMPROOT"
	fi

	if [ -n "$LUKS_DEVICE_IS_OPEN" ]; then
		cryptsetup_luks_close
	fi
}

write_rootfs() {
	if [ -n "${IMAGE_OUT-}" ]; then
		print_message "Writing $ROOTFS rootfs to $IMAGE_OUT..."
	else
		print_message "Writing $ROOTFS rootfs to $MEDIA..."
	fi
	echo
	$TAR "$ROOTFS" -C "$TMPROOT" ||
	  fatal_error "failed to write rootfs"
	sync
	print_done
}

setup_fstab() {
	if [ -n "$ROOTPART_UID" ]; then
		print_message "Updating fstab and extlinux.conf..."
		grep -e '[[:space:]]/[[:space:]]' $TMPROOT/etc/fstab &&
		  sed -i "s/LABEL=ROOT/UUID=$ROOTPART_UID/" "$TMPROOT"/etc/fstab ||
		  echo "UUID=$ROOTPART_UID	/	$ROOT_FSTYPE relatime	1 1" >> "$TMPROOT/etc/fstab"
		if [ -f "$TMPROOT/boot/extlinux/extlinux.conf" ]; then
			sed -i "s/LABEL=ROOT/UUID=$ROOTPART_UID/" "$TMPROOT/boot/extlinux/extlinux.conf"
		fi
		if [ -n "$BOOTPART_UID" ]; then
			mkdir -p "$TMPROOT/boot"
			echo "UUID=$BOOTPART_UID /boot ext4 nodev,nosuid,noexec,relatime 1 2" >> "$TMPROOT/etc/fstab"
		fi
		if [ -n "$FIRMPART_UID" ]; then
			mkdir -p "$TMPROOT/boot/efi"
			echo "UUID=$FIRMPART_UID /boot/efi vfat umask=0,quiet,showexec,iocharset=utf8,codepage=866 1 2" >> "$TMPROOT/etc/fstab"
		fi
		print_done
	fi
}

sync_partitions() {
	print_message "Informing kernel about partition table changes..."
	partprobe "$MEDIA" || fatal_error "failed to partprobe $MEDIA"
	sync
	sleep 3
	print_done
}

update_cmdline_txt() {
	[ -f "$TMPROOT/usr/share/u-boot/rpi_4/cmdline.txt" ] &&
		RPI4_UBOOT=$TMPROOT/usr/share/u-boot/rpi_4/cmdline.txt
	[ -f "$TMPROOT/usr/share/u-boot/rpi_4_32b/cmdline.txt" ] &&
		RPI4_UBOOT=$TMPROOT/usr/share/u-boot/rpi_4_32b/cmdline.txt
	[ -f "$TMPROOT/boot/efi/cmdline.txt" ] &&
		RPI4_UBOOT=$TMPROOT/boot/efi/cmdline.txt
	[ -n "${RPI4_UBOOT-}" ] &&
		sed -i "s/LABEL=ROOT/UUID=$ROOTPART_UID/" "$RPI4_UBOOT"
}

image_to_media() {
	print_message "Writing $IMAGE to $MEDIA..."
	"$CAT" "$IMAGE" |
		log_errtty dd of=$MEDIA bs=4M iflag=fullblock oflag=direct status=progress &&
			print_done || print_fail
	sync
	unset MEDIATABLE
	yes | parted "${MEDIA}" print | grep gpt && MEDIATABLE="gpt"
	yes | parted "${MEDIA}" print | grep msdos && MEDIATABLE="msdos"
	partprobe "$MEDIA"
	if [ "$MEDIATABLE" = "gpt" ]; then
		sgdisk -g "$MEDIA" 
	fi
	ROOTPART_NUM=$(ls ${MEDIA}${partsuffix}? |wc -l)
	if [ "$RESIZE" = 1 ]; then
		print_message "Resizing root partition ${MEDIA}${partsuffix}$ROOTPART_NUM..."
		parted -s "$MEDIA" resizepart "$ROOTPART_NUM" 100% &&
		  resize2fs -f "$MEDIA${partsuffix}$ROOTPART_NUM" ||
		  fatal_error "root partition ${MEDIA}${partsuffix}${ROOTPART_NUM} resize failed!!!"
		  e2fsck -f "$MEDIA${partsuffix}$ROOTPART_NUM" 
		print_done
	fi
	partprobe "$MEDIA"
}

set_path_to_uboot() {
	[ -n "$UBOOT" ] || UBOOT="${TMPROOT-}/usr/share/u-boot/$TARGET"
}

mount_chroot() {
	mount --bind /dev "${TMPROOT-}/dev"
	mount --bind /dev/pts "${TMPROOT-}/dev/pts"
	mount --bind /sys "${TMPROOT-}/sys"
	mount --bind /proc "${TMPROOT-}/proc"
	mount tmpfs -t tmpfs -o mode=755 "${TMPROOT-}/tmp"
	chroot "${TMPROOT-}" /bin/true ||
		fatal_error "test chroot to rootfs failed!!!"
}

unmount_chroot() {
	! mountpoint -q "${TMPROOT-}/tmp" || umount "${TMPROOT-}/tmp"
	! mountpoint -q "${TMPROOT-}/proc" || umount "${TMPROOT-}/proc"
	! mountpoint -q "${TMPROOT-}/sys" || umount "${TMPROOT-}/sys"
	! mountpoint -q "${TMPROOT-}/dev/pts" || umount "${TMPROOT-}/dev/pts"
	! mountpoint -q "${TMPROOT-}/dev" || umount "${TMPROOT-}/dev"
}

exec_chroot() {
	chroot "${TMPROOT-}" "$@"
}

rpm_install() {
	mkdir -p "${TMPROOT-}"/tmp/rpms &&
	cp "$1"/*.rpm "${TMPROOT-}"/tmp/rpms/ &&
	exec_chroot rpm -Uhv $(exec_chroot find /tmp/rpms -name '*.rpm')
}

add_dtb() {
	while read dtbdir; do
		if basename "$dtbdir" | grep -qe "^$DTB_KVER"; then
			print_message "Adding dtb to $(basename $dtbdir)..."
			mkdir -p "$dtbdir"/"$DTB_VENDOR"/
			cp "$DTB_FILE" "$dtbdir"/"$DTB_VENDOR"/
			! mountpoint -q "${TMPROOT-}/boot/efi" ||
				cp "$DTB_FILE" "${TMPROOT-}/boot/efi/"
			found="yes"
		fi
	done < <(find "${TMPROOT-}"/boot/devicetree -maxdepth 1 -mindepth 1 -type d)
	[ -z "$found" ] && fatal_error "Could not find any devicetree dirs for $DTB_KVER kernels"
	unset found
	print_done
}
