#!/bin/bash
# Netboot image builder for diskless K3s nodes
# Builds Ubuntu Noble base system with SquashFS

set -e

# Get the directory where this script is located (works on any machine)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

BUILD_DIR="$SCRIPT_DIR/build"
IMAGE_DIR="$SCRIPT_DIR/images"
HTTP_DIR="$SCRIPT_DIR/http"
VERSION=$(date +%Y%m%d-%H%M)

echo "Building netboot image version $VERSION"

# Clean previous build - unmount any stray mounts first
if [ -d "$BUILD_DIR/rootfs" ]; then
    echo "Cleaning up previous build mounts..."
    mount | grep "$BUILD_DIR/rootfs" | awk '{print $3}' | while read mount; do
        umount -l "$mount" 2>/dev/null || true
    done
fi

rm -rf $BUILD_DIR/rootfs
mkdir -p $BUILD_DIR/rootfs

# Copy custom initramfs config into rootfs BEFORE debootstrap, so it's available during chroot
echo "Setting up custom initramfs configuration..."
INITRAMFS_CONFIG="$SCRIPT_DIR/initramfs"
if [ ! -d "$INITRAMFS_CONFIG" ]; then
    echo "ERROR: Custom initramfs config not found at $INITRAMFS_CONFIG"
    exit 1
fi
# These dirs will be created by debootstrap, so we'll copy after that

# Create base Ubuntu system
echo "Running debootstrap (this will take several minutes)..."
debootstrap --arch=amd64 --variant=minbase --components=main,universe,multiverse \
    noble $BUILD_DIR/rootfs \
    http://archive.ubuntu.com/ubuntu

# Chroot and configure
cat << 'CHROOT_SCRIPT' > $BUILD_DIR/rootfs/setup.sh
#!/bin/bash
set -e

# Set keyboard layout and encoding to Norwegian UTF-8 (non-interactive)
export DEBIAN_FRONTEND=noninteractive
export DEBCONF_NONINTERACTIVE_SEEN=true
echo "keyboard-configuration keyboard-configuration/layout select Norwegian" | debconf-set-selections
echo "keyboard-configuration keyboard-configuration/variant select Norwegian" | debconf-set-selections
echo "locales locales/default_environment_locale select en_US.UTF-8" | debconf-set-selections
echo "locales locales/locales_to_be_generated multiselect en_US.UTF-8 UTF-8, nb_NO.UTF-8 UTF-8" | debconf-set-selections

# Update and upgrade
apt-get update
apt-get upgrade -y

# Install essential packages
apt-get install -y \
    linux-image-generic \
    linux-firmware \
    cloud-initramfs-rooturl \
    busybox-initramfs \
    initramfs-tools \
    keyboard-configuration \
    systemd \
    systemd-sysv \
    dbus \
    udev \
    kmod \
    iproute2 \
    iputils-ping \
    netplan.io \
    openssh-server \
    curl \
    wget \
    ca-certificates \
    gnupg \
    sudo \
    locales

# K3s prerequisites
apt-get install -y \
    apparmor \
    apparmor-utils \
    iptables \
    nftables \
    conntrack \
    socat \
    ethtool \
    nfs-common

# Container runtime prerequisites
apt-get install -y \
    containerd \
    runc

# Useful tools
apt-get install -y \
    htop \
    iotop \
    vim \
    less \
    rsync \
    git \
    squashfs-tools

# Clean up
apt-get clean
rm -rf /var/lib/apt/lists/*
rm -rf /tmp/*
rm -rf /var/tmp/*

# Configure hostname (will be overridden by netplan)
echo "k3s-node" > /etc/hostname

# Configure network with netplan
cat > /etc/netplan/01-netcfg.yaml <<EOF
network:
  version: 2
  renderer: networkd
  ethernets:
    all-en:
      match:
        name: en*
      dhcp4: true
      dhcp6: false
      dhcp4-overrides:
        use-hostname: true
EOF

# Enable systemd-networkd
systemctl enable systemd-networkd
systemctl enable systemd-resolved

# Configure SSH
sed -i 's/#PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
sed -i 's/#PubkeyAuthentication.*/PubkeyAuthentication yes/' /etc/ssh/sshd_config
systemctl enable ssh

# Create SSH directory for root
mkdir -p /root/.ssh
chmod 700 /root/.ssh

# Add your SSH public key here
cat >> /root/.ssh/authorized_keys <<SSHKEY
ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIMuTf8duvWFhU3CTlLFqr1+EeszA1Cu4e5Q07A9xKy1J lindahl@lindahl-Legion-5-Pro-16ACH6H
ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAILpyobN2sYtxnFFZCvl1nNK4MLc9oVTKgGhkgHy2u0tX lindahl@phoenix.home
SSHKEY

chmod 600 /root/.ssh/authorized_keys

# Disable password authentication completely
echo "root:*" | chpasswd -e

# Configure tmpfs mounts for ephemeral data
cat >> /etc/fstab <<FSTAB
# Ephemeral filesystems (RAM-based)
tmpfs /tmp tmpfs defaults,noatime,nosuid,nodev,noexec,mode=1777,size=2G 0 0
tmpfs /var/tmp tmpfs defaults,noatime,nosuid,nodev,noexec,mode=1777,size=1G 0 0
tmpfs /var/log tmpfs defaults,noatime,nosuid,nodev,noexec,mode=0755,size=1G 0 0
tmpfs /run tmpfs defaults,noatime,nosuid,nodev,noexec,mode=0755,size=512M 0 0
FSTAB

# Configure journald to use tmpfs
mkdir -p /etc/systemd/journald.conf.d
cat > /etc/systemd/journald.conf.d/tmpfs.conf <<JOURNAL
[Journal]
Storage=volatile
RuntimeMaxUse=256M
JOURNAL

# Configure systemd to not wait for network
systemctl disable systemd-networkd-wait-online.service

# Disable unnecessary services
systemctl disable apt-daily.timer
systemctl disable apt-daily-upgrade.timer
systemctl disable motd-news.timer

# Set timezone
ln -sf /usr/share/zoneinfo/UTC /etc/localtime

# Generate locales
echo "en_US.UTF-8 UTF-8" > /etc/locale.gen
locale-gen

# Build initramfs will be done outside the chroot after mount cleanup
echo "Initramfs build will be done after chroot cleanup"

CHROOT_SCRIPT

# Make script executable and run in chroot with proper mounts
chmod +x $BUILD_DIR/rootfs/setup.sh

# Mount required filesystems for DKMS compilation
mount -t proc proc $BUILD_DIR/rootfs/proc
mount -t sysfs sysfs $BUILD_DIR/rootfs/sys
mount -t devtmpfs devtmpfs $BUILD_DIR/rootfs/dev
mount -t devpts devpts $BUILD_DIR/rootfs/dev/pts

# Run setup script in chroot
chroot $BUILD_DIR/rootfs /setup.sh

# Unmount filesystems
umount -l $BUILD_DIR/rootfs/dev/pts
umount -l $BUILD_DIR/rootfs/dev
umount -l $BUILD_DIR/rootfs/sys
umount -l $BUILD_DIR/rootfs/proc

rm $BUILD_DIR/rootfs/setup.sh

# Copy custom initramfs config into rootfs while /proc is mounted
echo "Installing custom initramfs hooks and scripts..."
INITRAMFS_CONFIG="$SCRIPT_DIR/initramfs"
cp "$INITRAMFS_CONFIG/initramfs.conf" "$BUILD_DIR/rootfs/etc/initramfs-tools/"
cp "$INITRAMFS_CONFIG/modules" "$BUILD_DIR/rootfs/etc/initramfs-tools/"
cp -r "$INITRAMFS_CONFIG/hooks/"* "$BUILD_DIR/rootfs/usr/share/initramfs-tools/hooks/"
cp -r "$INITRAMFS_CONFIG/scripts/"* "$BUILD_DIR/rootfs/usr/share/initramfs-tools/scripts/"

# Build initramfs while /proc/sys/dev are still mounted
echo "Building custom netboot initramfs..."
KERNEL_VERSION=$(ls -1 $BUILD_DIR/rootfs/boot/vmlinuz-* | sed 's|.*/vmlinuz-||' | head -1)
chroot $BUILD_DIR/rootfs mkinitramfs -v -o /boot/initrd-netboot.img $KERNEL_VERSION
INITRAMFS_OUTPUT="$BUILD_DIR/rootfs/boot/initrd-netboot.img"
echo "Initramfs build complete. Size: $(du -h $INITRAMFS_OUTPUT | cut -f1)"

# Copy kernel and netboot initramfs
mkdir -p $IMAGE_DIR/$VERSION
cp $BUILD_DIR/rootfs/boot/vmlinuz-* $IMAGE_DIR/$VERSION/vmlinuz
cp $BUILD_DIR/rootfs/boot/initrd-netboot.img $IMAGE_DIR/$VERSION/initrd-netboot.img

echo "Creating squashfs image..."
mksquashfs $BUILD_DIR/rootfs \
    $IMAGE_DIR/$VERSION/filesystem.squashfs \
    -comp xz \
    -Xbcj x86 \
    -b 1M \
    -noappend \
    -no-progress

# Create version info file
cat > $IMAGE_DIR/$VERSION/version.txt <<EOF
Build Date: $(date)
Ubuntu Version: Noble (24.04)
Kernel: $(chroot $BUILD_DIR/rootfs dpkg -l | grep linux-image-generic | awk '{print $3}')
Image Size: $(du -h $IMAGE_DIR/$VERSION/filesystem.squashfs | awk '{print $1}')
EOF

echo "Image created successfully: $IMAGE_DIR/$VERSION/"
ls -lh $IMAGE_DIR/$VERSION/

# Create symlink to latest
ln -sfn $VERSION $IMAGE_DIR/latest

# Copy to HTTP directory
echo "Deploying to HTTP directory..."
rsync -av $IMAGE_DIR/$VERSION/ $HTTP_DIR/

# Fix permissions for web server access
echo "Setting permissions for HTTP serving..."
chmod 644 $HTTP_DIR/vmlinuz
chmod 644 $HTTP_DIR/initrd-netboot.img
chmod 644 $HTTP_DIR/filesystem.squashfs
chmod 644 $HTTP_DIR/version.txt

echo "Build complete! Version: $VERSION"
echo "Files available at: $HTTP_DIR/"
ls -lh $HTTP_DIR/vmlinuz $HTTP_DIR/initrd-netboot.img $HTTP_DIR/filesystem.squashfs