#!/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
	print_message "Formating "$ROOTPART" root partition to ext4..."
	echo 'y' | mkfs.ext4 -m 5 -b 4096 -i 8192 -I 256 "$ROOTPART" ||
		fatal_error "failed to format ROOT partition"
	sync
	print_done
}

set_bootflag() {
	if [ "$MEDIATABLE" = msdos ] && [ -n "$FIRMPART_NUM" ]; then
		if [ -z "$EFI" ]; 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 "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 START_ROOT_SECTOR="$START_SECTOR+128"
		echo 'y' | parted -s -a optimal "$MEDIA" \
			mkpart primary fat32 ${START_SECTOR}MiB ${START_ROOT_SECTOR}MiB ||
			fatal_error "failed to create firmware partition"
		sync
		print_done
	fi
	print_message "Creating root partition..."
	[ -n "$END_SECTOR" ] || END_SECTOR=100%
	[ -n "$START_ROOT_SECTOR" ] || START_ROOT_SECTOR="$START_SECTOR"
	echo 'y' | parted -s -a optimal "$MEDIA" \
		mkpart primary ext4 "$START_ROOT_SECTOR"MiB "$END_SECTOR" ||
		fatal_error "failed to create ROOT partition"
	sync
	print_done
}

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 "$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 "$ROOTPART_NUM" ] && ROOTPART="${MEDIA}${partsuffix}${ROOTPART_NUM}"
}

mount_partition() {
	# 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 "$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 "$TMPROOT" ] && [ -d "$TMPROOT" ]; then
		umount "$TMPROOT"
		rmdir "$TMPROOT"
	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	/	ext4 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 "$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
	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"
}
