#!/bin/bash
# Flash one planck node's NVMe stick with a fresh Raspberry Pi OS.
# Usage: flash-one.sh <nodename> <device>   e.g. flash-one.sh planck016 sda
# Run as root (via pkexec). Wipes the target device completely.
set -euo pipefail

NODE=$1
DEV=$2
IMG=/tmp/opencode/raspios-lite.img
PUBKEY=/home/adamcarr/.ssh/id_ed25519.pub
PASSFILE=/home/adamcarr/code/personal/pi-swarm/configs/node-password.txt
EEPROM_UPD=/tmp/opencode/netboot/pieeprom-revert.upd
EEPROM_SIG=/tmp/opencode/netboot/pieeprom-revert.sig

fail() { echo "ERROR: $*" >&2; exit 1; }

# --- safety checks --------------------------------------------------------
[[ "$NODE" =~ ^planck0(0[1-9]|1[0-9]|20)$ ]] || fail "bad node name: $NODE"
[[ -b /dev/$DEV ]] || fail "/dev/$DEV is not a block device"
[[ "$DEV" == nvme* ]] && fail "refusing NVMe device name (looks like internal disk)"
grep -q "^/dev/$DEV" /proc/mounts && fail "device has mounted partitions, unplug/replug first"

SIZE_BYTES=$(cat /sys/block/$DEV/size 2>/dev/null) || fail "cannot read /sys/block/$DEV/size"
SIZE_GB=$(( SIZE_BYTES * 512 / 1024 / 1024 / 1024 ))
(( SIZE_GB >= 200 && SIZE_GB <= 300 )) || fail "device is ${SIZE_GB}GB, expected ~240GB - wrong stick or wrong port?"

MODEL=$(lsblk -dn -o MODEL /dev/$DEV | head -1)
SERIAL=$(lsblk -dn -o SERIAL /dev/$DEV | head -1)
echo "Target: /dev/$DEV  ${SIZE_GB}GB  model=$MODEL serial=$SERIAL"
echo "Node:   $NODE"
echo "WIPING THIS DEVICE IN 3 SECONDS (ctrl-c to abort)"
sleep 3

# --- write image -----------------------------------------------------------
wipefs -a /dev/$DEV
dd if=$IMG of=/dev/$DEV bs=4M conv=fsync status=progress
partprobe /dev/$DEV
sleep 1

# --- grow root partition to fill the disk ---------------------------------
echo ",+" | sfdisk --no-reread -N 2 /dev/$DEV
partprobe /dev/$DEV
sleep 1
e2fsck -fp /dev/${DEV}2 >/dev/null
resize2fs /dev/${DEV}2 >/dev/null

# --- configure -------------------------------------------------------------
B=/mnt/planck-boot
R=/mnt/planck-root
mkdir -p $B $R
mount /dev/${DEV}1 $B
mount /dev/${DEV}2 $R

# enable ssh on first boot
touch $B/ssh

# planck020 only: its EEPROM is stuck in network-boot mode, these files
# tell the bootloader to flip it back to normal on first power-up
if [ "$NODE" = "planck020" ]; then
    cp $EEPROM_UPD $B/pieeprom.upd
    cp $EEPROM_SIG $B/pieeprom.sig
fi

# hostname
echo "$NODE" > $R/etc/hostname
sed -i "s/^127.0.1.1.*/127.0.1.1	$NODE/" $R/etc/hosts

# timezone (match old cluster)
rm -f $R/etc/localtime
ln -s /usr/share/zoneinfo/America/Los_Angeles $R/etc/localtime
echo "America/Los_Angeles" > $R/etc/timezone

# user adamcarr (uid 1000), password from passfile, groups like the old nodes
HASH=$(openssl passwd -6 "$(cat $PASSFILE)")
echo "adamcarr:x:1000:1000:Adam Carr:/home/adamcarr:/bin/bash" >> $R/etc/passwd
echo "adamcarr:$HASH:20000:0:99999:7:::" >> $R/etc/shadow
echo "adamcarr:x:1000:" >> $R/etc/group
awk -F: -v OFS=: '
    BEGIN{split("sudo adm dialout cdrom audio video plugdev games users input render netdev",t," "); for(i in t) want[t[i]]=1}
    want[$1] { if ($4=="") $4="adamcarr"; else $4=$4",adamcarr" }
    {print}' $R/etc/group > $R/etc/group.new && mv $R/etc/group.new $R/etc/group
mkdir -p $R/home/adamcarr/.ssh
cp $PUBKEY $R/home/adamcarr/.ssh/authorized_keys
chown -R 1000:1000 $R/home/adamcarr
chmod 700 $R/home/adamcarr/.ssh
chmod 600 $R/home/adamcarr/.ssh/authorized_keys
echo "adamcarr ALL=(ALL) NOPASSWD: ALL" > $R/etc/sudoers.d/010-adamcarr
chmod 440 $R/etc/sudoers.d/010-adamcarr

sync
umount $B $R
sync
eject /dev/$DEV 2>/dev/null || true

echo ""
echo "DONE: $NODE flashed. Safe to unplug now."