#!/usr/bin/env bash
# Deliberately guarded writer for the completed RK1 eMMC image.

set -Eeuo pipefail

IMAGE=""
TARGET=""
CONFIRM=""

usage() {
  cat <<'EOF'
Usage: sudo ./scripts/flash-emmc.sh --image IMAGE.img.xz \
       --target /dev/mmcblkN --confirm /dev/mmcblkN

The target and confirmation must match exactly. Only whole mmcblk devices are
accepted, and a device with mounted children or active swap is rejected.
EOF
}

while (($#)); do
  case "$1" in
    --image) IMAGE=${2:?}; shift 2 ;;
    --target) TARGET=${2:?}; shift 2 ;;
    --confirm) CONFIRM=${2:?}; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;;
  esac
done

for command_name in blockdev cmp dd lsblk sha256sum xz; do
  command -v "$command_name" >/dev/null || {
    echo "Missing command: $command_name" >&2
    exit 1
  }
done
[[ ${EUID:-$(id -u)} -eq 0 ]] || { echo 'Run this command as root' >&2; exit 1; }
[[ -f "$IMAGE" && "$IMAGE" == *.img.xz ]] || { echo 'A .img.xz image is required' >&2; exit 2; }
TARGET=$(readlink -f -- "$TARGET")
[[ "$TARGET" =~ ^/dev/mmcblk[0-9]+$ ]] || {
  echo 'Target must be a whole /dev/mmcblkN device (never a partition)' >&2
  exit 2
}
[[ "$CONFIRM" == "$TARGET" ]] || { echo '--confirm must exactly repeat the resolved target' >&2; exit 2; }
[[ -b "$TARGET" && "$(lsblk -dnro TYPE "$TARGET")" == disk ]] || {
  echo "Not a whole block device: $TARGET" >&2
  exit 1
}

if lsblk -nrpo MOUNTPOINTS "$TARGET" | grep -q '[^[:space:]]'; then
  echo "Target or one of its partitions is mounted: $TARGET" >&2
  lsblk -o NAME,PATH,SIZE,MODEL,SERIAL,TYPE,MOUNTPOINTS "$TARGET" >&2
  exit 1
fi
if command -v swapon >/dev/null && swapon --noheadings --raw --show=NAME \
  | grep -Eq "^${TARGET}(p[0-9]+)?$"; then
  echo "Target contains active swap: $TARGET" >&2
  exit 1
fi

xz --test "$IMAGE"
if [[ -f "$IMAGE.sha256" ]]; then
  (cd "$(dirname -- "$IMAGE")" && sha256sum --check --status "$(basename -- "$IMAGE.sha256")") || {
    echo 'Image checksum failed' >&2
    exit 1
  }
fi
image_bytes=$(xz --robot --list "$IMAGE" | awk -F '\t' '$1 == "file" { print $5 }')
target_bytes=$(blockdev --getsize64 "$TARGET")
[[ "$image_bytes" =~ ^[0-9]+$ && "$target_bytes" -ge "$image_bytes" ]] || {
  echo "Target is too small ($target_bytes bytes; image is $image_bytes bytes)" >&2
  exit 1
}

echo 'About to overwrite this entire device:' >&2
lsblk -d -o NAME,PATH,SIZE,MODEL,SERIAL,TRAN,TYPE "$TARGET" >&2
echo "Writing verified image: $IMAGE" >&2
xz --decompress --stdout "$IMAGE" \
  | dd of="$TARGET" bs=16M iflag=fullblock oflag=direct conv=fsync status=progress
sync

echo 'Reading the written region back for byte verification...' >&2
cmp --silent --bytes="$image_bytes" <(xz --decompress --stdout "$IMAGE") "$TARGET" || {
  echo 'Post-write verification FAILED; do not boot this device' >&2
  exit 1
}
blockdev --rereadpt "$TARGET" 2>/dev/null || true
echo "Flash and read-back verification passed: $TARGET" >&2