#!/bin/sh
# u-boot-licheerv-flash - (re)write the D1 boot chain to the raw boot area.
#
# Safe by default. Invoked from an apk trigger it does NOTHING unless you have
# opted in by creating /etc/u-boot-licheerv/autoflash.conf. Run by hand it asks
# before writing. Either way it only ever writes the disk that / is mounted on
# (unless you override DEVICE), refuses to clobber a partition, skips a no-op,
# backs up the old boot area, and verifies the write.
#
# Manual use:
#   u-boot-licheerv-flash            # prompt, auto-detect boot disk
#   u-boot-licheerv-flash -y         # no prompt
#   DEVICE=/dev/mmcblk0 u-boot-licheerv-flash   # force a target
set -u

BLOB=/usr/share/u-boot/lichee-rv-dock/u-boot-sunxi-with-spl.bin
CONF=/etc/u-boot-licheerv/autoflash.conf
BACKUP_DIR=/var/lib/u-boot-licheerv
SEEK_KIB=8                     # the SPL lives at 8 KiB on sunxi / D1

from_trigger=0
assume_yes=0
for a in "$@"; do
	case "$a" in
		--from-trigger) from_trigger=1 ;;
		-y|--yes)       assume_yes=1 ;;
		/dev/*)         DEVICE_ARG="$a" ;;
	esac
done

log()  { echo "u-boot-licheerv: $*"; }
warn() { echo "u-boot-licheerv: $*" >&2; }
# When fired from a trigger we must never fail the apk transaction: bail with 0.
give_up() { [ "$from_trigger" = 1 ] && exit 0 || exit 1; }

[ -f "$BLOB" ] || { warn "boot blob missing ($BLOB); nothing to do"; exit 0; }

# --- opt-in gate ---------------------------------------------------------------
DEVICE=auto
if [ -f "$CONF" ]; then
	# shellcheck disable=SC1090
	. "$CONF"
elif [ "$from_trigger" = 1 ]; then
	log "autoflash not enabled (no $CONF)."
	log "to update U-Boot now, run: u-boot-licheerv-flash"
	exit 0
fi
# An explicit device argument wins over the config/auto.
[ -n "${DEVICE_ARG:-}" ] && DEVICE="$DEVICE_ARG"

# --- resolve the target disk ---------------------------------------------------
resolve_disk() {  # partition path -> whole-disk path
	case "$1" in
		/dev/mmcblk*p[0-9]*) echo "${1%p[0-9]*}" ;;
		/dev/nvme*p[0-9]*)   echo "${1%p[0-9]*}" ;;
		/dev/sd[a-z][0-9]*)  echo "${1%%[0-9]*}" ;;
		*) echo "$1" ;;
	esac
}
if [ "${DEVICE:-auto}" = auto ]; then
	root_src=$(awk '$2=="/"{print $1; exit}' /proc/mounts 2>/dev/null || true)
	case "$root_src" in
		/dev/*) DEVICE=$(resolve_disk "$root_src") ;;
		*)
			warn "cannot derive the boot disk from root device '$root_src'."
			warn "set DEVICE=/dev/... in $CONF, or pass the disk as an argument."
			give_up ;;
	esac
fi

# --- refuse anything that isn't a real block device ----------------------------
if [ ! -b "$DEVICE" ]; then
	warn "target '$DEVICE' is not a block device; refusing."
	give_up
fi

blob_bytes=$(wc -c < "$BLOB")
blocks=$(( (blob_bytes + 1023) / 1024 ))
new_hash=$(sha256sum < "$BLOB" | cut -d' ' -f1)

# Hash exactly blob_bytes of the raw boot area (read whole KiB blocks, then trim
# to the blob length so we compare like-for-like; the write does not pad the
# final block, so a blocks*1024 readback would never match).
region_hash() {
	dd if="$DEVICE" bs=1024 skip="$SEEK_KIB" count="$blocks" status=none 2>/dev/null \
		| head -c "$blob_bytes" | sha256sum | cut -d' ' -f1
}

# --- refuse if partition 1 starts inside the region we'd overwrite -------------
# Read the first partition's start sector straight from the MBR (more reliable
# than sysfs, which doesn't always expose loop/partition geometry, and reflects
# the actual on-disk table). On a GPT disk the protective entry starts at sector
# 1, so this correctly refuses there too (writing the SPL at 8 KiB would trash
# GPT anyway; this board uses MBR by design).
mbr_first_start() {
	_min=""
	for _i in 0 1 2 3; do
		_off=$(( 446 + _i * 16 + 8 ))   # LBA-start field of each 16-byte entry
		_lba=$(dd if="$DEVICE" bs=1 skip="$_off" count=4 status=none 2>/dev/null \
			| od -An -tu1 \
			| awk '{print $1 + $2*256 + $3*65536 + $4*16777216}')
		[ -n "$_lba" ] && [ "$_lba" -gt 0 ] || continue
		[ -z "$_min" ] || [ "$_lba" -lt "$_min" ] && _min="$_lba"
	done
	[ -n "$_min" ] && echo "$_min"
}
first_start_sect=$(mbr_first_start)
if [ -n "$first_start_sect" ]; then
	first_start_bytes=$(( first_start_sect * 512 ))
	end_of_write=$(( SEEK_KIB * 1024 + blob_bytes ))
	if [ "$end_of_write" -ge "$first_start_bytes" ]; then
		warn "the blob ($blob_bytes B) written at ${SEEK_KIB}KiB would reach byte"
		warn "$end_of_write, but partition 1 starts at byte $first_start_bytes."
		warn "refusing to overwrite the partition. Move the partition or shrink U-Boot."
		give_up
	fi
fi

# --- skip if the on-disk boot chain is already identical -----------------------
cur_hash=$(region_hash)
if [ "$cur_hash" = "$new_hash" ]; then
	log "on-disk boot chain on $DEVICE is already current; nothing to do."
	exit 0
fi

# --- confirm when interactive --------------------------------------------------
if [ "$from_trigger" = 0 ] && [ "$assume_yes" = 0 ]; then
	printf "u-boot-licheerv: write U-Boot to %s at %s KiB? [y/N] " "$DEVICE" "$SEEK_KIB"
	read -r ans
	case "$ans" in [Yy]*) ;; *) log "aborted."; exit 1 ;; esac
fi

# --- back up the old region, write, verify ------------------------------------
mkdir -p "$BACKUP_DIR"
dd if="$DEVICE" bs=1024 skip="$SEEK_KIB" count="$blocks" \
	of="$BACKUP_DIR/boot-area.prev.img" status=none 2>/dev/null || true

log "writing $blob_bytes bytes to $DEVICE at ${SEEK_KIB}KiB"
if ! dd if="$BLOB" of="$DEVICE" bs=1024 seek="$SEEK_KIB" conv=notrunc,fsync status=none; then
	warn "the write to $DEVICE failed."
	give_up
fi

chk=$(region_hash)
if [ "$chk" = "$new_hash" ]; then
	log "OK: boot chain updated and verified on $DEVICE."
	exit 0
fi

# Verification failed: the write happened but didn't read back clean. From a
# trigger, don't fail the apk transaction (the kernel etc. installed fine); shout
# instead. Run by hand, return nonzero so callers notice.
warn "VERIFY FAILED after writing $DEVICE - the card may be failing."
warn "the previous boot area is saved at $BACKUP_DIR/boot-area.prev.img."
warn "if the board will not boot, recover over USB via FEL mode (xfel / sunxi-fel)."
give_up
