#!/bin/sh

verbose=
board=
device=
dryrun=
imagedir=
default_freq=528
freq=
skip_delay=

get_boot_blockdev() {
	# Find the blockdevice where /boot is mounted from
	mount | sed -n -E 's:.*(/dev/mmcblk[0-9])p[0-9] on /boot .*:\1:p'
}

get_defaults() {
	if [ -z "$board" -a -e /sys/firmware/devicetree/base/compatible ]; then
		case "$(cat /sys/firmware/devicetree/base/compatible 2>/dev/null)" in
		pine64,pinephone*) board=pinephone ;;
		pine64,pinetab*) board=pinetab ;;
		esac
	fi

	if [ -z "$device" ]; then
		case "$board" in
		pinephone|pinetab) device=$(get_boot_blockdev) ;;
		esac
	fi

	if [ -z "$imagedir" ]; then
		imagedir="$(realpath $(dirname $0))"
		[ -f "$imagedir/README.txt" ] || imagedir="/usr/share/u-boot"
	fi

	if [ -z "$freq" -a -e /proc/device-tree/memory/ram_freq ]; then
		proc_freq=$(cat /proc/device-tree/memory/ram_freq 2>/dev/null)
		if [ -n "$proc_freq" ]; then
			echo "Detected clock: $proc_freq MHz"
			default_freq="$proc_freq"
		else
			echo "WARNING: Failed to read clock from /proc/device-tree/memory/ram_freq"
			echo "Using default clock: $default_freq MHz"
		fi
	elif [ -z "$freq" ]; then
		echo Using default clock: $default_freq MHz
	else
		echo Changing clock: $freq MHz
	fi
}

die() {
	echo "ERROR: $@"
	exit 1
}

usage() {
	get_defaults

	cat <<EOF
usage: $0 [-n,--dry-run] [-i,--imagedir <imagedir>] [-b|--board <board-type>] [-d|--device <device>] [-r|--ram-freq <freq>] [-s|--skip-delay]

options:

 -b,--board <board>       Specify the board type: pinephone, pinetab
                          (current default: ${board:-none})

 -d,--device <device>     Specify the device where to install u-boot
                          (current default: ${device:-none})

 -i,--imagedir <imagedir> Specify u-boot image directory
                          (current default: ${imagedir:-none})

 -n,--dry-run             Print commands but don't execute them

 -r,--ram-freq            RAM clock speed, one of:
                          * 528 (default for new postmarketOS installations)
                          * 552 (default before v21.03 / edge 2021-04-05)
                          * 624

 -s,--skip-delay          Skip delay and flash the image immediately 

EOF
}

validate_checksum() {
	file="$1"
	file_sha512="$file.sha512"
	file_size=$(stat -c%s $file)
	device="$2"
	bs="$3"
	seek="$4"

	checksum0=$(cat $file_sha512 | awk {'print $1'})
	checksum1=$(dd if=$device bs=1 skip=$(($seek * $bs)) count=$file_size status=none | sha512sum | awk {'print $1'})
	if [ "$checksum0" != "$checksum1" ]; then
		echo "File: $checksum0"
		echo "Part: $checksum1"
		die "Checksum failed"
	fi
	echo "Successful U-Boot image checksum verification on $device :"
	echo -e "\t$checksum1"
}

while [ $# -gt 0 ]; do
	opt="$1"
	shift
	case "$opt" in
	-b|--board)
		case "$1" in
		pinephone) board="pinephone" ;;
		pinetab) board="pinetab" ;;
		pocketpc) board="pocketpc" ;;
		*) usage; exit 1;;
		esac
		shift
		;;
	-d|--device)
		device="$1"
		shift
		;;
	-i|--imagedir)
		imagedir="$1"
		shift
		;;
	-n|--dry-run)
		dryrun="echo"
		;;
	-r|--ram-freq)
		case "$1" in
		528) freq="528" ;;
		552) freq="552" ;;
		624) freq="624" ;;
		*) usage; exit 1;;
		esac
		shift
		;;
	-s|--skip-delay)
		skip_delay="yes"
		;;
        --)
                break
                ;;
        -*)
                usage
                exit 1
                ;;
        esac
done

get_defaults
if [ -z "$board" -o -z "$device" -o -z "$imagedir" -o  ! -e "$imagedir" ]; then
	usage
	exit 1
fi

if [ -z "$dryrun" ]; then
	if [ ! -z "$freq" ]; then
		echo "/!\ Disclaimer warning:"
		echo "Overclocking your device may render it unstable or prevent it to boot"
		echo "If you're unable to boot your device you'll have to flash an image to a mSD card"
		echo "then boot on it and update u-boot running this command:"
		echo "update-u-boot -d /dev/mmcblk2"
		echo 
		read -p "Continue? (y/n) [n]: " -n 1 -r
		echo 
		case "$REPLY" in
			y|Y) ;;
			*) exit 1 ;;
		esac
	fi
	if [ -z "$skip_delay" ]; then
		echo "Updating $board u-boot in $device in 3 seconds..."
		sleep 3
	else
		echo "Updating $board u-boot in $device"
	fi
fi

if [ -z "$freq" ]; then
	freq=$default_freq
fi

(
set -e
case "$board" in
pinephone|pinetab|pocketpc)
	[ -e "$imagedir/sourceparts-pocketpc/" ] || die "pocketpc images not installed, apk add u-boot-sourceparts-pocketpc"
	$dryrun dd if=$imagedir/sourceparts-pocketpc/u-boot-sunxi-with-spl-$freq.bin of=$device bs=1024 seek=8 oflag=direct status=none
	[ -z "$dryrun" ] && validate_checksum $imagedir/sourceparts-pocketpc/u-boot-sunxi-with-spl-$freq.bin $device 1024 8
	;;
esac
$dryrun sync
) || die "U-Boot installation in $device failed"

[ -z "$dryrun" ] && echo "Completed successfully."