#!/bin/bash
# SPDX-License-Identifier: GPL-2.0+

set -e

kbuild_root=$(dirname $(readlink -f "$0"))

declare -A cross_compile
declare -A extra_kcflags
declare -A target_root

if [[ -r $kbuild_root/config.sh ]] ; then
	source $kbuild_root/config.sh || exit 1
fi

#
# Options parsing
#

dry_run=0
partitions=
silent=0
sudo=sudo

while [[ $# != 0 ]] ; do
	option=$1
	shift

	case $option in
		--dry-run)
			dry_run=1
			sudo=echo
			;;
		--partitions)
			partitions=($1)
			shift
			;;
		--silent)
			silent=1
			;;
		--*|-*)
			echo "Unrecognized option $option"
			exit 1
			;;
		*)
			platform=$option
			;;
	esac
done

#
# Platform
#

if [[ -z "$platform" ]] ; then
	for f in $kbuild_root/platforms/*.sh ; do
		pattern=$(cat $f | grep "dir_pattern=" | head -1)
		pattern=${pattern/#dir_pattern=/}
		if [[ -n "$pattern" ]] ; then
			echo $PWD | grep -q "$pattern" || continue
			platform=$(basename -s .sh $f)
			break
		fi
	done
fi

if [[ ! -r "$kbuild_root/platforms/$platform.sh" ]] ; then
	echo "Unknown platform '$platform'. Supported platforms are"
	for f in $kbuild_root/platforms/*.sh ; do
		echo "- $(basename -s .sh $f)"
	done
	exit 1
fi

source "${kbuild_root}/platforms/$platform.sh" || exit 1

if [[ "${#flash_parts[@]}" -lt 1 ]] ; then
	echo "Platform '$platform' doesn't declare flash_parts, can't flash"
	exit 1
fi

#
# Architecture
#

if [[ ! -r "$kbuild_root/arch/$arch.sh" ]] ; then
	echo "Unknown architecture '$arch'. Supported architectures are"
	for f in $kbuild_root/arch/*.sh ; do
		echo "- $(basename -s .sh $f)"
	done
	exit 1
fi

source "${kbuild_root}/arch/$arch.sh" || exit 1

# Override the output directory for ARM for historical reasons
case $arch in
	arm)
		output_dir=$PWD/output/arm32
		;;
	*)
		output_dir=$PWD/output/$arch
		;;
esac

target_dir=${target_dir:-${target_root[$arch]}}
tftp_dir=${tftp_dir:-$tftp_root/$platform}

#
# Determine the image source and destination file names
#

_image_type=${image/:*/}

if [[ -r "$kbuild_root/images/${_image_type}.sh" ]] ; then
	source "${kbuild_root}/images/${_image_type}.sh" || exit 1
else
	_image_file=${_image_type}
fi

if echo "${image}" | grep -q ':' ; then
	_image_name=${image/*:/}
else
	_image_name=${_image_file}
fi

_image_file="$output_dir/arch/$arch/boot/${_image_file}"

#
# Validate the partition names and paths
#

if [[ ! $partitions ]] ; then
	partitions=(${!flash_parts[@]})
fi

for part_name in ${partitions[@]} ; do
	case $part_name in
	dt|kernel|modules)
		;;
	*)
		echo "Unknown partition name '$part_name'"
		exit 1
		;;
	esac

	ifs="$IFS"; IFS=";" read -ra flash_part <<< "${flash_parts[$part_name]}"; IFS="$ifs"
	part_path=${flash_part[2]}

	case $part_path in
	RAW|/*)
		;;
	*)
		echo "Unknown partition path '$part_path'"
		exit 1
		;;
	esac
done

#
# Locate the flash device
#

probe_device() {
	local dev=$1

	# Skip devices that don't exist.
	if [[ ! -b /dev/$dev ]] ; then
		return
	fi

	# Add the 'p' suffix for MMC devices.
	local partdev=$dev
	echo $dev | grep -q mmcblk && partdev=${partdev}p

	#
	# Safety check: make sure the number and name of the partitions we will
	# flash match the expectations.
	#
	local part_count=$(ls -1 /dev/ | grep "${partdev}[0-9]" | wc -l)
	if [[ $part_count != $flash_parts_num ]] ; then
		echo "Device /dev/$dev has $part_count partitions, $flash_parts_num expected" >&2
		return
	fi

	local part_name
	for part_name in ${!flash_parts[@]} ; do
		local flash_part
		IFS=";" read -ra flash_part <<< "${flash_parts[$part_name]}"

		local part_idx=${flash_part[0]}
		local part_label=${flash_part[1]}
		local part_dev=/dev/${partdev}${part_idx}

		local label=$(lsblk -n -o PARTLABEL ${part_dev})
		if [[ -z "$label" ]] ; then
			label=$(lsblk -n -o LABEL ${part_dev})
		fi

		if [[ "$label" != ${part_label} ]] ; then
			echo "Partition ${part_dev} has label '$label', expected '${part_label}'" >&2
			return
		fi
	done

	echo /dev/$partdev
}

disks="mmcblk0 sda"
dev=

for disk in $disks ; do
	dev=$(probe_device "$disk")
	if [[ -n "$dev" ]] ; then
		break
	fi
done

if [[ -z "$dev" ]] ; then
	echo "No flash device found"
	exit 1
fi

echo "Found flash device ${dev}"

#
# Safety check: make sure nothing is mounted on $flash_mount_point, and no
# partition from the device is mounted.
#

if mount | grep -q " on $flash_mount_point " ; then
	echo "File system already mounted on $flash_mount_point."
	exit 1
fi

if mount | grep -q "$dev" ; then
	echo "Please unmount partitions $dev* before proceeding."
	exit 1
fi

#
# Safety check: make sure the kernel modules are installed.
#

version=$(cat $output_dir/include/generated/utsrelease.h | cut -d '"' -f 2)
mod_path="$target_dir/lib/modules/$version"

if [[ ! -d "$mod_path" ]] ; then
	echo "Modules not installed in $mod_path"
	exit 1
fi

#
# Flash
#

for part_name in ${partitions[@]} ; do
	IFS=";" read -ra flash_part <<< "${flash_parts[$part_name]}"
	part_idx=${flash_part[0]}

	if [[ -z "$flash_msg" ]] ; then
		flash_msg="Will flash $part_name to ${dev}${part_idx}"
	else
		flash_msg="${flash_msg}, $part_name to ${dev}${part_idx}"
	fi
done

if [[ ${dry_run} == 1 ]] ; then
	flash_msg="${flash_msg} (dry-run)"
fi

echo "${flash_msg}."
echo -n "Proceed? [y/N] "

read answer

if [[ x$answer != xy ]] ; then
	echo "Aborted"
	exit 1
fi

flash_raw() {
	local dst="$1"
	local src="$3"

	${sudo} dd if="$src" of="$dst" bs=64k
}

flash_fs() {
	local dev="$1"
	local dst="${flash_mount_point}$2"
	local src="$3"

	${sudo} mount "$dev" "$flash_mount_point"
	trap "${sudo} umount $flash_mount_point" EXIT

	if [[ -f "$src" ]] ; then
		${sudo} cp "$src" "$dst"
	else
		${sudo} rsync -av --delete "$src" "$dst"
	fi

	trap - EXIT
	${sudo} umount "$flash_mount_point"
}

for part_name in ${partitions[@]} ; do
	IFS=";" read -ra flash_part <<< "${flash_parts[$part_name]}"
	part_idx=${flash_part[0]}
	part_path=${flash_part[2]}

	echo "Flashing ${part_name} (${part_path})"

	case $part_path in
	RAW)
		flash_cmd="flash_raw"
		;;
	/*)
		flash_cmd="flash_fs"
		;;
	esac

	case $part_name in
	dt)
		# For DT, only flash the first entry. This can be improved later if
		# needed.
		_dtb=${dtbs[0]}
		_dtb_src="$output_dir/arch/${arch}/boot/dts/${_dtb/:*/}"
		_dtb_dst=${_dtb/*:/}
		_dtb_dst=${_dtb_dst/*\//}

		$flash_cmd "${dev}${part_idx}" "${part_path}/${_dtb_dst}" "${_dtb_src}"
		;;
	kernel)
		$flash_cmd "${dev}${part_idx}" "${part_path}/${_image_name}" "${_image_file}"
		;;
	modules)
		$flash_cmd "${dev}${part_idx}" "${part_path}/" "${mod_path}"
		;;
	esac
done

echo "Files copied, syncing"
${sudo} sync
echo "Done"

[[ $silent == 0 ]] && play -q -n synth 0.1 sin 880