#!/bin/bash

# Ubuntu VM Start Script with Cloud Hypervisor and Btrfs Thin Provisioning
# Usage: ubuntu_vm_start.sh $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"

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}"
    exit 1
}

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

# 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 3 ]; then
    error "Usage: $0 <vm_name> <memory_mb> <cpu_cores>"
fi

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

# Validate arguments
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

log "Starting VM: $VM_NAME with ${MEMORY_MB}MB RAM and $CPU_CORES CPU cores"

# Check if cloud-hypervisor is available
if ! command -v cloud-hypervisor &> /dev/null; then
    error "cloud-hypervisor not found. Please install it first."
fi

# 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
fi

# Create base directory structure
mkdir -p "$VM_BASE_DIR"

# Check if the base directory is on btrfs
if ! btrfs filesystem show "$VM_BASE_DIR" &>/dev/null; then
    error "Base directory $VM_BASE_DIR is not on a btrfs filesystem. Please create a btrfs filesystem first."
fi

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

if [ ! -d "$VMS_SUBVOL" ]; then
    log "Creating VMs subvolume at $VMS_SUBVOL"
    btrfs subvolume create "$VMS_SUBVOL"
fi

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

# Download and prepare base image if it doesn't exist
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
    
    log "Converting qcow2 image to raw format..."
    qemu-img convert -p -f qcow2 -O raw "$TEMP_QCOW2" "$BASE_IMAGE_PATH"
    
    # Cleanup temporary file
    rm -f "$TEMP_QCOW2"
    
    log "Base image created at $BASE_IMAGE_PATH"
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
    chmod +x "$FIRMWARE_PATH"
    log "Firmware downloaded to $FIRMWARE_PATH"
fi

# Create VM subvolume by cloning from base
if [ -d "$VM_SUBVOL_PATH" ]; then
    warn "VM subvolume $VM_NAME already exists. Removing it..."
    btrfs subvolume delete "$VM_SUBVOL_PATH"
fi

log "Creating VM subvolume by cloning base subvolume..."
btrfs subvolume snapshot "$BASE_SUBVOL" "$VM_SUBVOL_PATH"

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

# Create cloud-init image for first boot
log "Creating cloud-init configuration..."
cat > "/tmp/user-data" << EOF
#cloud-config
users:
  - name: ubuntu
    sudo: ALL=(ALL) NOPASSWD:ALL
    shell: /bin/bash
    lock_passwd: false
    passwd: \$6\$rounds=4096\$saltsalt\$L9.LKkHxeed8Kn9.Kk8nNWn8W.XhHPyjKJJXYqKoTFJJy7P8dMCFK
    ssh_authorized_keys:
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC... # Add your SSH public key here
    groups: sudo
    home: /home/ubuntu

ssh_pwauth: true
disable_root: false
chpasswd:
  expire: false

# Enable SSH
ssh_authorized_keys: []

# Package updates and installs
package_update: true
package_upgrade: true
packages:
  - curl
  - wget
  - git
  - htop
  - vim

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

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

# 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

# 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

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

# 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
    # You may want to configure the bridge with an IP address
    # ip addr add 192.168.100.1/24 dev "$BRIDGE_NAME"
fi

# Create TAP interface for the VM
log "Creating TAP interface $TAP_NAME..."
ip tuntap add dev "$TAP_NAME" mode tap
ip link set dev "$TAP_NAME" up
ip link set dev "$TAP_NAME" master "$BRIDGE_NAME"

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

# Remove existing socket if it exists
rm -f "$VM_SOCKET"

# Start Cloud Hypervisor in background
cloud-hypervisor \
    --api-socket "$VM_SOCKET" \
    --memory "size=${MEMORY_MB}M" \
    --cpus "boot=$CPU_CORES" \
    --kernel "$FIRMWARE_PATH" \
    --disk "path=$VM_IMAGE_PATH" \
    --disk "path=$CLOUD_INIT_PATH,readonly=on" \
    --net "tap=$TAP_NAME,mac=52:54:00:$(printf '%02x:%02x:%02x' $((RANDOM%256)) $((RANDOM%256)) $((RANDOM%256)))" \
    --serial tty \
    --console off \
    --log-level info &

VM_PID=$!

log "VM $VM_NAME started with PID $VM_PID"
log "VM socket: $VM_SOCKET"
log "TAP interface: $TAP_NAME"
log "Bridge interface: $BRIDGE_NAME"

# Save VM information for management
VM_INFO_FILE="$VM_SUBVOL_PATH/vm-info.txt"
cat > "$VM_INFO_FILE" << EOF
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
STARTED=$(date)
EOF

log "VM information saved to $VM_INFO_FILE"

# Function to cleanup on exit
cleanup() {
    log "Cleaning up VM $VM_NAME..."
    if 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"
}

# Set trap for cleanup on script exit
trap cleanup EXIT INT TERM

log "VM $VM_NAME is running. Press Ctrl+C to stop."
log "To connect via SSH (once VM is booted): ssh ubuntu@<vm-ip>"
log "Default password: ubuntu (change after first login)"

# Wait for the VM process
wait "$VM_PID"