#!/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
}

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

	if [ -n "${BOOTFLAG-}" ]; then
		print_message "Setting bootflag at root partition..."
		parted -s "$MEDIA" set "$ROOTPART_NUM" "$BOOTFLAG" on ||
			fatal_error "failed to set ROOT on $BOOTFLAG"
		print_done
	fi
	
	if [ "$MEDIATABLE" = gpt -a -n "$EFI" ]; then
		print_message "Setting esp flag at EFI partition..."
		parted -s "$MEDIA" set 1 esp on ||
			fatal_error "failed to set esp flag on EFI partition"
	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=
	ROOTPART_NUM=1
	parted "${MEDIA}${partsuffix}1" print | grep "fat" &&
		FIRMPART_NUM=1
	[ -n "$FIRMPART_NUM" ] && ROOTPART_NUM=2
}

change_partition() {
	[ -n "$FIRMPART_NUM" ] && FIRMPART="${MEDIA}${partsuffix}${FIRMPART_NUM}"
	[ -n "$ROOTPART_NUM" ] && ROOTPART="${MEDIA}${partsuffix}${ROOTPART_NUM}"
	[ -n "$BBLPART_NUM" ] && BBLPART="${MEDIA}${partsuffix}${BBLPART_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
}

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 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"
}
