#!/bin/bash

wic=$1
dev=$2

if [ $# -ne 2 ]; then
	echo "$0 <wic> </dev/xxxxx>"
	exit 1
fi

# -------------------------------------

BOARD_SUPPORTED="am572x stm32mp157c-d2"
board_to_block () {
	local board=$1
	case $board in
		am572x)
			echo "2 3"
			;;
		stm32mp157c-d2)
			echo "7 8"
			;;
		*)
			;;
	esac
}

# -------------------------------------
loopdev=''
esp_size=''
esp_size_512=''
rootfs_size=''
rootfs_size_512=''
suffix=''
# mounts wic and returns loop device
mount_wic () {
	local dev=$(sudo losetup --show -f -P "$wic")
	echo "$dev"
}

# -------------------------------------
mount_and_get_size() {
	loopdev=$(mount_wic)
	[ -z "$loopdev" ] && echo "[ERROR] failed to mount WIC image (\"$wic\")" && exit 1

	local esp="$loopdev"p1
	local rootfs="$loopdev"p2

	local esp_basename=$(basename $esp)
	esp_size_512=$(cat /sys/class/block/$esp_basename/size)
	esp_size=$(expr $esp_size_512 / 2048)

	local rootfs_basename=$(basename $rootfs)
	rootfs_size_512=$(cat /sys/class/block/$rootfs_basename/size)
	rootfs_size=$(expr $rootfs_size_512 / 2048)
}
# -------------------------------------
cleanup () {
	#echo "[DEBUG] sudo losetup -d $loopdev"
	sudo losetup -d "$loopdev"
}
# -------------------------------------
device_suffix() {
	local device=$1
	case $device in
	/dev/sd*)
		# no need suffix
		;;
	/dev/mmc*)
		# need suffix
		suffix='p'
		;;
	*)
		echo "[ERROR] This kind of device are not supported by the script (device = $device)"
		exit 1
		;;
	esac
}
# -------------------------------------
# 1: device name
# 2: size of the partition to check on MB
_verify_device () {
	local device=$1
	local size=$2
	local size_512=$3
	if [ ! -e $device ]; then
		echo "[ERROR]: The device \"$device\" doesn't exist"
		cleanup
		exit 1
	fi
	if $(grep -q $device /proc/mounts) ;then
		echo "[ERROR]: The device \"$device\" are mounted on you PC, please umount it"
		echo "         before to flash it"
		cleanup
		exit 1
	fi
	dev_basename=$(basename $device)
	dev_size_512=$(cat /sys/class/block/$dev_basename/size)
	dev_size=$(expr $dev_size_512 / 2048)
	if [ "$dev_size_512" -lt "$size_512" ];then
		echo "[ERROR] The parition size of parition \"$device\" are too small"
		echo "[ERROR] "$size"MB minimum are requested for this parition [only "$dev_size"MB are available]"
		cleanup
		exit 1
	fi
}
# -------------------------------------
verify_devices () {
	local board=$1
	local device=$2

	# boot device
	boot_index=$(board_to_block "$board" | awk -F ' ' '{ print $1 }')
	boot_device="$device$suffix$boot_index"
	_verify_device $boot_device $esp_size $esp_size_512

	# root device
	root_index=$(board_to_block "$board" | awk -F ' ' '{ print $2 }')
	root_device="$device$suffix$root_index"
	_verify_device $root_device $rootfs_size $rootfs_size_512
}
# -------------------------------------
write_files () {
	local device=$1

	local esp="$loopdev"p1
	local rootfs="$loopdev"p2

	local esp_part=$(board_to_block "$board" | awk -F ' ' '{ print $1 }')
	local rootfs_part=$(board_to_block "$board" | awk -F ' ' '{ print $2 }')

	esp_part="$dev""$suffix""$esp_part"
	rootfs_part="$dev""$suffix""$rootfs_part"

	sudo dd if="$esp" of="$esp_part" bs=8M conv=fdatasync status=progress
	[ $? -ne 0 ] && echo "[ERROR] Failed to write ESP partition" && exit 1
	sudo dd if="$rootfs" of="$rootfs_part" bs=8M conv=fdatasync status=progress
	[ $? -ne 0 ] && echo "[ERROR] Failed to write rootfs partition" && exit 1
}

# -------------------------------------
_choice_shell() {
	local choice_name=$1
	local choice_list=$2
	local default_choice=$3
	#format list to have display aligned on column with '-' separation between name and description
	local options=$(echo "${choice_list}" | tr ' ' '\n')
	#change separator from 'space' to 'end of line' for 'select' command
	old_IFS=$IFS
	IFS=$'\n'
	local i=1
	unset LAUNCH_MENU_CHOICES
	for opt in $options; do
		printf "%3.3s. %s\n" $i $opt
		LAUNCH_MENU_CHOICES=(${LAUNCH_MENU_CHOICES[@]} $opt)
		i=$(($i+1))
	done
	IFS=$old_IFS
	# Item selection from list
	local selection=""
	while [ -z "$selection" ]; do
		echo -n "Which one would you like? [${default_choice}] "
		read -r -t 60 answer
		# Check that user has answered before timeout, else break
		[ "$?" -gt "128" ] && break
		if [ -z "$answer" ] && [ -n "$default_choice" ] ; then
			selection=${default_choice}
			break
		fi
		if [[ $answer =~ ^[0-9]+$ ]]; then
			if [ $answer -gt 0 ] && [ $answer -le ${#LAUNCH_MENU_CHOICES[@]} ]; then
				selection=${LAUNCH_MENU_CHOICES[$(($answer-1))]}
				break
			fi
		fi
		echo "Invalid choice: $answer"
		echo "Please use numeric value between '1' and '$(echo "$options" | wc -l)'"
	done
	eval ${choice_name}=$(echo $selection | cut -d' ' -f1)
}
# -------------------------------------


trap cleanup SIGHUP SIGINT SIGTERM

# choice the board
_choice_shell board "$BOARD_SUPPORTED" "stm32mp157c-d2"
echo "BOARD selected : $board"

if [ ! -f "$wic" ];then
	echo "[ERROR] wic file doesn't exist, please indicate a correct path for wic file"
	exit 1
fi

# get device suffix
device_suffix $dev

mount_and_get_size

verify_devices "$board" "$dev"

write_files $dev

cleanup