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

# This function REQUIRES a path to the device and
# prints the path to root partition, if there is one.
#
# The presence of /etc/fstab has been taken as the
# main characteristic of the root partition.
#
find_root_partition() {
	local TMPDEVROOT DEVNAME=$1
	TMPDEVROOT=$(mktemp -d --tmpdir devroot.XXXXXXXX)
	[ -d "$TMPDEVROOT" ] ||
		fatal_error "Cannot create a temporary directory"
	for d in "${DEVNAME}"?*; do
		local TDIR=$(findmnt --noheading --output TARGET "$d")
		if [ -z "$TDIR" ]; then
			TDIR="${TMPDEVROOT}/${d##*/}"
			mkdir "$TDIR"
			mount -o ro $d "$TDIR" ||
				continue
		fi
		if [ -f "$TDIR/etc/fstab" ]; then
			echo "$d"
			break
		fi
	done
	[ `ls "$TMPDEVROOT" | wc -l` -gt 0 ] &&
		umount "$TMPDEVROOT"/*
	rm -rf "$TMPDEVROOT"
}

# Takes a part number from the path to partition and prints it.
#
partnum() {
	local PART=$1
	echo ${PART##${MEDIA}${partsuffix}}
}

# Prints the path of device which specified by LABEL/UUID/path.
#
dev_by_spec() {
	blkid --output device -t $1
}

# This function sets partitions' numbers and paths
# that are using everywhere in alt-rootfs-installer.
#
# Call it when precise partitions number are required.
#
find_partitions() {
	ROOTPART=$(find_root_partition "$MEDIA")
	ROOTPART_NUM=$(partnum "$ROOTPART")
	local TDIR=$(findmnt --noheading --output TARGET "$ROOTPART")
	local MOUNTED=0
	if [ ! -d "$TDIR" ]; then
		TDIR=$(mktemp -d --tmpdir 'find_fstab.XXXXXXXX')
		mount -o ro "$ROOTPART" "$TDIR"
		MOUNTED=1
	fi
	local FSTAB="$TDIR"/etc/fstab

	local BOOTSPEC=$(grep '/boot\s' "$FSTAB" | awk '{print $1}')
	BOOTPART=$(dev_by_spec "$BOOTSPEC")
	BOOTPART_NUM=$(partnum $BOOTPART)

	local FIRMSPEC=$(grep '/boot/efi\s' "$FSTAB" | awk '{print $1}')
	FIRMPART=$(dev_by_spec "$FIRMSPEC")
	FIRMPART_NUM=$(partnum $FIRMPART)

	if [ $MOUNTED = 1 ]; then
		umount "$TDIR"
		rm -rf "$TDIR"
	fi
}

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
	find_partitions
	if [ "$RESIZE" = 1 ]; then
		print_message "Resizing root partition $ROOTPART..."
		parted -s "$MEDIA" resizepart "$ROOTPART_NUM" 100% &&
		  resize2fs -f $ROOTPART ||
		  fatal_error "root partition $ROOTPART resize failed!!!"
		  e2fsck -f $ROOTPART
		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
}

write_partition_blob() {
	local name="$1"
	local file="$2"
	local of="$3"
	local bs="${4-4M}"

	print_message "Writing $name for ${TARGET-}..."
	local blob
	blob="$(readlink -e "$file")"
	[ -f "${blob-}" ] || fatal_error "${blob-} was not found"
	[ -b "${of-}" ] || fatal_error "${of-} is not a block device or does not exist"
	dd if="$blob" of="$of" bs="$bs" status=none conv=nocreat || \
		fatal_error "failed to write $name"
	sync
	print_done
}