#!/bin/bash

# Ubuntu VM Start Script with Cloud Hypervisor and Btrfs Thin Provisioning
# Usage: ubuntu_vm_start.sh $vm_number $name $mem $cores

set -e

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

# Configuration
BASE_IMAGE_NAME="ubuntu-24.04-server-cloudimg-amd64"
BASE_IMAGE_URL="https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img"
FIRMWARE_URL="https://github.com/cloud-hypervisor/rust-hypervisor-firmware/releases/download/0.5.0/hypervisor-fw"
VM_BASE_DIR="/var/lib/vms"
BTRFS_MOUNT_POINT="/var/lib/vms"
BASE_SUBVOL="$BTRFS_MOUNT_POINT/base"
VMS_SUBVOL="$BTRFS_MOUNT_POINT/vms"

# Network configuration
BRIDGE_NAME="br0"
BRIDGE_IP="192.168.100.1/24"
NETWORK="192.168.100.0/24"

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

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

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
    
    # If VM_NUMBER is set and we're in VM creation phase, clean up
    if [ -n "$VM_NUMBER" ] && [ -n "$VM_PID" ]; then
        cleanup_failed_vm "$VM_NUMBER"
    fi
    
    exit 1
}

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

# Test functions
test_step() {
    local step_name="$1"
    local test_command="$2"
    
    info "Testing: $step_name"
    if eval "$test_command"; then
        log "✓ Test passed: $step_name"
        return 0
    else
        error "✗ Test failed: $step_name"
        return 1
    fi
}

test_file_exists() {
    local file_path="$1"
    local description="$2"
    test_step "$description" "[ -f '$file_path' ]"
}

test_directory_exists() {
    local dir_path="$1"
    local description="$2"
    test_step "$description" "[ -d '$dir_path' ]"
}

test_command_exists() {
    local command="$1"
    local description="$2"
    test_step "$description" "command -v '$command' &> /dev/null"
}

test_network_interface() {
    local interface="$1"
    local description="$2"
    test_step "$description" "ip link show '$interface' &>/dev/null"
}

test_process_running() {
    local pid="$1"
    local description="$2"
    test_step "$description" "kill -0 '$pid' 2>/dev/null"
}

# Cleanup function for failed VM creation
cleanup_failed_vm() {
    local vm_number="$1"
    warn "VM creation failed, cleaning up..."
    
    # Call the delete script to clean up
    local delete_script="$(dirname "$0")/ubuntu_vm_delete.sh"
    if [ -f "$delete_script" ]; then
        log "Running cleanup script: $delete_script"
        "$delete_script" "$vm_number" || warn "Cleanup script failed, manual cleanup may be required"
    else
        warn "Delete script not found at $delete_script, manual cleanup required"
    fi
}

# Generate a proper password hash for 'ubuntu'
generate_password_hash() {
    # Generate salt and hash for password 'ubuntu'
    python3 -c "import crypt; print(crypt.crypt('ubuntu', crypt.mksalt(crypt.METHOD_SHA512)))"
}

# Wait for VM to boot and verify static IP
wait_for_vm_boot() {
    local vm_name="$1"
    local expected_ip="$2"
    local max_wait=180  # 3 minutes, increased from 120
    local count=0
    
    log "Waiting for VM '$vm_name' to boot with static IP $expected_ip..."
    
    while [ $count -lt $max_wait ]; do
        # Check if VM process is still running
        if ! kill -0 "$VM_PID" 2>/dev/null; then
            error "VM process died while waiting for boot. Check log: $VM_LOG_FILE"
        fi
        
        # Try to ping the expected static IP
        if ping -c 2 -W 3 "$expected_ip" >/dev/null 2>&1; then
            log "VM is responding at static IP address: $expected_ip"
            echo "$expected_ip"
            return 0
        fi
        
        # Also check ARP table for our MAC address
        local vm_ip=$(arp -a | grep "$VM_MAC" | sed 's/.*(\([^)]*\)).*/\1/' | head -1)
        if [ -n "$vm_ip" ] && [ "$vm_ip" = "$expected_ip" ]; then
            log "VM MAC address found in ARP table with expected IP: $expected_ip"
            echo "$expected_ip"
            return 0
        fi
        
        sleep 3
        count=$((count + 3))
        if [ $((count % 15)) -eq 0 ]; then
            info "Still waiting for VM to boot... ($count/$max_wait seconds)"
            info "VM process PID $VM_PID is still running"
            info "Expected static IP: $expected_ip"
            # Show recent log entries
            if [ -f "$VM_LOG_FILE" ]; then
                info "Recent VM log entries:"
                tail -3 "$VM_LOG_FILE" 2>/dev/null || true
            fi
        fi
    done
    
    warn "VM did not respond at expected IP $expected_ip within $max_wait seconds"
    warn "VM may still be booting - check manually with: ping $expected_ip"
    return 1
}

# Test IP connectivity
test_ip_connectivity() {
    local vm_ip="$1"
    local max_attempts=10
    local attempt=1
    
    log "Testing IP connectivity to $vm_ip..."
    
    while [ $attempt -le $max_attempts ]; do
        info "Ping attempt $attempt/$max_attempts to $vm_ip"
        
        # Test ping connectivity with timeout
        if ping -c 3 -W 2 "$vm_ip" >/dev/null 2>&1; then
            log "✓ IP connectivity successful to $vm_ip"
            return 0
        fi
        
        sleep 3
        attempt=$((attempt + 1))
    done
    
    error "✗ IP connectivity failed after $max_attempts attempts to $vm_ip"
    return 1
}

# Test SSH connectivity
test_ssh_connection() {
    local vm_ip="$1"
    local max_attempts=10
    local attempt=1
    
    log "Testing SSH connectivity to $vm_ip..."
    
    while [ $attempt -le $max_attempts ]; do
        info "SSH attempt $attempt/$max_attempts"
        
        # Test SSH connection with timeout
        if timeout 10 sshpass -p 'ubuntu' ssh -o StrictHostKeyChecking=no -o ConnectTimeout=5 ubuntu@"$vm_ip" 'echo "SSH connection successful"' 2>/dev/null; then
            log "✓ SSH connection successful to ubuntu@$vm_ip"
            return 0
        fi
        
        sleep 5
        attempt=$((attempt + 1))
    done
    
    error "✗ SSH connection failed after $max_attempts attempts"
    return 1
}

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    error "This script must be run as root for btrfs operations"
fi

# Parse arguments
if [ $# -ne 4 ]; then
    error "Usage: $0 <vm_number> <vm_name> <memory_mb> <cpu_cores>"
fi

VM_NUMBER="$1"
VM_NAME="$2"
MEMORY_MB="$3"
CPU_CORES="$4"

# Validate arguments
if ! [[ "$VM_NUMBER" =~ ^[0-9]+$ ]]; then
    error "VM number must be a number"
fi

if [ "$VM_NUMBER" -lt 1 ] || [ "$VM_NUMBER" -gt 200 ]; then
    error "VM number must be between 1 and 200"
fi

if ! [[ "$MEMORY_MB" =~ ^[0-9]+$ ]]; then
    error "Memory must be a number (in MB)"
fi

if ! [[ "$CPU_CORES" =~ ^[0-9]+$ ]]; then
    error "CPU cores must be a number"
fi

if [[ "$VM_NAME" =~ [^a-zA-Z0-9_-] ]]; then
    error "VM name can only contain alphanumeric characters, hyphens, and underscores"
fi

# Calculate static IP address based on VM number
VM_STATIC_IP="192.168.100.$VM_NUMBER"

log "Starting VM: $VM_NAME (number $VM_NUMBER) with ${MEMORY_MB}MB RAM and $CPU_CORES CPU cores"
log "VM will be assigned static IP: $VM_STATIC_IP"

# Comprehensive prerequisite checks
log "Performing prerequisite checks..."

# Check if cloud-hypervisor is available
test_command_exists "cloud-hypervisor" "Cloud Hypervisor installation"

# Check if qemu-img is available (for image conversion)
if ! command -v qemu-img &> /dev/null; then
    warn "qemu-img not found. Installing qemu-utils..."
    apt update && apt install -y qemu-utils
    test_command_exists "qemu-img" "QEMU tools installation"
fi

# Check for required tools
test_command_exists "curl" "curl installation"
test_command_exists "btrfs" "btrfs tools installation"
test_command_exists "ip" "iproute2 tools installation"

# Check if ethtool is available (for TAP interface optimization)
if ! command -v ethtool &> /dev/null; then
    log "Installing ethtool for network interface optimization..."
    apt update && apt install -y ethtool
    test_command_exists "ethtool" "ethtool installation"
fi

# Check if sshpass is available (for SSH testing)
if ! command -v sshpass &> /dev/null; then
    log "Installing sshpass for SSH testing..."
    apt update && apt install -y sshpass
    test_command_exists "sshpass" "sshpass installation"
fi

# Check if python3 is available (for password hashing)
test_command_exists "python3" "Python3 installation"

# Check if genisoimage or mkisofs is available
if ! command -v genisoimage &> /dev/null && ! command -v mkisofs &> /dev/null; then
    log "Installing genisoimage for cloud-init ISO creation..."
    apt update && apt install -y genisoimage
fi

log "✓ All prerequisites checked"

# Create base directory structure
log "Setting up storage structure..."
mkdir -p "$VM_BASE_DIR"
test_directory_exists "$VM_BASE_DIR" "VM base directory creation"

# Check if the base directory is on btrfs
FILESYSTEM_TYPE=$(stat -f -c %T "$VM_BASE_DIR" 2>/dev/null)
if [ "$FILESYSTEM_TYPE" != "btrfs" ]; then
    error "Base directory $VM_BASE_DIR is not on a btrfs filesystem (detected: $FILESYSTEM_TYPE). Please create a btrfs filesystem first."
fi

log "✓ Btrfs filesystem detected at $VM_BASE_DIR"

# Create base and vms subvolumes if they don't exist
if [ ! -d "$BASE_SUBVOL" ]; then
    log "Creating base subvolume at $BASE_SUBVOL"
    btrfs subvolume create "$BASE_SUBVOL"
    test_directory_exists "$BASE_SUBVOL" "Base subvolume creation"
else
    log "✓ Base subvolume already exists"
fi

if [ ! -d "$VMS_SUBVOL" ]; then
    log "Creating VMs subvolume at $VMS_SUBVOL"
    btrfs subvolume create "$VMS_SUBVOL"
    test_directory_exists "$VMS_SUBVOL" "VMs subvolume creation"
else
    log "✓ VMs subvolume already exists"
fi

# Verify subvolumes are properly created
test_step "Base subvolume verification" "btrfs subvolume show '$BASE_SUBVOL' &>/dev/null"
test_step "VMs subvolume verification" "btrfs subvolume show '$VMS_SUBVOL' &>/dev/null"

# Define paths
BASE_IMAGE_PATH="$BASE_SUBVOL/${BASE_IMAGE_NAME}.raw"
FIRMWARE_PATH="$BASE_SUBVOL/hypervisor-fw"
VM_SUBVOL_PATH="$VMS_SUBVOL/vm$VM_NUMBER"
VM_IMAGE_PATH="$VM_SUBVOL_PATH/vm$VM_NUMBER.raw"
CLOUD_INIT_PATH="$VM_SUBVOL_PATH/cloud-init.img"

# Download and prepare base image if it doesn't exist
log "Preparing base image and firmware..."
if [ ! -f "$BASE_IMAGE_PATH" ]; then
    log "Base image not found. Downloading Ubuntu cloud image..."
    
    # Download the qcow2 image
    TEMP_QCOW2="/tmp/${BASE_IMAGE_NAME}.img"
    if ! curl -L --fail --progress-bar -o "$TEMP_QCOW2" "$BASE_IMAGE_URL"; then
        error "Failed to download Ubuntu cloud image from $BASE_IMAGE_URL"
    fi
    test_file_exists "$TEMP_QCOW2" "Ubuntu cloud image download"
    
    log "Converting qcow2 image to raw format..."
    qemu-img convert -p -f qcow2 -O raw "$TEMP_QCOW2" "$BASE_IMAGE_PATH"
    test_file_exists "$BASE_IMAGE_PATH" "Base image conversion"
    
    # Verify the converted image
    image_info=$(qemu-img info "$BASE_IMAGE_PATH" 2>/dev/null)
    if echo "$image_info" | grep -q "file format: raw"; then
        log "✓ Base image successfully converted to raw format"
    else
        error "Base image conversion verification failed"
    fi
    
    # Cleanup temporary file
    rm -f "$TEMP_QCOW2"
    
    log "✓ Base image created at $BASE_IMAGE_PATH"
else
    log "✓ Base image already exists at $BASE_IMAGE_PATH"
    test_file_exists "$BASE_IMAGE_PATH" "Base image verification"
fi

# Download firmware if it doesn't exist
if [ ! -f "$FIRMWARE_PATH" ]; then
    log "Downloading Cloud Hypervisor firmware..."
    if ! curl -L --fail --progress-bar -o "$FIRMWARE_PATH" "$FIRMWARE_URL"; then
        error "Failed to download firmware from $FIRMWARE_URL"
    fi
    test_file_exists "$FIRMWARE_PATH" "Firmware download"
    chmod +x "$FIRMWARE_PATH"
    test_step "Firmware executable check" "[ -x '$FIRMWARE_PATH' ]"
    log "✓ Firmware downloaded to $FIRMWARE_PATH"
    else
        log "✓ Firmware already exists at $FIRMWARE_PATH"
        test_file_exists "$FIRMWARE_PATH" "Firmware verification"
    fi
    
    # Extract kernel and initrd from base image
    log "Extracting kernel and initrd for kernel boot..."
    EXTRACT_SCRIPT="$(dirname "$0")/extract_kernel.sh"
    if [ -f "$EXTRACT_SCRIPT" ]; then
        "$EXTRACT_SCRIPT"
    else
        warn "Kernel extraction script not found, attempting manual extraction..."
        # Fallback manual extraction
        KERNEL_PATH="$BASE_SUBVOL/vmlinuz"
        INITRD_PATH="$BASE_SUBVOL/initrd.img"
        if [ ! -f "$KERNEL_PATH" ] || [ ! -f "$INITRD_PATH" ]; then
            log "Extracting kernel and initrd manually..."
            TEMP_MOUNT=$(mktemp -d)
            losetup -P /dev/loop1 "$BASE_IMAGE_PATH"
            mount /dev/loop1p16 "$TEMP_MOUNT"
            cp "$TEMP_MOUNT/vmlinuz-6.8.0-60-generic" "$KERNEL_PATH" 2>/dev/null || true
            cp "$TEMP_MOUNT/initrd.img-6.8.0-60-generic" "$INITRD_PATH" 2>/dev/null || true
            umount "$TEMP_MOUNT"
            losetup -d /dev/loop1
            rmdir "$TEMP_MOUNT"
        fi
    fi
    
    test_file_exists "$BASE_SUBVOL/vmlinuz" "Kernel extraction"
    test_file_exists "$BASE_SUBVOL/initrd.img" "Initrd extraction"

# Create VM subvolume by cloning from base
log "Setting up VM-specific storage..."
if [ -d "$VM_SUBVOL_PATH" ]; then
    warn "VM subvolume $VM_NAME already exists. Removing it..."
    if btrfs subvolume show "$VM_SUBVOL_PATH" &>/dev/null; then
        btrfs subvolume delete "$VM_SUBVOL_PATH"
    else
        rm -rf "$VM_SUBVOL_PATH"
    fi
    test_step "VM subvolume cleanup" "[ ! -d '$VM_SUBVOL_PATH' ]"
fi

log "Creating VM subvolume by cloning base subvolume..."
btrfs subvolume snapshot "$BASE_SUBVOL" "$VM_SUBVOL_PATH"
test_directory_exists "$VM_SUBVOL_PATH" "VM subvolume creation"
test_step "VM subvolume verification" "btrfs subvolume show '$VM_SUBVOL_PATH' &>/dev/null"

# Copy the base image to VM subvolume (this will be a CoW copy initially)
log "Creating VM disk image (thin provisioned)..."
cp --reflink=always "$BASE_IMAGE_PATH" "$VM_IMAGE_PATH"
test_file_exists "$VM_IMAGE_PATH" "VM disk image creation"

# Verify the image copy
vm_image_size=$(stat -c%s "$VM_IMAGE_PATH" 2>/dev/null)
base_image_size=$(stat -c%s "$BASE_IMAGE_PATH" 2>/dev/null)
if [ "$vm_image_size" = "$base_image_size" ]; then
    log "✓ VM disk image successfully created (size: $vm_image_size bytes)"
else
    error "VM disk image size mismatch (VM: $vm_image_size, Base: $base_image_size)"
fi

# Create cloud-init image for first boot
log "Creating cloud-init configuration..."

# Generate a random MAC address for the VM (used in cloud-init and hypervisor)
VM_MAC="52:54:00:$(printf '%02x:%02x:%02x' $((RANDOM%256)) $((RANDOM%256)) $((RANDOM%256)))"
log "Generated MAC address for VM: $VM_MAC"

# Generate proper password hash for 'ubuntu'
PASSWORD_HASH=$(generate_password_hash)
test_step "Password hash generation" "[ -n '$PASSWORD_HASH' ]"

cat > "/tmp/user-data" << EOF
#cloud-config
users:
  - name: ubuntu
    sudo: ALL=(ALL) NOPASSWD:ALL
    shell: /bin/bash
    lock_passwd: false
    passwd: $PASSWORD_HASH
    groups: sudo
    home: /home/ubuntu

# Enable SSH with password authentication
ssh_pwauth: true
disable_root: false
chpasswd:
  expire: false

# SSH configuration
ssh_authorized_keys: []

# Network configuration with static IP
network:
  config: disabled

write_files:
  - path: /etc/netplan/50-cloud-init.yaml
    content: |
      network:
        version: 2
        ethernets:
          ens3:
            dhcp4: false
            addresses:
              - $VM_STATIC_IP/24 # Use variable for IP
            # gateway4: 192.168.100.1 # Deprecated
            routes: # Use modern routes syntax
              - to: default
                via: 192.168.100.1
            nameservers:
              addresses:
                - 8.8.8.8
                - 1.1.1.1

# Package updates and installs
package_update: true
package_upgrade: false
packages:
  - openssh-server
  - curl
  - wget
  - git
  - htop
  - vim
  - net-tools

# Ensure SSH service is enabled and started
# Also configure static IP as backup
runcmd:
  - systemctl enable ssh
  - systemctl start ssh
  - systemctl status ssh
  - netplan apply
  - chmod 0600 /etc/netplan/50-cloud-init.yaml || echo "Failed to chmod /etc/netplan/50-cloud-init.yaml"
  - sleep 2 # Allow netplan apply to settle
  #- ip addr flush dev ens3
  #- ip addr add $VM_STATIC_IP/24 dev ens3
  #- ip route add default via 192.168.100.1
  - ip addr show ens3
  - ip route show
  - ping 192.168.100.2 -c 3
  - echo "WERE THERE"


# Final message
final_message: "Cloud-init setup complete. VM is ready for SSH access!"
EOF

test_file_exists "/tmp/user-data" "Cloud-init user-data creation"

# Create meta-data file
cat > "/tmp/meta-data" << EOF
instance-id: $VM_NAME
local-hostname: $VM_NAME
EOF

test_file_exists "/tmp/meta-data" "Cloud-init meta-data creation"

# Create cloud-init ISO
log "Creating cloud-init ISO..."
if command -v genisoimage &> /dev/null; then
    genisoimage -output "$CLOUD_INIT_PATH" -volid cidata -joliet -rock /tmp/user-data /tmp/meta-data
elif command -v mkisofs &> /dev/null; then
    mkisofs -o "$CLOUD_INIT_PATH" -V cidata -J -r /tmp/user-data /tmp/meta-data
else
    error "Neither genisoimage nor mkisofs found. Please install genisoimage or cdrtools."
fi

test_file_exists "$CLOUD_INIT_PATH" "Cloud-init ISO creation"

# Verify the ISO was created properly
iso_size=$(stat -c%s "$CLOUD_INIT_PATH" 2>/dev/null)
if [ "$iso_size" -gt 0 ]; then
    log "✓ Cloud-init ISO created successfully (size: $iso_size bytes)"
else
    error "Cloud-init ISO creation failed or resulted in empty file"
fi

# Cleanup temporary files
rm -f /tmp/user-data /tmp/meta-data

log "✓ Cloud-init ISO created at $CLOUD_INIT_PATH"

# Resize the VM disk to give it more space (optional, expand to 20GB)
log "Resizing VM disk to 20GB..."
qemu-img resize "$VM_IMAGE_PATH" 20G

# Verify disk resize
new_size=$(qemu-img info "$VM_IMAGE_PATH" | grep "virtual size" | awk '{print $3}')
if echo "$new_size" | grep -q "20"; then
    log "✓ VM disk successfully resized to 20GB"
else
    warn "VM disk resize verification failed, but continuing..."
fi

# Create network configuration
TAP_NAME="tap-$VM_NAME"

log "Setting up network configuration..."

# Check if bridge exists, create if not
if ! ip link show "$BRIDGE_NAME" &>/dev/null; then
    log "Creating bridge interface $BRIDGE_NAME..."
    ip link add name "$BRIDGE_NAME" type bridge
    ip link set dev "$BRIDGE_NAME" up
    
    # Configure bridge with IP address for VM network
    log "Configuring bridge IP address..."
    ip addr add "$BRIDGE_IP" dev "$BRIDGE_NAME"
    
    test_network_interface "$BRIDGE_NAME" "Bridge interface creation"
    test_step "Bridge IP configuration" "ip addr show '$BRIDGE_NAME' | grep -q '192.168.100.1'"
else
    log "✓ Bridge interface $BRIDGE_NAME already exists"
    # Ensure bridge has IP configured
    if ! ip addr show "$BRIDGE_NAME" | grep -q "192.168.100.1"; then
        log "Adding IP address to existing bridge..."
        ip addr add "$BRIDGE_IP" dev "$BRIDGE_NAME" 2>/dev/null || true
    fi
fi

# Create TAP interface for the VM
TAP_NAME="tap-vm$VM_NUMBER"
log "Creating TAP interface $TAP_NAME..."

# Remove existing TAP interface if it exists
if ip link show "$TAP_NAME" &>/dev/null; then
    warn "TAP interface $TAP_NAME already exists, removing it..."
    ip link delete "$TAP_NAME" 2>/dev/null || true
    sleep 1  # Give time for cleanup
fi

# Create TAP interface with proper configuration for Cloud Hypervisor
ip tuntap add dev "$TAP_NAME" mode tap user root
test_network_interface "$TAP_NAME" "TAP interface creation"

# Set TAP interface up
ip link set dev "$TAP_NAME" up
sleep 1  # Give the interface a moment to come up
test_step "TAP interface up" "ip link show '$TAP_NAME' | grep -q 'UP'"

# Attach to bridge
ip link set dev "$TAP_NAME" master "$BRIDGE_NAME"
sleep 1  # Give the bridge attachment a moment to complete
test_step "TAP interface bridge attachment" "ip link show '$TAP_NAME' | grep -q 'master'"

# Disable offloading features that can cause issues with Cloud Hypervisor
ethtool -K "$TAP_NAME" tx off rx off tso off gso off gro off lro off 2>/dev/null || warn "Could not disable TAP interface offloading (ethtool not available)"

log "✓ Network interfaces configured successfully"

# Ensure basic networking is set up (simplified version of setup_vm_network.sh)
log "Ensuring basic VM networking is configured..."

# Enable IP forwarding
echo 1 > /proc/sys/net/ipv4/ip_forward

# Set up basic NAT rules (remove existing first to avoid duplicates)
iptables -t nat -D POSTROUTING -s "$NETWORK" -j MASQUERADE 2>/dev/null || true
iptables -D FORWARD -i "$BRIDGE_NAME" -j ACCEPT 2>/dev/null || true
iptables -D FORWARD -o "$BRIDGE_NAME" -j ACCEPT 2>/dev/null || true

# Add new rules
iptables -t nat -A POSTROUTING -s "$NETWORK" -j MASQUERADE
iptables -A FORWARD -i "$BRIDGE_NAME" -j ACCEPT
iptables -A FORWARD -o "$BRIDGE_NAME" -j ACCEPT

log "✓ Basic NAT and forwarding rules configured"

# Check if dnsmasq is running for DHCP
if ! systemctl is-active --quiet dnsmasq 2>/dev/null; then
    warn "dnsmasq is not running. VMs may not get IP addresses automatically."
    warn "Consider running: sudo ./setup_vm_network.sh"
fi

# Start the VM with Cloud Hypervisor
log "Starting VM $VM_NAME..."
VM_SOCKET="/tmp/cloud-hypervisor-vm$VM_NUMBER.sock"
VM_LOG_FILE="/tmp/cloud-hypervisor-vm$VM_NUMBER.log"

# Remove existing socket and log file if they exist
rm -f "$VM_SOCKET" "$VM_LOG_FILE"


# Start Cloud Hypervisor in background with error handling
log "Launching Cloud Hypervisor..."

# Try to start Cloud Hypervisor and capture any error output
log "Starting Cloud Hypervisor with kernel boot:"
log "cloud-hypervisor --memory size=${MEMORY_MB}M --cpus boot=$CPU_CORES --kernel $KERNEL_PATH --initramfs $INITRD_PATH --cmdline 'root=LABEL=cloudimg-rootfs ro console=tty1 console=ttyS0' --disk path=$VM_IMAGE_PATH path=$CLOUD_INIT_PATH,readonly=on --net tap=$TAP_NAME,mac=$VM_MAC --serial file=$VM_LOG_FILE --console off --event-monitor path=${VM_LOG_FILE}.events"

# Use kernel boot instead of firmware boot to properly pass root device
KERNEL_PATH="$BASE_SUBVOL/vmlinuz"
INITRD_PATH="$BASE_SUBVOL/initrd.img"

cloud-hypervisor \
    --api-socket "$VM_SOCKET" \
    --memory "size=${MEMORY_MB}M" \
    --cpus "boot=$CPU_CORES" \
    --kernel "$KERNEL_PATH" \
    --initramfs "$INITRD_PATH" \
    --cmdline "root=LABEL=cloudimg-rootfs ro console=tty1 console=ttyS0" \
    --disk "path=$VM_IMAGE_PATH" "path=$CLOUD_INIT_PATH,readonly=on" \
    --net "tap=$TAP_NAME,mac=$VM_MAC" \
    --serial tty \
    --console off \
    --event-monitor "path=${VM_LOG_FILE}.events" &

VM_PID=$!

# Check if the process started successfully
if [ -z "$VM_PID" ]; then
    error "Failed to get VM process ID"
fi

# Verify VM process started
sleep 2
if ! test_process_running "$VM_PID" "VM process startup"; then
    error "VM process failed to start or died immediately. Check log: $VM_LOG_FILE"
fi

log "✓ VM $VM_NAME started successfully with PID $VM_PID"
log "VM socket: $VM_SOCKET"
log "VM log file: $VM_LOG_FILE"
log "TAP interface: $TAP_NAME"
log "Bridge interface: $BRIDGE_NAME"
log "VM MAC address: $VM_MAC"

# Wait for VM to initialize and check if it's running properly
log "Waiting for VM to initialize..."
init_wait_count=0
while [ $init_wait_count -lt 10 ]; do
    sleep 1
    init_wait_count=$((init_wait_count + 1))
    
    # Check if VM process is still running
    if ! kill -0 "$VM_PID" 2>/dev/null; then
        error "VM process died during initialization. Check log: $VM_LOG_FILE"
    fi
    
    if [ $((init_wait_count % 3)) -eq 0 ]; then
        info "VM initializing... ($init_wait_count/10 seconds)"
    fi
done

log "✓ VM initialization completed"

# Save VM information for management
log "Saving VM configuration..."
VM_INFO_FILE="$VM_SUBVOL_PATH/vm-info.txt"
cat > "$VM_INFO_FILE" << EOF
VM_NUMBER=$VM_NUMBER
VM_NAME=$VM_NAME
VM_PID=$VM_PID
VM_SOCKET=$VM_SOCKET
TAP_NAME=$TAP_NAME
BRIDGE_NAME=$BRIDGE_NAME
MEMORY_MB=$MEMORY_MB
CPU_CORES=$CPU_CORES
VM_IMAGE_PATH=$VM_IMAGE_PATH
CLOUD_INIT_PATH=$CLOUD_INIT_PATH
VM_MAC=$VM_MAC
VM_LOG_FILE=$VM_LOG_FILE
VM_STATIC_IP=$VM_STATIC_IP
STARTED="$(date '+%Y-%m-%d %H:%M:%S')"
EOF

test_file_exists "$VM_INFO_FILE" "VM info file creation"
log "✓ VM information saved to $VM_INFO_FILE"

# Function to cleanup on exit (only for interactive mode)
cleanup_on_exit() {
    log "Cleaning up VM $VM_NAME..."
    if [ -n "$VM_PID" ] && kill -0 "$VM_PID" 2>/dev/null; then
        kill "$VM_PID"
        wait "$VM_PID" 2>/dev/null
    fi
    ip link delete "$TAP_NAME" 2>/dev/null || true
    rm -f "$VM_SOCKET"
}

# Test VM boot and connectivity
log "Testing VM boot and connectivity..."

# Show the static IP that will be used
log "VM $VM_NAME will use static IP address: $VM_STATIC_IP"

# Wait for VM to boot and verify static IP
VM_IP=$(wait_for_vm_boot "$VM_NAME" "$VM_STATIC_IP")
if [ $? -ne 0 ] || [ -z "$VM_IP" ]; then
    error "VM failed to boot or respond at static IP $VM_STATIC_IP. Check log: $VM_LOG_FILE"
fi

log "✓ VM booted successfully and is using IP: $VM_IP"

# Test IP connectivity first
log "Testing IP connectivity before SSH..."
if test_ip_connectivity "$VM_IP"; then
    log "✓ IP connectivity test passed for $VM_IP"
else
    error "IP connectivity test failed for $VM_IP"
fi

# Test SSH connectivity
if test_ssh_connection "$VM_IP"; then
    log "🎉 SUCCESS: VM $VM_NAME is fully operational!"
    log "✓ VM is running with PID $VM_PID"
    log "✓ VM has IP address: $VM_IP"
    log "✓ SSH is working: ssh ubuntu@$VM_IP (password: ubuntu)"
    log "✓ VM info saved to: $VM_INFO_FILE"
    
    # Display relevant IP configuration lines from VM log
    if [ -f "$VM_LOG_FILE" ]; then
        info "Relevant IP configuration from VM log ($VM_LOG_FILE):"
        grep -E "inet .*ens3|default via" "$VM_LOG_FILE" | tail -n 5 || true
    fi
    
    echo ""
    info "VM $VM_NAME is ready for use!"
    info "Connect via SSH: ssh ubuntu@$VM_IP"
    info "Default password: ubuntu (please change after first login)"
    info "To stop the VM: sudo $(dirname "$0")/ubuntu_vm_delete.sh $VM_NUMBER"
    echo ""
    
    # Don't set trap for successful VMs - let them run
    log "VM $VM_NAME will continue running in the background."
    log "Use 'sudo $(dirname "$0")/ubuntu_vm_delete.sh $VM_NUMBER' to stop and delete it."
    
else
    error "SSH connectivity test failed. VM will be deleted for retry."
fi

# If we reach here, the VM is working properly
log "VM startup and testing completed successfully!"