#!/bin/bash
# Writes the contest image to USB sticks, all at the same time.
# ERASES the listed sticks. Only USB disks are accepted.
#
#   sudo ./write_sticks.sh build/acpc2026.qcow2 /dev/sdb /dev/sdc ...
#
# The image can be the qcow2 file of build.sh (the ICPC image it is based on
# must stay where it was during the build) or a raw .img file.
set -uo pipefail

usage() {
    echo "Usage: sudo $0 IMAGE DEVICE [DEVICE...]" >&2
    echo >&2
    echo "USB disks currently attached:" >&2
    lsblk -d -o NAME,SIZE,MODEL,TRAN | awk 'NR == 1 || $NF == "usb"' >&2
    exit 2
}

[ $# -ge 2 ] || usage
IMG=$1
shift
[ -f "$IMG" ] || { echo "Image not found: $IMG" >&2; exit 1; }
[ "$(id -u)" -eq 0 ] || { echo "Run this script as root (sudo)." >&2; exit 1; }
command -v qemu-img >/dev/null || { echo "qemu-img is missing." >&2; exit 1; }

FORMAT=$(qemu-img info --output=json "$IMG" | python3 -I -c 'import json, sys; print(json.load(sys.stdin)["format"])')
SIZE=$(qemu-img info --output=json "$IMG" | python3 -I -c 'import json, sys; print(json.load(sys.stdin)["virtual-size"])')

echo "Image: $IMG ($FORMAT, $((SIZE / 1000000000)) GB)"
echo
for dev in "$@"; do
    [ -b "$dev" ] || { echo "Not a disk: $dev" >&2; exit 1; }
    [ "$(lsblk -dn -o TYPE "$dev")" = disk ] || { echo "Not a whole disk: $dev" >&2; exit 1; }
    [ "$(lsblk -dn -o TRAN "$dev")" = usb ] || { echo "Not a USB disk, refused: $dev" >&2; exit 1; }
    if lsblk -n -o MOUNTPOINTS "$dev" | grep -q '[^[:space:]]'; then
        echo "A partition of $dev is mounted, unmount it first." >&2
        exit 1
    fi
    devsize=$(blockdev --getsize64 "$dev")
    [ "$devsize" -ge "$SIZE" ] || { echo "$dev is too small ($((devsize / 1000000000)) GB)" >&2; exit 1; }
    echo "  $dev  $(lsblk -dn -o SIZE,MODEL "$dev")"
done
echo
echo "ALL DATA ON THESE STICKS WILL BE ERASED."
read -rp "Type ERASE to continue: " answer
[ "$answer" = ERASE ] || { echo "Cancelled."; exit 1; }

logdir=$(mktemp -d)
write_one() {
    local dev=$1 start=$(date +%s)
    if qemu-img convert -n -f "$FORMAT" -O raw -t none "$IMG" "$dev" > "$logdir/$(basename "$dev").log" 2>&1; then
        sync
        echo "  done    $dev ($(( $(date +%s) - start )) s)"
    else
        echo "  FAILED  $dev (see $logdir/$(basename "$dev").log)"
        return 1
    fi
}

echo "Writing..."
fail=0
pids=()
for dev in "$@"; do
    write_one "$dev" &
    pids+=($!)
done
for pid in "${pids[@]}"; do wait "$pid" || fail=1; done

if [ $fail -eq 0 ]; then
    echo "All sticks are written."
    rm -rf "$logdir"
else
    echo "At least one stick failed. Logs: $logdir"
fi
exit $fail