#!/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
# When run with ssh target, rsyncs lxc/ and runs this script inside the container.

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

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"
  ssh "$TARGET" "bash /tmp/cm4-network-boot-lxc/setup.sh"
  echo "Done."
  exit 0
fi

# --- Running inside the LXC from here ---
echo "Configuring network boot (DHCP + TFTP on eth1, NAT via eth0) ..."

# 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); 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 on eth1 only (provisioning LAN)
# TFTP and PXE options are in network-boot-pxe.conf, controlled by toggle-network-boot-dhcp.sh
interface=eth1
bind-interfaces
dhcp-range=10.20.50.100,10.20.50.200,12h
log-dhcp
log-queries
port=0
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: 10.20.50.0/24 -> 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 10.20.50.0/24 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 10.20.50.0/24 -o eth0 -j MASQUERADE 2>/dev/null || \
    iptables -t nat -A POSTROUTING -s 10.20.50.0/24 -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 (10.20.50.1), range 10.20.50.100-200"
echo "  - NAT: 10.20.50.0/24 -> eth0 (internet)"
echo "  - TFTP root: /srv/tftpboot (RPi boot files; initrd.img if provided)"