#!/bin/bash
# Verify netboot image build completeness and correctness
# Run this after 'make build' to validate the generated image

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
HTTP_DIR="$SCRIPT_DIR/http"
TMPDIR=$(mktemp -d)
EXIT_CODE=0

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

error() {
    echo -e "${RED}✗${NC} $1"
    EXIT_CODE=1
}

success() {
    echo -e "${GREEN}✓${NC} $1"
}

warning() {
    echo -e "${YELLOW}!${NC} $1"
}

info() {
    echo "  $1"
}

cleanup() {
    rm -rf "$TMPDIR"
}
trap cleanup EXIT

echo "Verifying netboot image build..."
echo ""

# Check 1: Required files exist
echo "Checking required files..."
for file in vmlinuz initrd-netboot.img filesystem.squashfs version.txt boot.ipxe; do
    if [ -f "$HTTP_DIR/$file" ]; then
        success "$file exists"
    else
        error "$file missing"
    fi
done
echo ""

# Check 2: File types are correct
echo "Checking file types..."
if file "$HTTP_DIR/vmlinuz" | grep -q "Linux kernel"; then
    success "vmlinuz is a valid Linux kernel"
else
    error "vmlinuz is not a valid kernel"
fi

if file "$HTTP_DIR/initrd-netboot.img" | grep -q "cpio archive"; then
    success "initrd-netboot.img is a valid initramfs"
else
    error "initrd-netboot.img is not a valid initramfs"
fi

if file "$HTTP_DIR/filesystem.squashfs" | grep -q "Squashfs filesystem"; then
    success "filesystem.squashfs is a valid squashfs image"
else
    error "filesystem.squashfs is not a valid squashfs"
fi
echo ""

# Check 3: File permissions for HTTP serving
echo "Checking file permissions..."
for file in vmlinuz initrd-netboot.img filesystem.squashfs; do
    PERMS=$(stat -c "%a" "$HTTP_DIR/$file")
    if [ "$PERMS" = "644" ]; then
        success "$file has correct permissions (644)"
    else
        warning "$file has permissions $PERMS (expected 644)"
    fi
done
echo ""

# Check 4: File sizes are reasonable
echo "Checking file sizes..."
KERNEL_SIZE=$(stat -c %s "$HTTP_DIR/vmlinuz")
INITRD_SIZE=$(stat -c %s "$HTTP_DIR/initrd-netboot.img")
SQUASHFS_SIZE=$(stat -c %s "$HTTP_DIR/filesystem.squashfs")

if [ "$KERNEL_SIZE" -gt 5000000 ] && [ "$KERNEL_SIZE" -lt 50000000 ]; then
    success "vmlinuz size reasonable: $(numfmt --to=iec-i --suffix=B $KERNEL_SIZE)"
else
    warning "vmlinuz size unusual: $(numfmt --to=iec-i --suffix=B $KERNEL_SIZE)"
fi

if [ "$INITRD_SIZE" -gt 10000000 ] && [ "$INITRD_SIZE" -lt 100000000 ]; then
    success "initrd size reasonable: $(numfmt --to=iec-i --suffix=B $INITRD_SIZE)"
else
    warning "initrd size unusual: $(numfmt --to=iec-i --suffix=B $INITRD_SIZE)"
fi

if [ "$SQUASHFS_SIZE" -gt 500000000 ] && [ "$SQUASHFS_SIZE" -lt 2000000000 ]; then
    success "squashfs size reasonable: $(numfmt --to=iec-i --suffix=B $SQUASHFS_SIZE)"
else
    warning "squashfs size unusual: $(numfmt --to=iec-i --suffix=B $SQUASHFS_SIZE)"
fi
echo ""

# Check 5: Initramfs contains custom netboot script and binaries
echo "Checking initramfs contents..."
unmkinitramfs "$HTTP_DIR/initrd-netboot.img" "$TMPDIR/initramfs" 2>/dev/null

if [ -f "$TMPDIR/initramfs/main/scripts/netboot" ]; then
    success "Custom netboot script present in initramfs"

    # Verify it matches source
    if diff -q "$SCRIPT_DIR/initramfs/scripts/netboot" "$TMPDIR/initramfs/main/scripts/netboot" >/dev/null 2>&1; then
        success "netboot script matches source"
    else
        warning "netboot script differs from source"
    fi
else
    error "Custom netboot script missing from initramfs"
fi

# Check required binaries
for binary in wget curl awk; do
    if [ -f "$TMPDIR/initramfs/main/usr/bin/$binary" ]; then
        success "$binary binary present"
    else
        error "$binary binary missing from initramfs"
    fi
done

# Check unsquashfs in both possible locations
if [ -f "$TMPDIR/initramfs/main/usr/bin/unsquashfs" ] || [ -f "$TMPDIR/initramfs/main/usr/sbin/unsquashfs" ]; then
    success "unsquashfs binary present"
else
    error "unsquashfs binary missing from initramfs"
fi

if [ -f "$TMPDIR/initramfs/main/usr/sbin/switch_root" ]; then
    success "switch_root binary present"
else
    error "switch_root binary missing from initramfs"
fi
echo ""

# Check 6: Required kernel modules configured
echo "Checking kernel modules configuration..."
if [ -f "$TMPDIR/initramfs/main/conf/modules" ]; then
    MODULES_FILE="$TMPDIR/initramfs/main/conf/modules"

    for module in squashfs overlay r8125 ext4 isofs af_packet nls_iso8859-1; do
        if grep -q "^$module" "$MODULES_FILE"; then
            success "Module $module configured"
        else
            error "Module $module not configured"
        fi
    done
else
    warning "modules configuration file not found"
fi
echo ""

# Check 7: Squashfs root filesystem contents
echo "Checking squashfs root filesystem..."

# Check critical paths exist (handle usr-merge symlinks)
for path in boot etc root usr var; do
    if unsquashfs -ll "$HTTP_DIR/filesystem.squashfs" | grep -q "squashfs-root/$path$"; then
        success "/$path directory exists"
    else
        error "/$path directory missing"
    fi
done

# Check bin and lib (may be symlinks due to usr-merge)
for path in bin lib; do
    if unsquashfs -ll "$HTTP_DIR/filesystem.squashfs" | grep -qE "squashfs-root/$path( |->|$)"; then
        success "/$path exists (directory or symlink)"
    else
        error "/$path missing"
    fi
done

# Check SSH authorized_keys
if unsquashfs -cat "$HTTP_DIR/filesystem.squashfs" root/.ssh/authorized_keys >/dev/null 2>&1; then
    KEY_COUNT=$(unsquashfs -cat "$HTTP_DIR/filesystem.squashfs" root/.ssh/authorized_keys 2>/dev/null | grep -c "^ssh-")
    success "SSH authorized_keys present ($KEY_COUNT keys)"
else
    error "SSH authorized_keys missing or invalid"
fi

# Check netplan config
if unsquashfs -cat "$HTTP_DIR/filesystem.squashfs" etc/netplan/01-netcfg.yaml >/dev/null 2>&1; then
    success "Netplan configuration present"
else
    error "Netplan configuration missing"
fi

# Check K3s dependencies
for pkg in containerd runc; do
    if unsquashfs -ll "$HTTP_DIR/filesystem.squashfs" | grep -q "usr/bin/$pkg"; then
        success "K3s dependency: $pkg"
    else
        warning "K3s dependency missing: $pkg (may be in /usr/sbin)"
    fi
done
echo ""

# Check 8: Version info
echo "Build information:"
if [ -f "$HTTP_DIR/version.txt" ]; then
    cat "$HTTP_DIR/version.txt" | sed 's/^/  /'
else
    warning "version.txt not found"
fi
echo ""

# Check 9: iPXE boot configuration
echo "Checking iPXE configuration..."
if [ -f "$HTTP_DIR/boot.ipxe" ]; then
    success "boot.ipxe present"

    # Extract server URL from boot.ipxe
    SERVER_URL=$(grep "http://" "$HTTP_DIR/boot.ipxe" | head -1 | grep -oP 'http://[^/]+')
    if [ -n "$SERVER_URL" ]; then
        info "Boot server configured: $SERVER_URL"
    fi

    # Check if all required files are referenced
    for file in vmlinuz initrd-netboot.img filesystem.squashfs; do
        if grep -q "$file" "$HTTP_DIR/boot.ipxe"; then
            success "boot.ipxe references $file"
        else
            error "boot.ipxe missing reference to $file"
        fi
    done

    # Check for required kernel parameters
    if grep -q "overlayroot=tmpfs" "$HTTP_DIR/boot.ipxe"; then
        success "boot.ipxe configures overlayroot"
    else
        warning "boot.ipxe missing overlayroot parameter"
    fi
else
    error "boot.ipxe missing"
fi
echo ""

# Summary
echo "========================================="
if [ $EXIT_CODE -eq 0 ]; then
    echo -e "${GREEN}✓ All checks passed!${NC}"
    echo "Image is ready for deployment."
else
    echo -e "${RED}✗ Some checks failed${NC}"
    echo "Review errors above before deploying."
fi
echo "========================================="

exit $EXIT_CODE