#!/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

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?"

P1_SIZE=$(lsblk -dn -o SIZE /dev/${DEV}1 2>/dev/null | tr -dc 0-9)
[[ "$P1_SIZE" == "512M" ]] && fail "boot partition is already 512M - this stick was already flashed. Plugged in the right one?"

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
for i in $(seq 1 10); do
    partprobe /dev/$DEV 2>/dev/null || true
    udevadm settle 2>/dev/null || true
    umount /dev/${DEV}1 /dev/${DEV}2 2>/dev/null || true
    SIZE2=$(blockdev --getsize64 /dev/${DEV}2 2>/dev/null || echo 0)
    [ "$SIZE2" -gt 200000000000 ] && break
    sleep 1
done
[ "$SIZE2" -gt 200000000000 ] || fail "kernel never saw the resized partition"
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
for i in 1 2 3; do
    umount /dev/${DEV}1 /dev/${DEV}2 2>/dev/null || true
    mount /dev/${DEV}1 $B 2>/dev/null && mount /dev/${DEV}2 $R 2>/dev/null && break
    sleep 2
done
mountpoint -q $B && mountpoint -q $R || fail "could not mount partitions (desktop keeps grabbing them)"

# enable ssh on first boot
touch $B/ssh

# 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

# remove image's locked default 'pi' user (imager normally renames it;
# leaving it in place would collide with adamcarr's uid and break sudo)
sed -i "/^pi:/d" $R/etc/passwd $R/etc/shadow
awk -F: -v OFS=: '
    {r=""; n=split($4,m,","); for(i=1;i<=n;i++){ if(m[i]!="pi" && m[i]!="") r=r (r==""?"":",") m[i] }; $4=r}
    {print}' $R/etc/group > $R/etc/group.new && mv $R/etc/group.new $R/etc/group
rm -rf $R/home/pi

# 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."