#!/usr/bin/env bash
# pifleet-netbootctl — idempotent PXE/NFS lifecycle manager for diskless Pi 4/5 fleet
# Requires: root on the netboot management server
set -Eeuo pipefail

VERSION="1.0.0"
NETBOOT_ROOT="${NETBOOT_ROOT:-/srv/netboot}"
TFTP_ROOT="${TFTP_ROOT:-$NETBOOT_ROOT/tftpboot}"
NFS_ROOT="${NFS_ROOT:-$NETBOOT_ROOT/nfs}"
BASE_IMAGE="${BASE_IMAGE:-$NFS_ROOT/base-image}"
NODES_DIR="${NODES_DIR:-$NFS_ROOT/nodes}"
INVENTORY="${INVENTORY:-$NETBOOT_ROOT/inventory/nodes.tsv}"
DNSMASQ_D="${DNSMASQ_D:-/etc/dnsmasq.d}"
EXPORTS_D="${EXPORTS_D:-/etc/exports.d}"
SERVER_IP="${SERVER_IP:-}"
NFS_OPTS="${NFS_OPTS:-vers=3,tcp,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noatime}"
BOOT_ORDER_PI4="${BOOT_ORDER_PI4:-0xf21}"
BOOT_ORDER_PI5="${BOOT_ORDER_PI5:-0xf241}"

log()  { printf '[%s] %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)" "$*"; }
die()  { log "ERROR: $*"; exit 1; }
need_root() { [[ "${EUID}" -eq 0 ]] || die "run as root"; }

detect_server_ip() {
  if [[ -n "${SERVER_IP}" ]]; then
    printf '%s' "$SERVER_IP"
    return
  fi
  # Prefer primary default-route IPv4
  local ip
  ip="$(ip -4 route get 1.1.1.1 2>/dev/null | awk '{for(i=1;i<=NF;i++) if($i=="src"){print $(i+1); exit}}' || true)"
  [[ -n "$ip" ]] || ip="$(hostname -I 2>/dev/null | awk '{print $1}')"
  [[ -n "$ip" ]] || die "could not detect SERVER_IP; export SERVER_IP=x.x.x.x"
  printf '%s' "$ip"
}

usage() {
  cat <<EOF
pifleet-netbootctl v${VERSION}

Usage:
  $0 setup-server                 Install packages, dirs, dnsmasq proxy-DHCP, NFS
  $0 import-base IMAGE.{img,xz}   Unpack Raspberry Pi OS Lite 64-bit as golden base
  $0 add-node SERIAL HOSTNAME MAC [IP] [MODEL]
  $0 remove-node SERIAL
  $0 sync-node SERIAL             Re-sync boot files from base; preserve /etc identity
  $0 rebuild-exports              Rewrite NFS exports from inventory
  $0 rebuild-dnsmasq              Rewrite per-node dnsmasq tag mappings
  $0 list                         Show inventory
  $0 status                       Service + mount health

Inventory columns (TAB-separated): SERIAL HOSTNAME MAC IP MODEL
EOF
}

ensure_dirs() {
  install -d -m 755 \
    "$NETBOOT_ROOT" "$TFTP_ROOT" "$NFS_ROOT" "$NODES_DIR" \
    "$NETBOOT_ROOT/inventory" "$NETBOOT_ROOT/keys" "$NETBOOT_ROOT/cloud-init" \
    "$EXPORTS_D" "$DNSMASQ_D"
  [[ -f "$INVENTORY" ]] || printf 'SERIAL\tHOSTNAME\tMAC\tIP\tMODEL\n' > "$INVENTORY"
}

install_packages() {
  export DEBIAN_FRONTEND=noninteractive
  apt-get update -qq
  apt-get install -y -qq dnsmasq nfs-kernel-server rpcbind rsync xz-utils kpartx \
    qemu-user-static binfmt-support curl ca-certificates
}

write_sysctl() {
  cat > /etc/sysctl.d/99-pifleet-netboot.conf <<'SYSCTL'
# High concurrent NFS root boots (Pi fleet)
net.core.rmem_max = 16777216
net.core.wmem_max = 16777216
net.core.rmem_default = 1048576
net.core.wmem_default = 1048576
net.core.netdev_max_backlog = 5000
net.ipv4.tcp_rmem = 4096 87380 16777216
net.ipv4.tcp_wmem = 4096 65536 16777216
net.ipv4.neigh.default.gc_thresh1 = 4096
net.ipv4.neigh.default.gc_thresh2 = 8192
net.ipv4.neigh.default.gc_thresh3 = 16384
net.ipv4.neigh.default.base_reachable_time_ms = 600000
fs.file-max = 2097152
# SunRPC / NFS server
sunrpc.tcp_slot_table_entries = 128
sunrpc.tcp_max_slot_table_entries = 128
vm.swappiness = 10
vm.dirty_ratio = 20
vm.dirty_background_ratio = 5
SYSCTL
  sysctl --system >/dev/null
}

write_dnsmasq_base() {
  local sip
  sip="$(detect_server_ip)"
  cat > "$DNSMASQ_D/pifleet-pxe.conf" <<EOF
# PiFleet Proxy-DHCP + TFTP (does not replace primary DHCP)
port=0
log-dhcp
enable-tftp
tftp-root=${TFTP_ROOT}
tftp-no-fail
# ProxyDHCP: answer PXE only; lease IP from existing router
dhcp-range=${sip},proxy
# Raspberry Pi PXE (PXEClient)
pxe-service=0,"Raspberry Pi Boot"
# Per-node boot directories resolved by MAC → serial tag (generated)
conf-dir=${DNSMASQ_D}/pifleet-nodes,*.conf
EOF
  # Disable default dnsmasq if conflicting with systemd-resolved — leave to operator
  systemctl enable --now dnsmasq
  systemctl restart dnsmasq
  log "dnsmasq proxy-DHCP listening; tftp-root=${TFTP_ROOT}; server=${sip}"
}

write_nfs_base_export() {
  # Base image read-only for cloning; nodes RW
  cat > "$EXPORTS_D/pifleet-base.exports" <<EOF
${BASE_IMAGE} 10.0.0.0/8(ro,sync,no_subtree_check,no_root_squash,fsid=1000)
${BASE_IMAGE} 192.168.0.0/16(ro,sync,no_subtree_check,no_root_squash,fsid=1000)
${BASE_IMAGE} 172.16.0.0/12(ro,sync,no_subtree_check,no_root_squash,fsid=1000)
EOF
  exportfs -ra
  systemctl enable --now nfs-kernel-server rpcbind
}

cmd_setup_server() {
  need_root
  ensure_dirs
  install_packages
  write_sysctl
  write_dnsmasq_base
  write_nfs_base_export
  # NFS thread scaling
  mkdir -p /etc/nfs.conf.d
  cat > /etc/nfs.conf.d/pifleet.conf <<'NFS'
[nfsd]
threads=32
UDP=n
vers2=n
vers3=y
vers4=y
NFS
  systemctl restart nfs-kernel-server
  log "setup-server complete. Next: $0 import-base /path/to/raspios-lite-arm64.img.xz"
}

mount_image_root() {
  local img="$1" mnt="$2"
  mkdir -p "$mnt"
  local loop
  loop="$(losetup --show -fP "$img")"
  # Prefer partition 2 (root) on standard Raspberry Pi OS images
  local part="${loop}p2"
  [[ -b "$part" ]] || part="${loop}p1"
  mount -o ro "$part" "$mnt"
  printf '%s' "$loop"
}

cmd_import_base() {
  need_root
  ensure_dirs
  local src="${1:-}" 
  [[ -n "$src" && -f "$src" ]] || die "usage: import-base IMAGE.{img,xz}"
  local work="$NETBOOT_ROOT/work"
  mkdir -p "$work"
  local img="$src"
  if [[ "$src" == *.xz ]]; then
    img="$work/base.img"
    log "decompressing $src ..."
    xz -T0 -dc "$src" > "$img"
  fi
  local mnt="$work/mnt"
  umount "$mnt" 2>/dev/null || true
  local loop
  loop="$(mount_image_root "$img" "$mnt")"
  log "syncing golden image → ${BASE_IMAGE}"
  mkdir -p "$BASE_IMAGE"
  rsync -aHAX --delete --info=progress2 "$mnt"/ "$BASE_IMAGE"/
  umount "$mnt"
  losetup -d "$loop" || true

  # Diskless base adjustments (applied once to golden image)
  prepare_base_image_for_nfs
  log "base image ready at $BASE_IMAGE"
}

prepare_base_image_for_nfs() {
  local root="$BASE_IMAGE"
  [[ -d "$root" ]] || die "missing base image $root"

  # fstab: tmpfs for ephemeral; drop SD card mounts
  cat > "$root/etc/fstab" <<'FSTAB'
proc            /proc           proc    defaults          0       0
tmpfs           /tmp            tmpfs   nosuid,nodev,size=256M  0 0
tmpfs           /var/tmp        tmpfs   nosuid,nodev,size=64M   0 0
tmpfs           /var/log        tmpfs   nosuid,nodev,size=128M  0 0
tmpfs           /var/tmp/log    tmpfs   nosuid,nodev,size=32M   0 0
FSTAB

  # Ensure SSH + systemd-networkd friendly defaults
  mkdir -p "$root/etc/ssh/sshd_config.d" "$root/etc/systemd/system/multi-user.target.wants"
  cat > "$root/etc/ssh/sshd_config.d/99-pifleet.conf" <<'SSH'
PasswordAuthentication no
KbdInteractiveAuthentication no
PermitRootLogin no
PubkeyAuthentication yes
SSH

  # First-boot identity seeder (reads kernel cmdline + serial)
  install -d -m 755 "$root/usr/local/lib/pifleet" "$root/usr/local/lib/systemd/system"
  cat > "$root/usr/local/lib/pifleet/apply-node-identity.sh" <<'IDENT'
#!/bin/bash
set -euo pipefail
# Apply hostname, SSH keys, optional static IP from files dropped at provision time
SERIAL="$(tr -d '\0' </proc/device-tree/serial-number 2>/dev/null || true)"
SERIAL="${SERIAL:-unknown}"
ID_DIR="/etc/pifleet/identity"
mkdir -p /etc/pifleet

if [[ -f "$ID_DIR/hostname" ]]; then
  HOST="$(tr -d ' \n' < "$ID_DIR/hostname")"
  if [[ -n "$HOST" ]]; then
    hostnamectl set-hostname "$HOST" 2>/dev/null || echo "$HOST" > /etc/hostname
    if grep -q '^10\.77\.' /etc/hosts 2>/dev/null; then :; fi
    sed -i "/pifleet-node/d" /etc/hosts 2>/dev/null || true
    echo "127.0.1.1 $HOST $HOST.pifleet-node" >> /etc/hosts
  fi
fi

if [[ -d "$ID_DIR/ssh" ]]; then
  mkdir -p /home/pi/.ssh /root/.ssh
  if [[ -f "$ID_DIR/ssh/authorized_keys" ]]; then
    install -m 600 -o pi -g pi "$ID_DIR/ssh/authorized_keys" /home/pi/.ssh/authorized_keys 2>/dev/null \
      || install -m 600 "$ID_DIR/ssh/authorized_keys" /home/pi/.ssh/authorized_keys
    install -m 600 "$ID_DIR/ssh/authorized_keys" /root/.ssh/authorized_keys
  fi
  chmod 700 /home/pi/.ssh 2>/dev/null || true
fi

# Optional: systemd network static (if provided)
if [[ -f "$ID_DIR/eth0.network" ]]; then
  install -m 644 "$ID_DIR/eth0.network" /etc/systemd/network/10-pifleet-eth0.network
  networkctl reload 2>/dev/null || true
fi

echo "$SERIAL" > /etc/pifleet/serial
IDENT
  chmod 755 "$root/usr/local/lib/pifleet/apply-node-identity.sh"

  cat > "$root/usr/local/lib/systemd/system/pifleet-identity.service" <<'UNIT'
[Unit]
Description=PiFleet node identity (hostname, SSH keys, network)
After=local-fs.target
Before=ssh.service

[Service]
Type=oneshot
ExecStart=/usr/local/lib/pifleet/apply-node-identity.sh
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
UNIT
  ln -sfn /usr/local/lib/systemd/system/pifleet-identity.service \
    "$root/etc/systemd/system/multi-user.target.wants/pifleet-identity.service"

  # Prefer predictable interface names off for classic eth0 on Pi
  mkdir -p "$root/etc/systemd/network"
}

inventory_has() {
  local serial="$1"
  awk -F'\t' -v s="$serial" 'NR>1 && $1==s {found=1} END{exit !found}' "$INVENTORY"
}

inventory_get() {
  local serial="$1"
  awk -F'\t' -v s="$serial" 'NR>1 && $1==s {print; exit}' "$INVENTORY"
}

inventory_upsert() {
  local serial="$1" hostname="$2" mac="$3" ip="$4" model="$5"
  local tmp
  tmp="$(mktemp)"
  awk -F'\t' -v s="$serial" 'NR==1 || $1!=s {print}' "$INVENTORY" > "$tmp"
  printf '%s\t%s\t%s\t%s\t%s\n' "$serial" "$hostname" "$mac" "$ip" "$model" >> "$tmp"
  mv "$tmp" "$INVENTORY"
}

normalize_mac() {
  echo "$1" | tr 'A-F' 'a-f' | sed 's/-/:/g'
}

seed_boot_files() {
  local serial="$1" node_nfs="$2"
  local boot="$TFTP_ROOT/$serial"
  local sip
  sip="$(detect_server_ip)"
  mkdir -p "$boot"

  # Copy firmware/kernel from node (or base) boot partition files if present
  local boot_src=""
  if [[ -d "$node_nfs/boot/firmware" ]]; then
    boot_src="$node_nfs/boot/firmware"
  elif [[ -d "$node_nfs/boot" ]]; then
    boot_src="$node_nfs/boot"
  elif [[ -d "$BASE_IMAGE/boot/firmware" ]]; then
    boot_src="$BASE_IMAGE/boot/firmware"
  else
    boot_src="$BASE_IMAGE/boot"
  fi

  if [[ -d "$boot_src" ]]; then
    rsync -a --delete \
      --include='*.dat' --include='*.elf' --include='*.bin' \
      --include='kernel*.img' --include='*.dtb' --include='overlays/***' \
      --include='config.txt' --include='cmdline.txt' \
      --exclude='*' \
      "$boot_src"/ "$boot"/ 2>/dev/null || rsync -a "$boot_src"/ "$boot"/
  fi

  # Pi netboot: TFTP path is /SERIAL/... relative to tftp-root when bootloader uses serial prefix
  cat > "$boot/cmdline.txt" <<EOF
console=serial0,115200 console=tty1 root=/dev/nfs nfsroot=${sip}:${node_nfs},${NFS_OPTS} rw ip=dhcp rootwait elevator=deadline cgroup_enable=memory
EOF

  # Ensure config.txt enables network / correct arm_64bit
  if [[ ! -f "$boot/config.txt" ]]; then
    cat > "$boot/config.txt" <<'CFG'
arm_64bit=1
kernel=kernel8.img
dtparam=audio=on
camera_auto_detect=1
display_auto_detect=1
dtoverlay=vc4-kms-v3d
max_framebuffers=2
disable_overscan=1
CFG
  fi

  # Also place a copy of bootcode.bin / start files at tftp root for older boot flow
  if [[ -f "$boot/bootcode.bin" ]]; then
    cp -a "$boot/bootcode.bin" "$TFTP_ROOT/bootcode.bin" 2>/dev/null || true
  fi
}

write_node_identity() {
  local node_nfs="$1" hostname="$2" mac="$3" ip="$4"
  local id="$node_nfs/etc/pifleet/identity"
  mkdir -p "$id/ssh"
  printf '%s\n' "$hostname" > "$id/hostname"
  printf '%s\n' "$mac" > "$id/mac"
  printf '%s\n' "$ip" > "$id/ip"

  # SSH keys: fleet key + optional per-node
  local fleet_keys="$NETBOOT_ROOT/keys/authorized_keys"
  if [[ -f "$fleet_keys" ]]; then
    install -m 600 "$fleet_keys" "$id/ssh/authorized_keys"
  else
    # Generate fleet keypair once
    if [[ ! -f "$NETBOOT_ROOT/keys/fleet_ed25519" ]]; then
      ssh-keygen -t ed25519 -N '' -f "$NETBOOT_ROOT/keys/fleet_ed25519" -C "pifleet-fleet@$(hostname)"
    fi
    cat "$NETBOOT_ROOT/keys/fleet_ed25519.pub" > "$id/ssh/authorized_keys"
    chmod 600 "$id/ssh/authorized_keys"
  fi

  if [[ -n "$ip" && "$ip" != "-" ]]; then
    cat > "$id/eth0.network" <<EOF
[Match]
Name=eth0 en*

[Network]
Address=${ip}/24
Gateway=$(echo "$ip" | awk -F. '{print $1"."$2"."$3".1"}')
DNS=1.1.1.1
DNS=8.8.8.8
DHCP=no
EOF
  fi

  # Embed serial for inventory match
  mkdir -p "$node_nfs/etc/pifleet"
}

write_node_export() {
  local serial="$1" path="$2"
  # Unique fsid from serial hash
  local fsid
  fsid="$(printf '%s' "$serial" | cksum | awk '{print ($1 % 900000) + 1001}')"
  cat > "$EXPORTS_D/pifleet-node-${serial}.exports" <<EOF
${path} 10.0.0.0/8(rw,sync,no_subtree_check,no_root_squash,fsid=${fsid})
${path} 192.168.0.0/16(rw,sync,no_subtree_check,no_root_squash,fsid=${fsid})
${path} 172.16.0.0/12(rw,sync,no_subtree_check,no_root_squash,fsid=${fsid})
EOF
}

write_node_dnsmasq() {
  local serial="$1" mac="$2"
  mac="$(normalize_mac "$mac")"
  # Tag client by MAC; set dhcp-option for vendor — Pi uses TFTP path based on serial once bootloader runs
  # Map MAC → directory via symlink at tftp for alternate flows
  mkdir -p "$DNSMASQ_D/pifleet-nodes"
  cat > "$DNSMASQ_D/pifleet-nodes/${serial}.conf" <<EOF
# Node ${serial}
dhcp-host=${mac},set:pi_${serial}
# Optional: fixed IP when not using proxy-only mode — leave commented for ProxyDHCP
# dhcp-host=${mac},10.77.0.50,${serial}
EOF
  # Convenience: lowercase MAC path some firmwares probe
  local macflat
  macflat="$(echo "$mac" | tr -d ':')"
  ln -sfn "$serial" "$TFTP_ROOT/${macflat}" 2>/dev/null || true
}

cmd_add_node() {
  need_root
  ensure_dirs
  [[ -d "$BASE_IMAGE" ]] || die "base image missing; run import-base first"
  local serial="${1:-}" hostname="${2:-}" mac="${3:-}" ip="${4:--}" model="${5:-pi4}"
  [[ -n "$serial" && -n "$hostname" && -n "$mac" ]] || die "usage: add-node SERIAL HOSTNAME MAC [IP] [MODEL]"
  serial="$(echo "$serial" | tr 'A-F' 'a-f')"
  mac="$(normalize_mac "$mac")"
  local node_nfs="$NODES_DIR/$serial"

  log "provisioning node serial=$serial hostname=$hostname mac=$mac ip=$ip model=$model"
  mkdir -p "$node_nfs"
  rsync -aHAX --delete --info=stats2 "$BASE_IMAGE"/ "$node_nfs"/
  write_node_identity "$node_nfs" "$hostname" "$mac" "$ip"
  seed_boot_files "$serial" "$node_nfs"
  write_node_export "$serial" "$node_nfs"
  write_node_dnsmasq "$serial" "$mac"
  inventory_upsert "$serial" "$hostname" "$mac" "$ip" "$model"
  exportfs -ra
  systemctl restart dnsmasq
  log "node $serial ready — set EEPROM BOOT_ORDER=${BOOT_ORDER_PI4} (Pi4) or ${BOOT_ORDER_PI5} (Pi5)"
  log "  tftp: $TFTP_ROOT/$serial"
  log "  nfs:  $node_nfs"
}

cmd_remove_node() {
  need_root
  local serial="${1:-}"
  [[ -n "$serial" ]] || die "usage: remove-node SERIAL"
  serial="$(echo "$serial" | tr 'A-F' 'a-f')"
  local node_nfs="$NODES_DIR/$serial"
  # Fail if exported and busy
  if mount | grep -q "$node_nfs"; then
    die "node path appears busy/mounted by clients: $node_nfs"
  fi
  rm -f "$EXPORTS_D/pifleet-node-${serial}.exports"
  rm -f "$DNSMASQ_D/pifleet-nodes/${serial}.conf"
  rm -rf "$TFTP_ROOT/$serial"
  # Keep NFS data unless FORCE=1
  if [[ "${FORCE:-0}" == "1" ]]; then
    rm -rf "$node_nfs"
  else
    log "left $node_nfs in place (FORCE=1 to delete rootfs)"
  fi
  local tmp
  tmp="$(mktemp)"
  awk -F'\t' -v s="$serial" 'NR==1 || $1!=s {print}' "$INVENTORY" > "$tmp"
  mv "$tmp" "$INVENTORY"
  exportfs -ra
  systemctl restart dnsmasq
  log "removed node $serial from boot services"
}

cmd_sync_node() {
  need_root
  local serial="${1:-}"
  [[ -n "$serial" ]] || die "usage: sync-node SERIAL"
  serial="$(echo "$serial" | tr 'A-F' 'a-f')"
  inventory_has "$serial" || die "unknown serial $serial"
  local node_nfs="$NODES_DIR/$serial"
  local id_backup
  id_backup="$(mktemp -d)"
  [[ -d "$node_nfs/etc/pifleet" ]] && cp -a "$node_nfs/etc/pifleet" "$id_backup/" || true
  rsync -aHAX --delete \
    --exclude '/etc/pifleet/***' \
    --exclude '/home/***' \
    --exclude '/var/lib/pifleet/***' \
    "$BASE_IMAGE"/ "$node_nfs"/
  [[ -d "$id_backup/pifleet" ]] && cp -a "$id_backup/pifleet" "$node_nfs/etc/" || true
  rm -rf "$id_backup"
  seed_boot_files "$serial" "$node_nfs"
  log "synced $serial from base (identity preserved)"
}

cmd_rebuild_exports() {
  need_root
  ensure_dirs
  write_nfs_base_export
  awk -F'\t' 'NR>1 && $1!="" {
    s=$1; path="'"$NODES_DIR"'/" s;
    print s, path
  }' "$INVENTORY" | while read -r s path; do
    [[ -d "$path" ]] && write_node_export "$s" "$path"
  done
  exportfs -ra
  log "exports rebuilt"
}

cmd_rebuild_dnsmasq() {
  need_root
  ensure_dirs
  write_dnsmasq_base
  mkdir -p "$DNSMASQ_D/pifleet-nodes"
  awk -F'\t' 'NR>1 && $1!="" {print $1"\t"$3}' "$INVENTORY" | while IFS=$'\t' read -r s m; do
    write_node_dnsmasq "$s" "$m"
  done
  systemctl restart dnsmasq
  log "dnsmasq rebuilt"
}

cmd_list() {
  ensure_dirs
  column -t -s $'\t' "$INVENTORY" 2>/dev/null || cat "$INVENTORY"
}

cmd_status() {
  echo "== services =="
  systemctl is-active dnsmasq nfs-kernel-server rpcbind 2>/dev/null || true
  echo "== exports =="
  exportfs -v 2>/dev/null | head -40 || true
  echo "== tftp =="
  ls -la "$TFTP_ROOT" 2>/dev/null | head -20 || true
  echo "== inventory =="
  cmd_list
}

main() {
  local cmd="${1:-}"
  shift || true
  case "$cmd" in
    setup-server) cmd_setup_server "$@" ;;
    import-base) cmd_import_base "$@" ;;
    add-node) cmd_add_node "$@" ;;
    remove-node) cmd_remove_node "$@" ;;
    sync-node) cmd_sync_node "$@" ;;
    rebuild-exports) cmd_rebuild_exports "$@" ;;
    rebuild-dnsmasq) cmd_rebuild_dnsmasq "$@" ;;
    list) cmd_list "$@" ;;
    status) cmd_status "$@" ;;
    -h|--help|help|"") usage ;;
    *) die "unknown command: $cmd" ;;
  esac
}

main "$@"