#!/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}* >> "$LOG" 2>&1
}

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

format_partition() {
	if [ -n "$FIRMPART" ]; then
		print_message "Formating firmware partition to FAT32..."
		echo "mkfs.fat -F 32 $FIRMPART"
		echo 'y' | mkfs.fat -F 32 $FIRMPART >> "$LOG" 2>&1 ||
		  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" >> "$LOG" 2>&1 ||
		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
}

create_partition() {
	if [ -n "$FIRMPART" ]; then
		print_message "Creating firmware partition..."
		echo 'y' | parted -s -a optimal "$MEDIA" \
			mkpart primary fat32 2MiB ${START_SECTOR} >> "$LOG" 2>&1 ||
			fatal_error "failed to create firmware partition"
		sync
		print_done
	fi
	print_message "Creating root partition..."
	echo 'y' | parted -s -a optimal "$MEDIA" \
		mkpart primary ext4 "$START_SECTOR" "$END_SECTOR" >> "$LOG" 2>&1 ||
		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
}

mount_partition() {
	# Get UUID
	ROOTPART_UID="$(blkid $ROOTPART | sed -E 's;.*\sUUID="([^"]*)".*;\1;')"
	TMPROOT="$(mktemp -d --tmpdir 'rootpart.XXXXXXXX' 2>>"$LOG")"
	print_message "Mounting root partition to the temporary directory..."
	mount "$ROOTPART" "$TMPROOT" >> "$LOG" 2>&1 ||
		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="$(mktemp -d --tmpdir 'firmpart.XXXXXXXX' 2>>"$LOG")"
		mount "$FIRMPART" "$TMPFIRM" >> "$LOG" 2>&1 ||
			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
	dd if="$ROOTFS" status=progress | $TAR - -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 >> "$LOG" 2>&1 &&
		  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
	[ -n "${RPI4_UBOOT-}" ] &&
		sed -i "s/LABEL=ROOT/UUID=$ROOTPART_UID/" "$RPI4_UBOOT"
}

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