#!/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"

# Decrypt secrets from phoenix (requires SSH access as the invoking user, not root)
echo "Decrypting secrets from phoenix..."
SECRETS_FILE="$SCRIPT_DIR/secrets/netboot.sops.yaml"
SUDO_USER_HOME=$(getent passwd "${SUDO_USER:-$USER}" | cut -d: -f6)
if [ -f "$SECRETS_FILE" ]; then
    # Run SSH as the original user (not root) to use their SSH keys
    ROOT_PW_HASH=$(sudo -u "${SUDO_USER:-$USER}" bash -c "cat '$SECRETS_FILE' | ssh phoenix 'sops -d --input-type yaml --output-type yaml /dev/stdin'" | grep root_password_hash | cut -d' ' -f2)
    if [ -z "$ROOT_PW_HASH" ]; then
        echo "WARNING: Failed to decrypt root password, console login will be disabled"
        ROOT_PW_HASH="*"
    fi
else
    echo "WARNING: No secrets file found at $SECRETS_FILE, console login will be disabled"
    ROOT_PW_HASH="*"
fi

# 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

# Write root password hash to temp file for chroot to read
# Use /root/ not /tmp/ because systemd installation may mount tmpfs over /tmp
mkdir -p "$BUILD_DIR/rootfs/root"
if [ -n "$ROOT_PW_HASH" ] && [ "$ROOT_PW_HASH" != "*" ]; then
    echo "$ROOT_PW_HASH" > "$BUILD_DIR/rootfs/root/.pw_hash"
    echo "Root password hash written to rootfs"
else
    echo "*" > "$BUILD_DIR/rootfs/root/.pw_hash"
    echo "WARNING: No valid password hash, console login will be disabled"
fi

# 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

# Enable noble-updates and noble-security for HWE kernel
cat > /etc/apt/sources.list.d/ubuntu.sources << 'SOURCES'
Types: deb
URIs: http://archive.ubuntu.com/ubuntu
Suites: noble noble-updates noble-backports
Components: main universe multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg

Types: deb
URIs: http://security.ubuntu.com/ubuntu
Suites: noble-security
Components: main universe multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg
SOURCES

apt-get update
apt-get upgrade -y

# Install HWE kernel (6.14+) for RTL8125BP XID 689 support
apt-get install -y \
    linux-image-generic-hwe-24.04 \
    linux-firmware \
    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 \
    open-iscsi

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

# Vulkan drivers for GPU compute workloads (ollama, llama.cpp)
apt-get install -y \
    mesa-vulkan-drivers \
    libvulkan1 \
    vulkan-tools

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

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

# Don't set static hostname - let DHCP provide it via networkd
# Empty /etc/hostname allows transient hostname from DHCP
echo "" > /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 - disable socket activation, use traditional daemon
sed -i 's/#PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
sed -i 's/#PubkeyAuthentication.*/PubkeyAuthentication yes/' /etc/ssh/sshd_config
# Disable socket activation (Ubuntu 24.04 default) and use traditional sshd
systemctl disable ssh.socket 2>/dev/null || true
rm -f /etc/systemd/system/ssh.service.requires/ssh.socket 2>/dev/null || true
rm -f /etc/systemd/system/sockets.target.wants/ssh.socket 2>/dev/null || true
systemctl enable ssh

# Fix SSH host key permissions (must be 0600 for private keys, sshd refuses otherwise)
chmod 600 /etc/ssh/ssh_host_*_key
chmod 644 /etc/ssh/ssh_host_*_key.pub

# 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

# Set root password from decrypted hash (for console login only)
ROOT_PW_HASH=$(cat /root/.pw_hash)
echo "root:$ROOT_PW_HASH" | chpasswd -e
rm -f /root/.pw_hash

# 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/"

# Install node storage setup service
echo "Installing node storage setup service..."
FILES_DIR="$SCRIPT_DIR/files"
cp "$FILES_DIR/setup-node-storage" "$BUILD_DIR/rootfs/usr/local/bin/"
chmod +x "$BUILD_DIR/rootfs/usr/local/bin/setup-node-storage"
cp "$FILES_DIR/setup-node-storage.service" "$BUILD_DIR/rootfs/etc/systemd/system/"
# Enable the service (create symlink manually since we can't run systemctl)
mkdir -p "$BUILD_DIR/rootfs/etc/systemd/system/multi-user.target.wants"
ln -sf /etc/systemd/system/setup-node-storage.service \
    "$BUILD_DIR/rootfs/etc/systemd/system/multi-user.target.wants/setup-node-storage.service"

# Install DHCP hostname service
echo "Installing DHCP hostname service..."
cp "$FILES_DIR/set-hostname-from-dhcp" "$BUILD_DIR/rootfs/usr/local/bin/"
chmod +x "$BUILD_DIR/rootfs/usr/local/bin/set-hostname-from-dhcp"
cp "$FILES_DIR/set-hostname-from-dhcp.service" "$BUILD_DIR/rootfs/etc/systemd/system/"
ln -sf /etc/systemd/system/set-hostname-from-dhcp.service \
    "$BUILD_DIR/rootfs/etc/systemd/system/multi-user.target.wants/set-hostname-from-dhcp.service"

# Download and install K3s binary
echo "Downloading K3s binary..."
K3S_VERSION="v1.34.3+k3s1"
curl -sfL "https://github.com/k3s-io/k3s/releases/download/${K3S_VERSION}/k3s" \
    -o "$BUILD_DIR/rootfs/usr/local/bin/k3s"
chmod +x "$BUILD_DIR/rootfs/usr/local/bin/k3s"
echo "K3s $K3S_VERSION installed"

# Install K3s agent service
echo "Installing K3s agent service..."
# Create K3s directories first (will be bind-mounted from NVMe at runtime)
mkdir -p "$BUILD_DIR/rootfs/etc/rancher/k3s"
mkdir -p "$BUILD_DIR/rootfs/etc/rancher/node"
mkdir -p "$BUILD_DIR/rootfs/var/lib/rancher/k3s/agent"
cp "$FILES_DIR/k3s-agent.service" "$BUILD_DIR/rootfs/etc/systemd/system/"
cp "$FILES_DIR/k3s-agent.env" "$BUILD_DIR/rootfs/etc/rancher/k3s/"
# Enable the service
ln -sf /etc/systemd/system/k3s-agent.service \
    "$BUILD_DIR/rootfs/etc/systemd/system/multi-user.target.wants/k3s-agent.service"

# 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