#!/usr/bin/env bash
# Setup network boot on the provisioning LXC: DHCP + TFTP on eth1, NAT so LAN uses eth0 for internet.
# Run inside the LXC (as root), or from your machine: ./setup-network-boot-on-lxc.sh root@10.130.60.141 [SUBNET]
# SUBNET optional: A.B.C.D/PREFIX (e.g. 10.100.1.1/24). When run with ssh target, writes lan-subnet.conf on LXC if SUBNET given.
# When run with ssh target, rsyncs lxc/ and runs this script inside the container. Subnet is read from /opt/cm4-provisioning/lan-subnet.conf.

set -e
TARGET="${1:-}"
SUBNET_ARG="${2:-}"

if [[ -n "$TARGET" ]]; then
  # Run remotely: sync lxc/ and script, then execute inside LXC
  SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
  REPO_DIR="$(dirname "$SCRIPT_DIR")"
  echo "Syncing lxc config and script to $TARGET ..."
  rsync -a "$REPO_DIR/lxc/" "$TARGET:/tmp/cm4-network-boot-lxc/" --exclude='.git'
  if [[ -f "$REPO_DIR/network-boot-initramfs/initrd.img" ]]; then
    echo "Copying initrd.img to $TARGET ..."
    scp "$REPO_DIR/network-boot-initramfs/initrd.img" "$TARGET:/tmp/cm4-network-boot-lxc/initrd.img"
  else
    echo "Note: network-boot-initramfs/initrd.img not found (run build.sh first); skipping."
  fi
  scp "$SCRIPT_DIR/setup-network-boot-on-lxc.sh" "$TARGET:/tmp/cm4-network-boot-lxc/setup.sh"
  # If SUBNET_ARG given, write lan-subnet.conf on LXC so inner script uses the set subnet
  if [[ -n "$SUBNET_ARG" ]]; then
    if [[ "$SUBNET_ARG" =~ ^([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+)/([0-9]+)$ ]]; then
      LAN_GW="${BASH_REMATCH[1]}"
      PREFIX="${BASH_REMATCH[2]}"
      BASE_3="${LAN_GW%.*}"
      LAN_CIDR="${BASE_3}.0/${PREFIX}"
      DHCP_RANGE_START="${BASE_3}.100"
      DHCP_RANGE_END="${BASE_3}.200"
      ssh "$TARGET" "mkdir -p /opt/cm4-provisioning && echo 'LAN_GW=$LAN_GW' > /opt/cm4-provisioning/lan-subnet.conf && echo 'LAN_CIDR=$LAN_CIDR' >> /opt/cm4-provisioning/lan-subnet.conf && echo 'DHCP_RANGE_START=$DHCP_RANGE_START' >> /opt/cm4-provisioning/lan-subnet.conf && echo 'DHCP_RANGE_END=$DHCP_RANGE_END' >> /opt/cm4-provisioning/lan-subnet.conf"
      echo "Wrote lan-subnet.conf on LXC (LAN_GW=$LAN_GW, DHCP ${DHCP_RANGE_START}-${DHCP_RANGE_END})."
    else
      echo "Warning: SUBNET must be A.B.C.D/PREFIX (e.g. 10.100.1.1/24); ignoring '$SUBNET_ARG'."
    fi
  fi
  ssh "$TARGET" "bash /tmp/cm4-network-boot-lxc/setup.sh"
  echo "Done."
  exit 0
fi

# --- Running inside the LXC from here ---
# LAN subnet: use /opt/cm4-provisioning/lan-subnet.conf (written by deploy-to-proxmox.sh or passed as SUBNET when running remotely)
# Optional first arg when running locally: A.B.C.D/PREFIX to set/write lan-subnet.conf
LAN_CONF="/opt/cm4-provisioning/lan-subnet.conf"
if [[ "$1" =~ ^([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+)/([0-9]+)$ ]]; then
  LAN_GW="${BASH_REMATCH[1]}"
  PREFIX="${BASH_REMATCH[2]}"
  BASE_3="${LAN_GW%.*}"
  LAN_CIDR="${BASE_3}.0/${PREFIX}"
  DHCP_RANGE_START="${BASE_3}.100"
  DHCP_RANGE_END="${BASE_3}.200"
  mkdir -p /opt/cm4-provisioning
  echo "LAN_GW=$LAN_GW" > "$LAN_CONF"
  echo "LAN_CIDR=$LAN_CIDR" >> "$LAN_CONF"
  echo "DHCP_RANGE_START=$DHCP_RANGE_START" >> "$LAN_CONF"
  echo "DHCP_RANGE_END=$DHCP_RANGE_END" >> "$LAN_CONF"
  echo "Using set subnet: $LAN_CIDR (gateway $LAN_GW), DHCP ${DHCP_RANGE_START}-${DHCP_RANGE_END}."
elif [[ -f "$LAN_CONF" ]]; then
  source "$LAN_CONF"
else
  LAN_GW="10.20.50.1"
  LAN_CIDR="10.20.50.0/24"
  DHCP_RANGE_START="10.20.50.100"
  DHCP_RANGE_END="10.20.50.200"
  echo "No lan-subnet.conf and no SUBNET argument; using defaults: $LAN_CIDR."
fi
echo "Configuring network boot (DHCP + TFTP on eth1, NAT via eth0) — LAN $LAN_CIDR (gateway $LAN_GW), DHCP ${DHCP_RANGE_START}-${DHCP_RANGE_END} ..."

# 1) Install dnsmasq
if ! command -v dnsmasq >/dev/null 2>&1; then
  apt-get update -qq && DEBIAN_FRONTEND=noninteractive apt-get install -y -qq dnsmasq
fi

# 2) dnsmasq config for eth1 only (DHCP + TFTP + DNS); PXE options in network-boot-pxe.conf (toggle with toggle-network-boot-dhcp.sh)
mkdir -p /etc/dnsmasq.d
cat > /etc/dnsmasq.d/network-boot.conf << DNSMASQ
# DHCP + DNS on eth1 only (provisioning LAN)
# TFTP and PXE options in network-boot-pxe.conf, controlled by toggle-network-boot-dhcp.sh
interface=eth1
bind-interfaces
dhcp-range=${DHCP_RANGE_START},${DHCP_RANGE_END},12h
# DNS: file.server resolves to this host (eth1) so scripts can use http://file.server/...
address=/file.server/${LAN_GW}
# Explicitly send this host as DNS server to DHCP clients (option 6) so they use LXC DNS and resolve file.server
dhcp-option=6,${LAN_GW}
# Other DNS queries forwarded via LXC's resolv.conf
log-dhcp
log-queries
DNSMASQ
mkdir -p /opt/cm4-provisioning
if [ -f /tmp/cm4-network-boot-lxc/toggle-network-boot-dhcp.sh ]; then
  cp /tmp/cm4-network-boot-lxc/toggle-network-boot-dhcp.sh /opt/cm4-provisioning/
  chmod +x /opt/cm4-provisioning/toggle-network-boot-dhcp.sh
  /opt/cm4-provisioning/toggle-network-boot-dhcp.sh enable
fi

# 3) TFTP root: fetch Raspberry Pi 4 boot files from GitHub if missing
mkdir -p /srv/tftpboot
if [[ ! -f /srv/tftpboot/start4cd.elf ]]; then
  echo "Fetching Raspberry Pi firmware boot files from GitHub ..."
  if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then
    apt-get update -qq && DEBIAN_FRONTEND=noninteractive apt-get install -y -qq curl
  fi
  tmpdir=$(mktemp -d)
  trap "rm -rf $tmpdir" EXIT
  if command -v curl >/dev/null 2>&1; then
    curl -sL "https://github.com/raspberrypi/firmware/archive/refs/heads/master.tar.gz" -o "$tmpdir/firmware.tar.gz"
  else
    wget -q -O "$tmpdir/firmware.tar.gz" "https://github.com/raspberrypi/firmware/archive/refs/heads/master.tar.gz"
  fi
  tar xzf "$tmpdir/firmware.tar.gz" -C "$tmpdir"
  cp -a "$tmpdir/firmware-master/boot/." /srv/tftpboot/
  rm -rf "$tmpdir"
  echo "Copied RPi boot files to /srv/tftpboot"
else
  echo "TFTP root already has boot files (start4cd.elf present), skipping fetch."
fi

# 3b) Copy provisioning initrd.img to TFTP root if provided
if [[ -f /tmp/cm4-network-boot-lxc/initrd.img ]]; then
  cp /tmp/cm4-network-boot-lxc/initrd.img /srv/tftpboot/initrd.img
  echo "Copied initrd.img to /srv/tftpboot"
  if [[ -f /srv/tftpboot/config.txt ]] && ! grep -q 'initramfs initrd.img' /srv/tftpboot/config.txt 2>/dev/null; then
    echo "" >> /srv/tftpboot/config.txt
    echo "# Provisioning initramfs (network-boot-initramfs)" >> /srv/tftpboot/config.txt
    echo "initramfs initrd.img followkernel" >> /srv/tftpboot/config.txt
    echo "Added initramfs line to config.txt"
  fi
fi

# 4) IP forwarding (LAN clients use WAN)
echo 'net.ipv4.ip_forward=1' > /etc/sysctl.d/99-cm4-network-boot.conf
sysctl -p /etc/sysctl.d/99-cm4-network-boot.conf 2>/dev/null || sysctl -w net.ipv4.ip_forward=1

# 5) NAT: LAN subnet -> eth0 (masquerade)
if command -v nft >/dev/null 2>&1; then
  mkdir -p /etc/nftables.d
  cat > /etc/nftables.d/nat-lan.conf << NFT
table ip nat {
	chain postrouting {
		type nat hook postrouting priority srcnat; policy accept;
		ip saddr ${LAN_CIDR} oifname "eth0" masquerade
	}
}
NFT
  if ! nft list table ip nat 2>/dev/null | grep -q postrouting; then
    nft -f /etc/nftables.d/nat-lan.conf
  fi
  # Ensure main config includes our drop-in (Debian)
  if [[ -f /etc/nftables.conf ]] && ! grep -q 'nftables.d/nat-lan' /etc/nftables.conf 2>/dev/null; then
    echo 'include "/etc/nftables.d/nat-lan.conf"' >> /etc/nftables.conf
  fi
  echo "NAT rule added (nftables) and saved to /etc/nftables.d/nat-lan.conf"
else
  # Fallback iptables
  iptables -t nat -C POSTROUTING -s "${LAN_CIDR}" -o eth0 -j MASQUERADE 2>/dev/null || \
    iptables -t nat -A POSTROUTING -s "${LAN_CIDR}" -o eth0 -j MASQUERADE
  echo "NAT rule added (iptables)."
fi

# 6) Enable and start dnsmasq
systemctl enable dnsmasq
systemctl restart dnsmasq

echo "Network boot setup done."
echo "  - DHCP + TFTP on eth1 ($LAN_GW), range ${DHCP_RANGE_START}-${DHCP_RANGE_END}"
echo "  - NAT: ${LAN_CIDR} -> eth0 (internet)"
echo "  - TFTP root: /srv/tftpboot (RPi boot files; initrd.img if provided)"