#!/bin/bash

# Extract kernel and initrd from Ubuntu cloud image
# This script ensures kernel boot works automatically

set -e

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

log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
    exit 1
}

BASE_SUBVOL="/var/lib/vms/base"
BASE_IMAGE_PATH="$BASE_SUBVOL/ubuntu-24.04-server-cloudimg-amd64.raw"
KERNEL_PATH="$BASE_SUBVOL/vmlinuz"
INITRD_PATH="$BASE_SUBVOL/initrd.img"

# Check if kernel and initrd already exist
if [ -f "$KERNEL_PATH" ] && [ -f "$INITRD_PATH" ]; then
    log "Kernel and initrd already extracted"
    exit 0
fi

# Check if base image exists
if [ ! -f "$BASE_IMAGE_PATH" ]; then
    error "Base image not found at $BASE_IMAGE_PATH"
fi

log "Extracting kernel and initrd from base image..."

# Create temporary mount point
TEMP_MOUNT=$(mktemp -d)
trap "umount '$TEMP_MOUNT' 2>/dev/null || true; losetup -d /dev/loop1 2>/dev/null || true; rmdir '$TEMP_MOUNT'" EXIT

# Mount the image
losetup -P /dev/loop1 "$BASE_IMAGE_PATH"

# Mount the boot partition (partition 16 contains kernel/initrd)
mount /dev/loop1p16 "$TEMP_MOUNT"

# Copy kernel and initrd
if [ -f "$TEMP_MOUNT/vmlinuz-6.8.0-60-generic" ] && [ -f "$TEMP_MOUNT/initrd.img-6.8.0-60-generic" ]; then
    cp "$TEMP_MOUNT/vmlinuz-6.8.0-60-generic" "$KERNEL_PATH"
    cp "$TEMP_MOUNT/initrd.img-6.8.0-60-generic" "$INITRD_PATH"
    log "✓ Kernel and initrd extracted successfully"
else
    error "Kernel or initrd not found in boot partition"
fi

# Cleanup is handled by trap
log "Kernel extraction completed"