#!/bin/bash

# Cloud Hypervisor VM with Bridge Networking and Multi-Instance Support
# Usage: ./run-vm.sh [start|stop|restart|status|config|init]

set -e

# Default Configuration (can be overridden by .chconfig)
VM_NAME="myvm"
MEMORY_SIZE="8192M"
CPU_COUNT="4"
DISK_COUNT="1"
DISK_1="10G:vm-disk.img"
BRIDGE_NAME="zosbr"
VM_IP="192.168.1.100"
VM_MASK="255.255.255.0"
KERNEL_PATH="output/vmlinuz.efi"
INITRD_PATH="output/initrd.img"
CONSOLE_ARGS="console=ttyS0,115200n8"

# Runtime configuration
INSTANCE_ID=""
TAP_INTERFACE=""
VM_MAC=""
PID_FILE=""
CONFIG_FILE=".chconfig"
DISK_PATHS=()

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

log() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
    exit 1
}

info() {
    echo -e "${BLUE}[CONFIG]${NC} $1"
}

generate_instance_id() {
    # 8-character instance ID for files and collision resistance
    INSTANCE_ID=$(head /dev/urandom | tr -dc a-z0-9 | head -c 8)
}

generate_mac() {
    # Generate random MAC with VMware OUI prefix
    VM_MAC=$(printf "52:54:00:%02x:%02x:%02x" $((RANDOM%256)) $((RANDOM%256)) $((RANDOM%256)))
}

load_config() {
    if [ -f "$CONFIG_FILE" ]; then
        log "Loading configuration from $CONFIG_FILE"
        
        # Source the config file
        source "$CONFIG_FILE"
        
        # Validate required settings
        [ -z "$VM_NAME" ] && error "VM_NAME not set in config"
        [ -z "$MEMORY_SIZE" ] && error "MEMORY_SIZE not set in config"
        [ -z "$CPU_COUNT" ] && error "CPU_COUNT not set in config"
        [ -z "$DISK_COUNT" ] && DISK_COUNT=0
        
        # Validate numeric values
        if ! [[ "$CPU_COUNT" =~ ^[0-9]+$ ]]; then
            error "CPU_COUNT must be a number"
        fi
        
        if ! [[ "$DISK_COUNT" =~ ^[0-9]+$ ]]; then
            error "DISK_COUNT must be a number"
        fi
        
        info "VM Name: $VM_NAME"
        info "Memory: $MEMORY_SIZE"
        info "CPUs: $CPU_COUNT"
        info "Disk Count: $DISK_COUNT"
        info "Bridge: $BRIDGE_NAME"
        info "Kernel: $KERNEL_PATH"
        if [ -f "$INITRD_PATH" ]; then
            info "Initrd: $INITRD_PATH"
        fi
    else
        warn "No config file found ($CONFIG_FILE), using defaults"
    fi
    
    # Generate instance-specific configuration
    generate_instance_id
    generate_mac
    
    # TAP interface naming - stay under 15 char limit
    # Format: ch-XXXXXXXX (11 chars max) - safe and collision-resistant
    TAP_INTERFACE="ch-${INSTANCE_ID}"
    
    PID_FILE="/tmp/${VM_NAME}-${INSTANCE_ID}.pid"
    
    info "Instance ID: $INSTANCE_ID"
    info "TAP Interface: $TAP_INTERFACE (${#TAP_INTERFACE} chars)"
    info "MAC Address: $VM_MAC"
    
    # Verify TAP name length (safety check)
    if [ ${#TAP_INTERFACE} -gt 15 ]; then
        error "TAP interface name too long: $TAP_INTERFACE (${#TAP_INTERFACE} chars, max 15)"
    fi
    
    # Process disk configuration
    process_disk_config
}

process_disk_config() {
    DISK_PATHS=()
    
    for ((i=1; i<=DISK_COUNT; i++)); do
        disk_var="DISK_$i"
        disk_config="${!disk_var}"
        
        if [ -z "$disk_config" ]; then
            warn "DISK_$i not defined, skipping"
            continue
        fi
        
        # Parse "size:name" format
        disk_size="${disk_config%%:*}"
        disk_name="${disk_config#*:}"
        
        # Auto-generate name if empty
        if [ -z "$disk_name" ]; then
            disk_name="${VM_NAME}-${INSTANCE_ID}-disk${i}.img"
        else
            # Add instance ID to prevent conflicts
            disk_name="${VM_NAME}-${INSTANCE_ID}-${disk_name}"
        fi
        
        DISK_PATHS+=("$disk_size:$disk_name")
        info "Disk $i: $disk_size -> $disk_name"
    done
}

create_config() {
    if [ -f "$CONFIG_FILE" ]; then
        read -p "Config file $CONFIG_FILE already exists. Overwrite? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log "Keeping existing config file"
            return
        fi
    fi
    
    log "Creating default config file: $CONFIG_FILE"
    
    cat > "$CONFIG_FILE" << 'EOF'
# Cloud Hypervisor VM Configuration
# Lines starting with # are comments

# Basic VM configuration
VM_NAME=myvm
MEMORY_SIZE=8192M
CPU_COUNT=4

# Disk configuration (array format)
# Format: "size:name" - if name is empty, auto-generated
# Set DISK_COUNT=0 for no disks
DISK_COUNT=2
DISK_1="10G:data.img"
DISK_2="5G:swap.img"
# DISK_3="20G:"  # Auto-generated name
# DISK_4="1G:temp.img"

# Network configuration
BRIDGE_NAME=zosbr
VM_IP=192.168.1.100
VM_MASK=255.255.255.0

# Paths
KERNEL_PATH=output/vmlinuz.efi
INITRD_PATH=output/initrd.img

# Console
CONSOLE_ARGS=console=ttyS0,115200n8
EOF
    
    log "Config file created. Edit $CONFIG_FILE and run again."
    log ""
    log "Example configurations:"
    log "  Small VM:  MEMORY_SIZE=2G, CPU_COUNT=1, DISK_COUNT=1"
    log "  Large VM:  MEMORY_SIZE=16G, CPU_COUNT=8, DISK_COUNT=3"
    log "  No disks:  DISK_COUNT=0"
}

check_requirements() {
    log "Checking requirements..."
    
    command -v cloud-hypervisor >/dev/null 2>&1 || error "cloud-hypervisor not found"
    command -v ip >/dev/null 2>&1 || error "ip command not found"
    
    if [ "$DISK_COUNT" -gt 0 ]; then
        command -v qemu-img >/dev/null 2>&1 || error "qemu-img not found (needed for disk creation)"
    fi
    
    [ -f "$KERNEL_PATH" ] || error "Kernel file not found: $KERNEL_PATH"
    
    if [ "$EUID" -ne 0 ]; then
        error "This script must be run as root (for network setup)"
    fi
}

create_disks() {
    if [ "$DISK_COUNT" -eq 0 ]; then
        log "No disks configured"
        return
    fi
    
    log "Creating disk images..."
    
    for disk_config in "${DISK_PATHS[@]}"; do
        disk_size="${disk_config%%:*}"
        disk_path="${disk_config#*:}"
        
        if [ ! -f "$disk_path" ]; then
            log "Creating disk: $disk_path ($disk_size)"
            qemu-img create -f raw "$disk_path" "$disk_size"
        else
            log "Disk already exists: $disk_path"
            
            # Show disk info
            DISK_INFO=$(qemu-img info "$disk_path" 2>/dev/null | grep "virtual size" || echo "Unable to get size")
            info "  $DISK_INFO"
        fi
    done
}

setup_bridge() {
    log "Setting up bridge network: $BRIDGE_NAME"
    
    # Create bridge if it doesn't exist
    if ! ip link show "$BRIDGE_NAME" >/dev/null 2>&1; then
        log "Creating bridge: $BRIDGE_NAME"
        ip link add name "$BRIDGE_NAME" type bridge
        ip link set "$BRIDGE_NAME" up
        
        # Configure bridge IP (adjust as needed)
        BRIDGE_IP=$(echo "$VM_IP" | sed 's/\.[0-9]*$/.1/')
        ip addr add "${BRIDGE_IP}/24" dev "$BRIDGE_NAME" 2>/dev/null || true
        info "Bridge IP: $BRIDGE_IP"
    else
        log "Bridge $BRIDGE_NAME already exists"
    fi
}

setup_tap() {
    log "Setting up TAP interface: $TAP_INTERFACE"
    
    # Remove existing TAP if it exists (very unlikely with 8-char random IDs)
    if ip link show "$TAP_INTERFACE" >/dev/null 2>&1; then
        warn "TAP interface $TAP_INTERFACE already exists, removing..."
        ip link delete "$TAP_INTERFACE" 2>/dev/null || true
    fi
    
    # Create TAP interface
    ip tuntap add "$TAP_INTERFACE" mode tap
    ip link set "$TAP_INTERFACE" up
    
    # Add TAP to bridge
    ip link set "$TAP_INTERFACE" master "$BRIDGE_NAME"
    
    log "TAP interface $TAP_INTERFACE added to bridge $BRIDGE_NAME"
}

cleanup_tap() {
    if [ -n "$TAP_INTERFACE" ]; then
        log "Cleaning up TAP interface: $TAP_INTERFACE"
        
        if ip link show "$TAP_INTERFACE" >/dev/null 2>&1; then
            ip link delete "$TAP_INTERFACE"
            log "TAP interface $TAP_INTERFACE removed"
        fi
    fi
}

cleanup_disks() {
    if [ "$1" = "--remove-disks" ]; then
        log "Removing disk images..."
        for disk_config in "${DISK_PATHS[@]}"; do
            disk_path="${disk_config#*:}"
            if [ -f "$disk_path" ]; then
                rm -f "$disk_path"
                log "Removed: $disk_path"
            fi
        done
    fi
}

start_vm() {
    log "Starting VM: $VM_NAME"
    
    load_config
    check_requirements
    create_disks
    setup_bridge
    setup_tap
    
    # Build cloud-hypervisor command
    CH_CMD="cloud-hypervisor"
    CH_CMD="$CH_CMD --kernel $KERNEL_PATH"
    CH_CMD="$CH_CMD --memory size=$MEMORY_SIZE"
    CH_CMD="$CH_CMD --cpus boot=$CPU_COUNT"
    CH_CMD="$CH_CMD --net tap=$TAP_INTERFACE,mac=$VM_MAC"
    CH_CMD="$CH_CMD --serial tty"
    CH_CMD="$CH_CMD --console off"
    CH_CMD="$CH_CMD --cmdline '$CONSOLE_ARGS'"
    
    # Add initrd if it exists
    if [ -f "$INITRD_PATH" ]; then
        CH_CMD="$CH_CMD --initramfs $INITRD_PATH"
        log "Using initrd: $INITRD_PATH"
    fi
    
    # Add disks
    for disk_config in "${DISK_PATHS[@]}"; do
        disk_path="${disk_config#*:}"
        CH_CMD="$CH_CMD --disk path=$disk_path"
        log "Attached disk: $disk_path"
    done
    
    log "Starting cloud-hypervisor..."
    log "Command: $CH_CMD"
    
    # Save instance info
    cat > "${PID_FILE}.info" << EOF
VM_NAME=$VM_NAME
INSTANCE_ID=$INSTANCE_ID
TAP_INTERFACE=$TAP_INTERFACE
VM_MAC=$VM_MAC
MEMORY_SIZE=$MEMORY_SIZE
CPU_COUNT=$CPU_COUNT
DISK_COUNT=$DISK_COUNT
BRIDGE_NAME=$BRIDGE_NAME
EOF
    
    # Start VM in background and save PID
    eval "$CH_CMD" &
    VM_PID=$!
    echo $VM_PID > "$PID_FILE"
    
    log "VM started with PID: $VM_PID"
    log "Instance ID: $INSTANCE_ID"
    log "TAP Interface: $TAP_INTERFACE"
    log "Connect to serial console (Ctrl+C to exit)"
    
    # Wait for VM process
    wait $VM_PID
    log "VM process ended"
    
    # Cleanup on exit
    cleanup_tap
    rm -f "$PID_FILE" "${PID_FILE}.info"
}

stop_vm() {
    local instance_pattern="${1:-}"
    local remove_disks=""
    
    if [ "$1" = "--remove-disks" ] || [ "$2" = "--remove-disks" ]; then
        remove_disks="--remove-disks"
        if [ "$1" = "--remove-disks" ]; then
            instance_pattern=""
        fi
    fi
    
    # If no specific instance, load config and stop current
    if [ -z "$instance_pattern" ]; then
        load_config
        log "Stopping VM: $VM_NAME (Instance: $INSTANCE_ID)"
        
        if [ -f "$PID_FILE" ]; then
            VM_PID=$(cat "$PID_FILE")
            if kill -0 "$VM_PID" 2>/dev/null; then
                log "Terminating VM process: $VM_PID"
                kill "$VM_PID"
                sleep 2
                
                # Force kill if still running
                if kill -0 "$VM_PID" 2>/dev/null; then
                    warn "Force killing VM process: $VM_PID"
                    kill -9 "$VM_PID"
                fi
            fi
            
            # Load TAP interface from info file
            if [ -f "${PID_FILE}.info" ]; then
                source "${PID_FILE}.info"
            fi
            
            cleanup_tap
            
            if [ -n "$remove_disks" ]; then
                cleanup_disks --remove-disks
            fi
            
            rm -f "$PID_FILE" "${PID_FILE}.info"
        else
            warn "No PID file found"
        fi
    else
        # Stop specific instance by pattern
        log "Stopping VM instances matching: $instance_pattern"
        
        for pid_file in /tmp/*${instance_pattern}*.pid; do
            if [ -f "$pid_file" ]; then
                VM_PID=$(cat "$pid_file")
                info_file="${pid_file}.info"
                
                if [ -f "$info_file" ]; then
                    source "$info_file"
                    log "Stopping $VM_NAME (Instance: $INSTANCE_ID)"
                fi
                
                if kill -0 "$VM_PID" 2>/dev/null; then
                    kill "$VM_PID"
                    sleep 1
                    if kill -0 "$VM_PID" 2>/dev/null; then
                        kill -9 "$VM_PID"
                    fi
                fi
                
                # Cleanup TAP if info available
                if [ -n "$TAP_INTERFACE" ] && ip link show "$TAP_INTERFACE" >/dev/null 2>&1; then
                    ip link delete "$TAP_INTERFACE"
                fi
                
                rm -f "$pid_file" "$info_file"
            fi
        done
    fi
    
    log "VM stopped"
}

status_vm() {
    local show_all="${1:-}"
    
    if [ "$show_all" = "--all" ]; then
        log "All running VM instances:"
        echo
        
        local found=0
        for pid_file in /tmp/*.pid; do
            # Look for files matching our naming pattern
            if [ -f "$pid_file" ] && [[ "$(basename "$pid_file")" =~ -[a-z0-9]{8}\.pid$ ]]; then
                VM_PID=$(cat "$pid_file" 2>/dev/null)
                info_file="${pid_file}.info"
                
                if kill -0 "$VM_PID" 2>/dev/null && [ -f "$info_file" ]; then
                    source "$info_file"
                    echo "  VM: $VM_NAME"
                    echo "    Instance ID: $INSTANCE_ID"
                    echo "    PID: $VM_PID"
                    echo "    Memory: $MEMORY_SIZE"
                    echo "    CPUs: $CPU_COUNT"
                    echo "    Disks: $DISK_COUNT"
                    echo "    TAP: $TAP_INTERFACE (${#TAP_INTERFACE} chars)"
                    echo "    Bridge: $BRIDGE_NAME"
                    echo
                    found=1
                fi
            fi
        done
        
        if [ $found -eq 0 ]; then
            log "No running VM instances found"
        fi
    else
        load_config
        if [ -f "$PID_FILE" ]; then
            VM_PID=$(cat "$PID_FILE")
            if kill -0 "$VM_PID" 2>/dev/null; then
                log "VM '$VM_NAME' (Instance: $INSTANCE_ID) is running with PID: $VM_PID"
                log "TAP Interface: $TAP_INTERFACE (${#TAP_INTERFACE} chars)"
                return 0
            else
                warn "VM PID file exists but process is not running"
                rm -f "$PID_FILE" "${PID_FILE}.info"
            fi
        fi
        
        log "VM '$VM_NAME' is not running"
        return 1
    fi
}

show_config() {
    load_config
    echo
    echo "Current Configuration:"
    echo "====================="
    echo "VM Name:      $VM_NAME"
    echo "Instance ID:  $INSTANCE_ID"
    echo "Memory:       $MEMORY_SIZE"
    echo "CPUs:         $CPU_COUNT"
    echo "Disk Count:   $DISK_COUNT"
    
    if [ "$DISK_COUNT" -gt 0 ]; then
        echo "Disks:"
        for i in "${!DISK_PATHS[@]}"; do
            disk_config="${DISK_PATHS[$i]}"
            disk_size="${disk_config%%:*}"
            disk_path="${disk_config#*:}"
            echo "  $((i+1)): $disk_size -> $disk_path"
        done
    fi
    
    echo "Bridge:       $BRIDGE_NAME"
    echo "TAP:          $TAP_INTERFACE (${#TAP_INTERFACE} chars)"
    echo "VM MAC:       $VM_MAC"
    echo "VM IP:        $VM_IP"
    echo "Kernel:       $KERNEL_PATH"
    echo "Initrd:       $INITRD_PATH"
    echo "Console:      $CONSOLE_ARGS"
    echo "PID File:     $PID_FILE"
    echo
}

# Trap to cleanup on script exit
trap 'cleanup_tap 2>/dev/null; rm -f "$PID_FILE" "${PID_FILE}.info" 2>/dev/null' EXIT INT TERM

case "${1:-start}" in
    start)
        if [ -f "$CONFIG_FILE" ]; then
            load_config
        fi
        start_vm
        ;;
    stop)
        if [ "$2" = "--all" ]; then
            stop_vm "*" "$3"
        else
            stop_vm "$2" "$3"
        fi
        ;;
    restart)
        stop_vm
        sleep 1
        start_vm
        ;;
    status)
        status_vm "$2"
        ;;
    config)
        show_config
        ;;
    init|create-config)
        create_config
        ;;
    *)
        echo "Usage: $0 {start|stop|restart|status|config|init}"
        echo ""
        echo "Commands:"
        echo "  start                    - Start the VM"
        echo "  stop [pattern]           - Stop VM(s) [matching pattern]"
        echo "  stop --all               - Stop all running VMs"
        echo "  stop --remove-disks      - Stop and remove disk files"
        echo "  restart                  - Restart the VM"
        echo "  status                   - Show VM status"
        echo "  status --all             - Show all running VMs"
        echo "  config                   - Show current configuration"
        echo "  init                     - Create default .chconfig file"
        echo ""
        echo "Examples:"
        echo "  $0 stop myvm             - Stop VMs with 'myvm' in name"
        echo "  $0 stop --all            - Stop all running VMs"
        echo "  $0 status --all          - List all running VMs"
        echo ""
        echo "Configuration file: $CONFIG_FILE"
        if [ -f "$CONFIG_FILE" ]; then
            VM_NAME_TEMP=$(grep "^VM_NAME=" "$CONFIG_FILE" 2>/dev/null | cut -d= -f2 || echo "unknown")
            MEM_SIZE_TEMP=$(grep "^MEMORY_SIZE=" "$CONFIG_FILE" 2>/dev/null | cut -d= -f2 || echo "unknown")
            DISK_COUNT_TEMP=$(grep "^DISK_COUNT=" "$CONFIG_FILE" 2>/dev/null | cut -d= -f2 || echo "unknown")
            echo "Current VM: $VM_NAME_TEMP (Memory: $MEM_SIZE_TEMP, Disks: $DISK_COUNT_TEMP)"
        else
            echo "No config file found. Run '$0 init' to create one."
        fi
        exit 1
        ;;
esac