#!/usr/bin/env bash
set -euo pipefail

# --- CONFIG (edit if needed) ---
# Leave NVME empty → script will auto-detect the SSK dock.
NVME="${NVME:-}"
FLAVOR="${FLAVOR:-desktop}"
# Persistent cache so the image survives reboots.
IMG_DIR="${IMG_DIR:-/var/cache/styx-rpi}"
IMG_FILE="${IMG_FILE:-ubuntu-24.04.3-preinstalled-${FLAVOR}-arm64+raspi.img}"
IMG_BOOT_MNT="${IMG_BOOT_MNT:-/mnt/img-boot}"
IMG_ROOT_MNT="${IMG_ROOT_MNT:-/mnt/img-root}"
TGT_ROOT="/mnt/target-root"
TGT_BOOT="/mnt/target-boot"

STYX_USER="styx"
STYX_HOSTNAME="titan-ag"
STYX_PASS="TempPass#123"                # will be forced to change on first login via cloud-init
SSH_PUBKEY="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOb8oMX6u0z3sH/p/WBGlvPXXdbGETCKzWYwR/dd6fZb titan-bastion"

# Video / input prefs
DSI_FLAGS="video=DSI-1:800x480@60D video=HDMI-A-1:off video=HDMI-A-2:off"

# --- Helpers ---
fatal(){ echo "ERROR: $*" >&2; exit 1; }
need(){ command -v "$1" >/dev/null || fatal "Missing tool: $1"; }

require_root(){ [[ $EUID -eq 0 ]] || exec sudo -E "$0" "$@"; }

part() {
  local n="$1"
  if [[ "$NVME" =~ [0-9]$ ]]; then
    echo "${NVME}p${n}"
  else
    echo "${NVME}${n}"
  fi
}

auto_detect_target_disk() {
  # If user already set NVME, validate and return
  if [[ -n "${NVME:-}" ]]; then
    [[ -b "$NVME" ]] || fatal "NVME='$NVME' is not a block device"
    return
  fi

  # Prefer stable by-id symlinks
  local byid
  byid=$(ls -1 /dev/disk/by-id/usb-SSK* 2>/dev/null | head -n1 || true)
  if [[ -n "$byid" ]]; then
    NVME=$(readlink -f "$byid")
  else
    # Heuristic via lsblk -S: look for USB with SSK/Ingram/Storage in vendor/model
    NVME=$(lsblk -S -p -o NAME,TRAN,VENDOR,MODEL | \
             awk '/ usb / && ($3 ~ /SSK|Ingram/i || $4 ~ /SSK|Storage/i){print $1; exit}')
  fi

  [[ -n "${NVME:-}" && -b "$NVME" ]] || fatal "Could not auto-detect SSK USB NVMe dock. Export NVME=/dev/sdX and re-run."
  echo "Auto-detected target disk: $NVME"
}

preflight_cleanup() {
  local img="$IMG_DIR/$IMG_FILE"

  # 1) Unmount image mountpoints and detach only loops for this IMG
  umount -lf "$IMG_BOOT_MNT" "$IMG_ROOT_MNT" 2>/dev/null || true
  # losetup -j exits non-zero if no association → tolerate it
  { losetup -j "$img" | cut -d: -f1 | xargs -r losetup -d; } 2>/dev/null || true

  # 2) Unmount our target mounts
  umount -lf "$TGT_ROOT/boot/firmware" "$TGT_BOOT" "$TGT_ROOT" 2>/dev/null || true

  # 3) Unmount the actual target partitions if mounted anywhere (tolerate 'not found')
  for p in "$(part 1)" "$(part 2)"; do
    # findmnt returns 1 when no match → capture and iterate if any
    while read -r mnt; do
      [ -n "$mnt" ] && umount -lf "$mnt" 2>/dev/null || true
    done < <(findmnt -rno TARGET -S "$p" 2>/dev/null || true)
  done

  # 4) Close dm-crypt mapping (if it exists)
  cryptsetup luksClose cryptroot 2>/dev/null || true
  dmsetup remove -f cryptroot 2>/dev/null || true

  # 5) Let udev settle
  command -v udevadm >/dev/null && udevadm settle || true
}

guard_target_device() {
  # Refuse to operate if NVME appears to be the current system disk
  local root_src root_disk
  root_src=$(findmnt -no SOURCE /)
  root_disk=$(lsblk -no pkname "$root_src" 2>/dev/null || true)
  if [[ -n "$root_disk" && "/dev/$root_disk" == "$NVME" ]]; then
    fatal "Refusing to operate on system disk ($NVME). Pick the external NVMe."
  fi
}

need_host_fido2() {
  if ! command -v fido2-token >/dev/null 2>&1; then
    echo "Host is missing fido2-token. On Arch: sudo pacman -S libfido2"
    echo "On Debian/Ubuntu host: sudo apt-get install fido2-tools"
    exit 1
  fi
}

ensure_image() {
  mkdir -p "$IMG_DIR"
  chmod 755 "$IMG_DIR"

  local BASE="https://cdimage.ubuntu.com/releases/noble/release"
  local XZ="ubuntu-24.04.3-preinstalled-${FLAVOR}-arm64+raspi.img.xz"

  # If the decompressed .img is missing, fetch/decompress into the cache.
  if [[ ! -f "$IMG_DIR/$IMG_FILE" ]]; then
    need curl; need unxz  # Arch: pacman -S curl xz   |  Ubuntu: apt-get install curl xz-utils
    if [[ ! -f "$IMG_DIR/$XZ" ]]; then
      echo "Fetching image…"
      curl -fL -o "$IMG_DIR/$XZ" "$BASE/$XZ"
    fi
    echo "Decompressing to $IMG_DIR/$IMG_FILE …"
    # Keep the .xz for future runs; stream-decompress to the .img
    if command -v unxz >/dev/null 2>&1; then
      unxz -c "$IMG_DIR/$XZ" > "$IMG_DIR/$IMG_FILE"
    else
      need xz
      xz -dc "$IMG_DIR/$XZ" > "$IMG_DIR/$IMG_FILE"
    fi
    sync
  else
    echo "Using cached image: $IMG_DIR/$IMG_FILE"
  fi
}

ensure_binfmt_aarch64(){
  # Register qemu-aarch64 for chrooted ARM64 apt runs
  if [[ ! -e /proc/sys/fs/binfmt_misc/qemu-aarch64 ]]; then
    need docker
    systemctl enable --now docker >/dev/null 2>&1 || true
    docker run --rm --privileged tonistiigi/binfmt --install arm64 >/dev/null
  fi
  if [[ ! -x /usr/local/bin/qemu-aarch64-static ]]; then
    docker rm -f qemu-static >/dev/null 2>&1 || true
    docker create --name qemu-static docker.io/multiarch/qemu-user-static:latest >/dev/null
    docker cp qemu-static:/usr/bin/qemu-aarch64-static /usr/local/bin/
    install -D -m755 /usr/local/bin/qemu-aarch64-static /usr/local/bin/qemu-aarch64-static
    docker rm qemu-static >/dev/null
  fi
}

open_image() {
  [[ -r "$IMG_DIR/$IMG_FILE" ]] || fatal "Image not found: $IMG_DIR/$IMG_FILE"
  mkdir -p "$IMG_BOOT_MNT" "$IMG_ROOT_MNT"

  # Pre-clean: detach any previous loop(s) for this image (tolerate absence)
  umount -lf "$IMG_BOOT_MNT" 2>/dev/null || true
  umount -lf "$IMG_ROOT_MNT" 2>/dev/null || true
  # If no loop is attached, losetup -j returns non-zero → swallow it
  mapfile -t OLD < <({ losetup -j "$IMG_DIR/$IMG_FILE" | cut -d: -f1; } 2>/dev/null || true)
  for L in "${OLD[@]:-}"; do losetup -d "$L" 2>/dev/null || true; done
  command -v udevadm >/dev/null && udevadm settle || true

  # Attach with partition scan; wait for partition nodes to exist
  LOOP=$(losetup --find --show --partscan "$IMG_DIR/$IMG_FILE") || fatal "losetup failed"
  command -v udevadm >/dev/null && udevadm settle || true
  for _ in {1..25}; do
    [[ -b "${LOOP}p1" && -b "${LOOP}p2" ]] && break
    sleep 0.1
    command -v udevadm >/dev/null && udevadm settle || true
  done
  [[ -b "${LOOP}p1" ]] || fatal "loop partitions not present for $LOOP"

  # Cleanup on exit: unmount first, then detach loop (tolerate absence)
  trap 'umount -lf "'"$IMG_BOOT_MNT"'" "'"$IMG_ROOT_MNT"'" 2>/dev/null; losetup -d "'"$LOOP"'" 2>/dev/null' EXIT

  # Mount image partitions read-only
  mount -o ro "${LOOP}p1" "$IMG_BOOT_MNT"
  mount -o ro "${LOOP}p2" "$IMG_ROOT_MNT"

  # Sanity checks without using failing pipelines
  # start*.elf must exist
  if ! compgen -G "$IMG_BOOT_MNT/start*.elf" > /dev/null; then
    fatal "start*.elf not found in image"
  fi
  # vmlinuz-* must exist
  if ! compgen -G "$IMG_ROOT_MNT/boot/vmlinuz-*" > /dev/null; then
    fatal "vmlinuz-* not found in image root"
  fi
}

confirm_and_wipe(){
  lsblk -o NAME,SIZE,MODEL,TRAN,LABEL "$NVME"
  read -rp "Type EXACTLY 'WIPE' to destroy ALL DATA on $NVME: " ACK
  [[ "$ACK" == "WIPE" ]] || fatal "Aborted"
  wipefs -a "$NVME"
  sgdisk -Zo "$NVME"
  # GPT: 1: 1MiB..513MiB vfat ESP; 2: rest LUKS
  parted -s "$NVME" mklabel gpt \
    mkpart system-boot fat32 1MiB 513MiB set 1 esp on \
    mkpart cryptroot 513MiB 100%
  partprobe "$NVME"; sleep 1
  mkfs.vfat -F32 -n system-boot "$(part 1)"
}

setup_luks(){
  echo "Create LUKS2 on $(part 2) (you will be prompted for a passphrase; keep it as fallback)"
  need cryptsetup
  cryptsetup luksFormat --type luks2 "$(part 2)"
  cryptsetup open "$(part 2)" cryptroot
  mkfs.ext4 -L rootfs /dev/mapper/cryptroot
}

mount_targets(){
  mkdir -p "$TGT_ROOT" "$TGT_BOOT"
  mount /dev/mapper/cryptroot "$TGT_ROOT"
  mkdir -p "$TGT_ROOT/boot/firmware"
  mount "$(part 1)" "$TGT_BOOT"
  mount --bind "$TGT_BOOT" "$TGT_ROOT/boot/firmware"
}

rsync_root_and_boot(){
  need rsync
  rsync -aAXH --numeric-ids --delete \
  --exclude='/boot/firmware' --exclude='/boot/firmware/**' \
  --exclude='/dev/*' --exclude='/proc/*' --exclude='/sys/*' \
  --exclude='/run/*' --exclude='/tmp/*' --exclude='/mnt/*' \
  --exclude='/media/*' --exclude='/lost+found' \
  "$IMG_ROOT_MNT"/ "$TGT_ROOT"/
  rsync -aH --delete "$IMG_BOOT_MNT"/ "$TGT_ROOT/boot/firmware"/
}

write_crypttab_fstab(){
  LUUID=$(blkid -s UUID -o value "$(part 2)")
  printf 'cryptroot UUID=%s none luks,discard,fido2-device=auto\n' "$LUUID" > "$TGT_ROOT/etc/crypttab"
  cat > "$TGT_ROOT/etc/fstab" <<EOF
/dev/mapper/cryptroot  /               ext4  defaults,discard,errors=remount-ro  0 1
LABEL=system-boot      /boot/firmware  vfat  defaults,umask=0077                 0 1
EOF
}

fix_firmware_files(){
  local C="$TGT_ROOT/boot/firmware/config.txt"
  local CL="$TGT_ROOT/boot/firmware/cmdline.txt"
  [[ -f "$C" ]] || fatal "missing $C"

  # Always boot the uncompressed Pi 5 kernel
  if grep -q '^kernel=' "$C"; then
    sed -i 's#^kernel=.*#kernel=kernel_2712.img#' "$C"
  else
    sed -i '1i kernel=kernel_2712.img' "$C"
  fi

  # Ensure initramfs and cmdline indirection are set
  grep -q '^initramfs ' "$C" || echo 'initramfs initrd.img followkernel' >> "$C"
  grep -q '^cmdline=cmdline.txt' "$C" || sed -i '1i cmdline=cmdline.txt' "$C"

  # Display & buses (Pi 5)
  grep -q '^dtoverlay=vc4-kms-v3d-pi5' "$C" || echo 'dtoverlay=vc4-kms-v3d-pi5' >> "$C"
  grep -q '^dtparam=i2c_arm=on'        "$C" || echo 'dtparam=i2c_arm=on' >> "$C"
  grep -q '^dtparam=pciex1=on'         "$C" || echo 'dtparam=pciex1=on' >> "$C"
  grep -q '^dtparam=pciex1_gen=2'      "$C" || echo 'dtparam=pciex1_gen=2' >> "$C"
  grep -q '^enable_uart=1'             "$C" || echo 'enable_uart=1' >> "$C"

  # Minimal, correct dracut hints using the bare UUID
  local LUUID; LUUID=$(blkid -s UUID -o value "$(part 2)")
  : > "$CL"
  {
    echo -n "rd.luks.uuid=$LUUID rd.luks.name=$LUUID=cryptroot "
    echo -n "root=/dev/mapper/cryptroot rootfstype=ext4 rootwait fixrtc "
    echo    "console=serial0,115200 console=tty1 ds=nocloud;s=file:///boot/firmware/ ${DSI_FLAGS} rd.debug"
  } >> "$CL"
}

seed_cloud_init(){
  # NoCloud seed to create user, lock down SSH, set hostname, and enable avahi.
  cat > "$TGT_ROOT/boot/firmware/user-data" <<EOF
#cloud-config
hostname: $STYX_HOSTNAME
manage_etc_hosts: true
users:
  - name: $STYX_USER
    gecos: "$STYX_USER"
    shell: /bin/bash
    groups: [sudo,video,i2c]
    sudo: ALL=(ALL) NOPASSWD:ALL
    lock_passwd: false
    ssh_authorized_keys:
      - $SSH_PUBKEY
chpasswd:
  list: |
    $STYX_USER:$STYX_PASS
  expire: true
ssh_pwauth: false
package_update: true
packages: [openssh-server, avahi-daemon]
runcmd:
  - systemctl enable --now ssh
  - systemctl enable --now avahi-daemon || true
EOF

  # Minimal meta-data for NoCloud
  date +%s | awk '{print "instance-id: iid-titan-ag-"$1"\nlocal-hostname: '"$STYX_HOSTNAME"'"}' \
    > "$TGT_ROOT/boot/firmware/meta-data"
}

prep_chroot_mounts(){
  for d in dev proc sys; do mount --bind "/$d" "$TGT_ROOT/$d"; done
  mount -t devpts devpts "$TGT_ROOT/dev/pts"
  # Replace the usual resolv.conf symlink with a real file for apt to work
  rm -f "$TGT_ROOT/etc/resolv.conf"
  cp /etc/resolv.conf "$TGT_ROOT/etc/resolv.conf"

  # Block service starts (no systemd in chroot)
  cat > "$TGT_ROOT/usr/sbin/policy-rc.d" <<'EOP'
#!/bin/sh
exit 101
EOP
  chmod +x "$TGT_ROOT/usr/sbin/policy-rc.d"

  # Ensure qemu static is present inside chroot
  install -D -m755 /usr/local/bin/qemu-aarch64-static "$TGT_ROOT/usr/bin/qemu-aarch64-static"
}

in_chroot(){
  chroot "$TGT_ROOT" /usr/bin/qemu-aarch64-static /bin/bash -lc '
set -euo pipefail
export DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC

# --- APT sources (ports) ---
cat > /etc/apt/sources.list <<'"'"'EOS'"'"'
deb http://ports.ubuntu.com/ubuntu-ports noble main restricted universe multiverse
deb http://ports.ubuntu.com/ubuntu-ports noble-updates main restricted universe multiverse
deb http://ports.ubuntu.com/ubuntu-ports noble-security main restricted universe multiverse
EOS

apt-get update

# --- Remove snaps and pin them off ---
apt-get -y purge snapd || true
rm -rf /snap /var/snap /var/lib/snapd /home/*/snap || true
mkdir -p /etc/apt/preferences.d
cat > /etc/apt/preferences.d/nosnap.pref <<'"'"'EOS'"'"'
Package: snapd
Pin: release *
Pin-Priority: -10
EOS

# --- Base tools (no flash-kernel; we use dracut) ---
apt-get install -y --no-install-recommends \
  openssh-client openssh-server openssh-sftp-server avahi-daemon \
  cryptsetup dracut fido2-tools libfido2-1 i2c-tools \
  python3-smbus python3-pil zbar-tools qrencode lm-sensors \
  file zstd lz4 || true

# Camera apps: try rpicam-apps; otherwise basic libcamera tools
apt-get install -y rpicam-apps || apt-get install -y libcamera-tools || true

# --- Persistent journal so we can read logs after failed boot ---
mkdir -p /etc/systemd/journald.conf.d
cat > /etc/systemd/journald.conf.d/99-persistent.conf <<'"'"'EOS'"'"'
[Journal]
Storage=persistent
EOS

# --- SSH hardening (ensure file exists even if package was half-installed) ---
if [ ! -f /etc/ssh/sshd_config ]; then
  mkdir -p /etc/ssh
  cat > /etc/ssh/sshd_config <<'"'"'EOS'"'"'
PermitRootLogin no
PasswordAuthentication no
KbdInteractiveAuthentication no
PubkeyAuthentication yes
# Accept defaults for the rest
EOS
fi
sed -i -e "s/^#\?PasswordAuthentication .*/PasswordAuthentication no/" \
       -e "s/^#\?KbdInteractiveAuthentication .*/KbdInteractiveAuthentication no/" \
       -e "s/^#\?PermitRootLogin .*/PermitRootLogin no/" \
       -e "s/^#\?PubkeyAuthentication .*/PubkeyAuthentication yes/" /etc/ssh/sshd_config || true

# --- Hostname & hosts ---
echo "'"$STYX_HOSTNAME"'" > /etc/hostname
if grep -q "^127\\.0\\.1\\.1" /etc/hosts; then
  sed -i "s/^127\\.0\\.1\\.1.*/127.0.1.1\t'"$STYX_HOSTNAME"'/" /etc/hosts
else
  echo -e "127.0.1.1\t'"$STYX_HOSTNAME"'" >> /etc/hosts
fi

# --- Enable services on first boot ---
mkdir -p /etc/systemd/system/multi-user.target.wants
ln -sf /lib/systemd/system/ssh.service /etc/systemd/system/multi-user.target.wants/ssh.service
ln -sf /lib/systemd/system/avahi-daemon.service /etc/systemd/system/multi-user.target.wants/avahi-daemon.service || true

# --- Ensure i2c group ---
getent group i2c >/dev/null || groupadd i2c

# --- Dracut configuration (generic, not host-only) ---
mkdir -p /etc/dracut.conf.d
cat > /etc/dracut.conf.d/00-hostonly.conf <<'"'"'EOS'"'"'
hostonly=no
EOS
cat > /etc/dracut.conf.d/10-systemd-crypt.conf <<'"'"'EOS'"'"'
add_dracutmodules+=" systemd crypt "
EOS
cat > /etc/dracut.conf.d/20-drivers.conf <<'"'"'EOS'"'"'
add_drivers+=" nvme xhci_pci xhci_hcd usbhid hid_generic hid "
EOS
cat > /etc/dracut.conf.d/30-fido2.conf <<'"'"'EOS'"'"'
install_items+="/usr/bin/systemd-cryptsetup /usr/bin/fido2-token /usr/lib/*/libfido2.so* /usr/lib/*/libcbor.so*"
EOS

# --- Build initramfs and place it where firmware expects it ---
KVER=$(ls -1 /lib/modules | sort -V | tail -n1)
dracut --force /boot/initramfs-$KVER.img $KVER
ln -sf initramfs-$KVER.img /boot/initrd.img
ln -sf initramfs-$KVER.img /boot/initrd.img-$KVER
cp -a /boot/initramfs-$KVER.img /boot/firmware/initrd.img

# --- Create uncompressed kernel for Pi 5 firmware ---
if [ -f "/usr/lib/linux-image-$KVER/Image" ]; then
  cp -a "/usr/lib/linux-image-$KVER/Image" /boot/firmware/kernel_2712.img
else
  FMT=$(file -b "/boot/vmlinuz-$KVER" || true)
  case "$FMT" in
    *Zstandard*|*zstd*) zstd -dc "/boot/vmlinuz-$KVER" > /boot/firmware/kernel_2712.img ;;
    *LZ4*)              lz4  -dc "/boot/vmlinuz-$KVER" > /boot/firmware/kernel_2712.img ;;
    *gzip*)             zcat     "/boot/vmlinuz-$KVER" > /boot/firmware/kernel_2712.img ;;
    *)                  cp -a "/boot/vmlinuz-$KVER" /boot/firmware/kernel_2712.img ;;
  esac
fi

# --- Ensure Pi 5 DTB is present on the boot partition ---
DTB=$(find /lib/firmware -type f -name "bcm2712-rpi-5-b.dtb" | sort | tail -n1 || true)
[ -n "$DTB" ] && cp -a "$DTB" /boot/firmware/

# --- Dracut hook to copy rdsosreport.txt to the FAT partition on failure ---
mkdir -p /usr/lib/dracut/modules.d/99copylog
cat > /usr/lib/dracut/modules.d/99copylog/module-setup.sh <<'"'"'EOS'"'"'
#!/bin/bash
check() { return 0; }
depends() { echo base; return 0; }
install() {
  # Guard $moddir for nounset; derive if absent
  local mdir="${moddir:-$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)}"
  inst_hook emergency 99 "$mdir/copylog.sh"
}
EOS
chmod +x /usr/lib/dracut/modules.d/99copylog/module-setup.sh

cat > /usr/lib/dracut/modules.d/99copylog/copylog.sh <<'"'"'EOS'"'"'
#!/bin/sh
set -e
for dev in /dev/nvme0n1p1 /dev/sda1 /dev/sdb1 /dev/mmcblk0p1; do
  [ -b "$dev" ] || continue
  mkdir -p /mnt/bootfat
  if mount -t vfat "$dev" /mnt/bootfat 2>/dev/null; then
    if [ -s /run/initramfs/rdsosreport.txt ]; then
      cp -f /run/initramfs/rdsosreport.txt /mnt/bootfat/rdsosreport.txt 2>/dev/null || true
      sync || true
    fi
    umount /mnt/bootfat || true
    break
  fi
done
EOS
chmod +x /usr/lib/dracut/modules.d/99copylog/copylog.sh

# Rebuild to ensure the copylog module is included
dracut --force /boot/initramfs-$KVER.img $KVER
ln -sf initramfs-$KVER.img /boot/initrd.img
cp -a /boot/initramfs-$KVER.img /boot/firmware/initrd.img

true
'
}

verify_boot_assets(){
  echo "---- verify boot assets on FAT ----"
  file "$TGT_ROOT/boot/firmware/kernel_2712.img" || true
  ls -lh "$TGT_ROOT/boot/firmware/initrd.img" || true
  echo "-- config.txt (key lines) --"
  grep -E '^(kernel|initramfs|cmdline)=|^dtoverlay=|^dtparam=' "$TGT_ROOT/boot/firmware/config.txt" || true
  echo "-- cmdline.txt --"
  cat "$TGT_ROOT/boot/firmware/cmdline.txt" || true
  echo "-- firmware blobs (sample) --"
  ls -1 "$TGT_ROOT/boot/firmware"/start*.elf "$TGT_ROOT/boot/firmware"/fixup*.dat | head -n 8 || true
  echo "-- Pi5 DTB --"
  ls -l "$TGT_ROOT/boot/firmware/"*rpi-5-b.dtb || true
}

enroll_fido_tokens(){
  echo "Enrolling FIDO2 Solo keys into $(part 2) ..."
  need systemd-cryptenroll
  need fido2-token

  # Collect all hidraw paths from both output styles (some distros print 'Device: /dev/hidrawX')
  mapfile -t DEVS < <(
    fido2-token -L \
      | sed -n 's,^\(/dev/hidraw[0-9]\+\):.*,\1,p; s,^Device:[[:space:]]\+/dev/hidraw\([0-9]\+\).*,/dev/hidraw\1,p' \
      | sort -u
  )

  if (( ${#DEVS[@]} == 0 )); then
    echo "No FIDO2 tokens detected; skipping enrollment (you can enroll later)."
    echo "Example later: systemd-cryptenroll $(part 2) --fido2-device=/dev/hidrawX --fido2-with-client-pin=no"
    return 0
  fi

  # Recommend keeping exactly ONE key plugged during first enrollment to avoid ambiguity.
  if (( ${#DEVS[@]} > 1 )); then
    echo "Note: multiple FIDO2 tokens present: ${DEVS[*]}"
    echo "If enrollment fails, try with only one key inserted."
  fi

  local rc=0
  for D in "${DEVS[@]}"; do
    echo "-> Enrolling $D (you should be asked to touch the key)"
    if ! SYSTEMD_LOG_LEVEL=debug systemd-cryptenroll "$(part 2)" \
          --fido2-device="$D" \
          --fido2-with-client-pin=no \
          --fido2-with-user-presence=yes \
          --fido2-with-user-verification=no \
          --label="solo-$(basename "$D")"; then
      echo "WARN: enrollment failed for $D"
      rc=1
    fi
  done

  echo "Tokens enrolled (if any):"
  systemd-cryptenroll "$(part 2)" --list || true
  return $rc
}

cleanup(){
  rm -f "$TGT_ROOT/usr/sbin/policy-rc.d" || true
  umount -lf "$TGT_ROOT/dev/pts" 2>/dev/null || true
  for d in dev proc sys; do umount -lf "$TGT_ROOT/$d" 2>/dev/null || true; done
  umount -lf "$TGT_ROOT/boot/firmware" 2>/dev/null || true
  umount -lf "$TGT_BOOT" 2>/dev/null || true
  umount -lf "$TGT_ROOT" 2>/dev/null || true
  cryptsetup close cryptroot 2>/dev/null || true
  umount -lf "$IMG_BOOT_MNT" 2>/dev/null || true
  umount -lf "$IMG_ROOT_MNT" 2>/dev/null || true
}

main(){
  require_root
  need losetup; need parted; need rsync
  auto_detect_target_disk
  echo "Target disk: $NVME"
  ensure_binfmt_aarch64
  ensure_image
  preflight_cleanup
  guard_target_device
  open_image
  confirm_and_wipe
  setup_luks
  mount_targets
  rsync_root_and_boot
  write_crypttab_fstab
  fix_firmware_files
  seed_cloud_init
  prep_chroot_mounts
  in_chroot
  verify_boot_assets
  need_host_fido2
  enroll_fido_tokens
  cleanup
  echo "✅ NVMe prepared."
  echo "   Install in the Pi 5 and boot with no SD."
  echo "   Expect LUKS to unlock automatically with a Solo key inserted;"
  echo "   passphrase fallback remains. Hostname: ${STYX_HOSTNAME}  User: ${STYX_USER}"
  echo "   On first boot, reach it via: ssh -i ~/.ssh/id_ed25519_titan styx@titan-ag.local"
}

main "$@"