#!/usr/bin/env bash
# flash-image.sh 1.0.0 - AppGoBox - MIT License
# Writes a raw disk image (.img, .img.xz, .img.gz, .img.zst, .zip with one .img) to a removable
# microSD card or USB drive on Linux or macOS, after optional SHA-256 verification.
# It ERASES the target device. It refuses vendor flashing packages (Rockchip update.img,
# Allwinner PhoenixCard, Amlogic USB burning images), mounted or system disks, and asks you to
# type the device path to confirm. Nothing is downloaded; nothing other than the target is written.
#
# Usage: sudo ./flash-image.sh [--dry-run] [--allow-fixed] IMAGE DEVICE [SHA256]
#   Linux DEVICE: /dev/sdX or /dev/mmcblkN      macOS DEVICE: /dev/diskN
set -euo pipefail

dry_run=0 allow_fixed=0
while [ $# -gt 0 ]; do
  case "$1" in
    --dry-run) dry_run=1; shift ;;
    --allow-fixed) allow_fixed=1; shift ;;
    -h|--help) sed -n '2,10p' "$0"; exit 0 ;;
    --) shift; break ;;
    -*) echo "Unknown option: $1" >&2; exit 2 ;;
    *) break ;;
  esac
done
if [ $# -lt 2 ] || [ $# -gt 3 ]; then sed -n '9,10p' "$0" >&2; exit 2; fi
image=$1 device=$2 expected=${3:-}
die() { echo "Error: $*" >&2; exit 1; }

[ -f "$image" ] || die "image not found: $image"
os=$(uname -s)

# 1. Optional checksum of the downloaded file.
if [ -n "$expected" ]; then
  expected=$(printf '%s' "$expected" | tr 'A-F' 'a-f')
  [[ "$expected" =~ ^[0-9a-f]{64}$ ]] || die "SHA256 must be 64 hex characters"
  echo "Checking SHA-256 of $image ..."
  if command -v sha256sum >/dev/null 2>&1; then actual=$(sha256sum -- "$image" | awk '{print $1}')
  else actual=$(shasum -a 256 -- "$image" | awk '{print $1}'); fi
  [ "$actual" = "$expected" ] || die "SHA-256 mismatch (got $actual). Do not write this file."
  echo "SHA-256 matches."
fi

# 2. How to read the image as a raw byte stream.
lower=$(printf '%s' "$image" | tr '[:upper:]' '[:lower:]')
case "$lower" in
  *.7z|*.rar|*.00[0-9]|*.part[0-9]*) die "extract this archive first (7-Zip), then pass the .img file" ;;
  *.xz)  reader=(xz -dc -- "$image") ;;
  *.gz)  reader=(gzip -dc -- "$image") ;;
  *.zst) reader=(zstd -dc -- "$image") ;;
  *.zip)
    entries=$(unzip -Z1 -- "$image" | grep -Ei '\.img$' || true)
    [ "$(printf '%s\n' "$entries" | grep -c .)" -eq 1 ] || die "zip must contain exactly one .img; extract it and choose the right file"
    reader=(unzip -p -- "$image" "$entries") ;;
  *.img|*.raw|*.bin) reader=(cat -- "$image") ;;
  *) die "unsupported file type: $image" ;;
esac
command -v "${reader[0]}" >/dev/null 2>&1 || die "missing tool: ${reader[0]}"

# 3. Refuse vendor flashing packages; require a partition table (MBR or GPT).
head_hex=$("${reader[@]}" 2>/dev/null | head -c 4608 | od -An -v -tx1 | tr -d ' \n' || true)
magic() { [ "${head_hex:$(( $1 * 2 )):${#2}}" = "$2" ]; }
magic 0 524b4657 && die "this is a Rockchip update.img; use RKDevTool / upgrade_tool / the vendor tool instead"
magic 0 524b4146 && die "this is a Rockchip firmware package; use the vendor flashing tool instead"
magic 0 494d414745575459 && die "this is an Allwinner PhoenixCard/PhoenixSuit image; use PhoenixCard instead"
{ magic 8 5619b527 || magic 8 27b51956; } && die "this is an Amlogic USB burning image; use USB Burning Tool / burn-tool instead"
if ! magic 510 55aa && ! magic 512 4546492050415254 && ! magic 4096 4546492050415254; then
  die "no MBR/GPT partition table found; this does not look like a raw disk image"
fi

# 4. Validate the target device.
size_bytes=0
if [ "$os" = Linux ]; then
  [[ "$device" =~ ^/dev/(sd[a-z]+|mmcblk[0-9]+)$ ]] || die "use a whole-disk device such as /dev/sdb or /dev/mmcblk0"
  [ -b "$device" ] || die "not a block device: $device"
  name=${device#/dev/}
  # Any mounted filesystem or swap on the disk or its children (partitions, LVM, RAID) blocks writing;
  # this also covers the disk that holds the running system.
  if lsblk -nro MOUNTPOINT "$device" | grep -q .; then die "$device (or a partition/volume on it) is mounted or in use; unmount it first"; fi
  removable=$(cat "/sys/block/$name/removable" 2>/dev/null || echo 0)
  tran=$(lsblk -dno TRAN "$device" 2>/dev/null || true)
  if [ "$removable" != 1 ] && [ "$tran" != usb ] && [[ "$name" != mmcblk* ]] && [ $allow_fixed = 0 ]; then
    die "$device is not removable/USB; pass --allow-fixed only if you are sure"
  fi
  size_bytes=$(lsblk -bdno SIZE "$device")
  lsblk -o NAME,SIZE,TYPE,TRAN,MODEL "$device"
  target=$device
elif [ "$os" = Darwin ]; then
  [[ "$device" =~ ^/dev/disk[0-9]+$ ]] || die "use a whole disk such as /dev/disk4"
  info=$(diskutil info "$device") || die "unknown disk: $device"
  echo "$info" | grep -Eq 'Internal:[[:space:]]+Yes' && [ $allow_fixed = 0 ] && die "$device is an internal disk"
  bootdisk=$(diskutil info / | awk -F: '/Part of Whole/ {gsub(/ /,"",$2); print $2}')
  [ "/dev/$bootdisk" != "$device" ] || die "$device holds the running system"
  size_bytes=$(echo "$info" | awk -F'[()]' '/Disk Size/ {split($2,a," "); print a[1]; exit}')
  echo "$info" | grep -E 'Device / Media Name|Disk Size|Protocol|Removable Media'
  target=${device/\/dev\/disk/\/dev\/rdisk}
else
  die "unsupported system: $os (on Windows use Raspberry Pi Imager, balenaEtcher or Rufus)"
fi
[ "${size_bytes:-0}" -gt 0 ] || die "could not read the size of $device"

# 5. Explicit confirmation, then write.
echo
echo "ALL DATA ON $device ($(( size_bytes / 1000000000 )) GB) WILL BE ERASED."
read -r -p "Type the device path ($device) to continue: " answer
[ "$answer" = "$device" ] || die "confirmation did not match; nothing written"
[ $dry_run = 1 ] && { echo "Dry run: would write ${reader[*]} | dd of=$target"; exit 0; }
[ "$(id -u)" = 0 ] || die "writing needs root; run again with sudo"
[ "$os" = Darwin ] && diskutil unmountDisk "$device" >/dev/null
if [ "$os" = Linux ]; then
  "${reader[@]}" | dd of="$target" bs=4M conv=fsync status=progress
else
  "${reader[@]}" | dd of="$target" bs=4m
fi
sync
echo "Done. Remove the card safely, insert it into the board and power on."