#!/bin/bash

# Ubuntu VM Management Script
# Usage: ubuntu_vm_manage.sh <command> [vm_number]

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
VM_BASE_DIR="/var/lib/vms"
VMS_SUBVOL="$VM_BASE_DIR/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}"
}

show_usage() {
    echo "Ubuntu VM Management Script"
    echo ""
    echo "Usage: $0 <command> [vm_number]"
    echo ""
    echo "Commands:"
    echo "  list                    - List all VMs and their status"
    echo "  status <vm_number>      - Show detailed status of a specific VM"
    echo "  console <vm_number>     - Connect to VM console (serial)"
    echo "  ssh <vm_number>         - SSH to VM (requires network setup)"
    echo "  stop <vm_number>        - Stop a running VM"
    echo "  start <vm_number>       - Start a stopped VM"
    echo "  delete <vm_number>      - Delete a VM completely"
    echo "  ip <vm_number>          - Show VM IP address"
    echo "  logs <vm_number>        - Show VM logs"
    echo ""
    echo "Examples:"
    echo "  $0 list"
    echo "  $0 status 1"
    echo "  $0 console 1"
    echo "  $0 ssh 1"
}

list_vms() {
    log "Listing all VMs..."
    echo ""
    printf "%-8s %-15s %-15s %-10s %-8s %-15s %-20s\n" "VM #" "IP ADDRESS" "VM NAME" "STATUS" "PID" "MEMORY" "STARTED"
    printf "%-8s %-15s %-15s %-10s %-8s %-15s %-20s\n" "----" "----------" "-------" "------" "---" "------" "-------"
    
    if [ ! -d "$VMS_SUBVOL" ]; then
        warn "No VMs directory found at $VMS_SUBVOL"
        return
    fi
    
    for vm_dir in "$VMS_SUBVOL"/vm*; do
        if [ -d "$vm_dir" ]; then
            vm_dirname=$(basename "$vm_dir")
            # Extract number from vm directory name (vm1, vm2, etc.)
            vm_number=${vm_dirname#vm}
            if [[ "$vm_number" =~ ^[0-9]+$ ]]; then
                vm_ip="192.168.100.$vm_number"
                vm_info_file="$vm_dir/vm-info.txt"
                
                if [ -f "$vm_info_file" ]; then
                    # Safely source the file with error handling
                    if source "$vm_info_file" 2>/dev/null; then
                        # Check if VM is running
                        if [ -n "$VM_PID" ] && kill -0 "$VM_PID" 2>/dev/null; then
                            status="${GREEN}RUNNING${NC}"
                            pid="$VM_PID"
                        else
                            status="${RED}STOPPED${NC}"
                            pid="N/A"
                        fi
                        
                        # Handle missing or malformed STARTED field
                        if [ -z "$STARTED" ]; then
                            STARTED="Unknown"
                        fi
                        
                        # Use VM_NAME from config, fallback to vm number
                        display_name="${VM_NAME:-vm$vm_number}"
                        
                        printf "%-8s %-15s %-15s %-18s %-8s %-15s %-20s\n" "$vm_number" "$vm_ip" "$display_name" "$(echo -e "$status")" "$pid" "${MEMORY_MB}MB" "$STARTED"
                    else
                        printf "%-8s %-15s %-15s %-18s %-8s %-15s %-20s\n" "$vm_number" "$vm_ip" "vm$vm_number" "$(printf "%b" "${RED}ERROR${NC}")" "N/A" "N/A" "Config Error"
                    fi
                else
                    printf "%-8s %-15s %-15s %-18s %-8s %-15s %-20s\n" "$vm_number" "$vm_ip" "vm$vm_number" "$(printf "%b" "${YELLOW}UNKNOWN${NC}")" "N/A" "N/A" "N/A"
                fi
            fi
        fi
    done
    echo ""
}

show_vm_status() {
    local vm_number="$1"
    local vm_dir="$VMS_SUBVOL/vm$vm_number"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    # Validate VM number
    if ! [[ "$vm_number" =~ ^[0-9]+$ ]]; then
        error "VM number must be a number"
    fi
    
    if [ ! -d "$vm_dir" ]; then
        error "VM number '$vm_number' not found"
    fi
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM info file not found for VM number '$vm_number'"
    fi
    
    source "$vm_info_file"
    
    log "VM Status for VM number: $vm_number"
    echo ""
    echo "VM Number:      $vm_number"
    echo "VM Name:        ${VM_NAME:-vm$vm_number}"
    echo "Static IP:      192.168.100.$vm_number"
    echo "Memory:         ${MEMORY_MB}MB"
    echo "CPU Cores:      $CPU_CORES"
    echo "Started:        $STARTED"
    echo "Image Path:     $VM_IMAGE_PATH"
    echo "Cloud-init:     $CLOUD_INIT_PATH"
    echo "Socket:         $VM_SOCKET"
    echo "TAP Interface:  $TAP_NAME"
    echo "Bridge:         $BRIDGE_NAME"
    
    # Check if VM is running
    if [ -n "$VM_PID" ] && kill -0 "$VM_PID" 2>/dev/null; then
        echo -e "Status:         ${GREEN}RUNNING${NC} (PID: $VM_PID)"
        
        # Show network info
        if ip link show "$TAP_NAME" &>/dev/null; then
            echo -e "Network:        ${GREEN}TAP interface active${NC}"
        else
            echo -e "Network:        ${RED}TAP interface not found${NC}"
        fi
        
        # Show socket info
        if [ -S "$VM_SOCKET" ]; then
            echo -e "API Socket:     ${GREEN}Available${NC}"
        else
            echo -e "API Socket:     ${RED}Not available${NC}"
        fi
    else
        echo -e "Status:         ${RED}STOPPED${NC}"
    fi
    echo ""
}

connect_console() {
    local vm_number="$1"
    local vm_dir="$VMS_SUBVOL/vm$vm_number"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    # Validate VM number
    if ! [[ "$vm_number" =~ ^[0-9]+$ ]]; then
        error "VM number must be a number"
    fi
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM number '$vm_number' not found"
    fi
    
    source "$vm_info_file"
    
    if [ -z "$VM_PID" ] || ! kill -0 "$VM_PID" 2>/dev/null; then
        error "VM number '$vm_number' is not running"
    fi
    
    info "Connecting to console for VM number '$vm_number'"
    info "Press Ctrl+A then X to exit console"
    echo ""
    
    # Connect to the VM's console via the API socket
    if [ -S "$VM_SOCKET" ]; then
        # Use socat to connect to the console
        if command -v socat &>/dev/null; then
            socat - UNIX-CONNECT:"$VM_SOCKET"
        else
            warn "socat not found. Installing..."
            apt update && apt install -y socat
            socat - UNIX-CONNECT:"$VM_SOCKET"
        fi
    else
        error "VM socket not found at $VM_SOCKET"
    fi
}

ssh_to_vm() {
    local vm_number="$1"
    local vm_ip="192.168.100.$vm_number"
    
    # Validate VM number
    if ! [[ "$vm_number" =~ ^[0-9]+$ ]]; then
        error "VM number must be a number"
    fi
    
    info "Attempting SSH connection to VM number $vm_number at $vm_ip"
    info "Default login: ubuntu / ubuntu"
    echo ""
    
    # Try to SSH directly
    if command -v sshpass &>/dev/null; then
        info "Using sshpass for automatic login..."
        sshpass -p 'ubuntu' ssh -o StrictHostKeyChecking=no ubuntu@"$vm_ip"
    else
        info "Manual SSH connection (enter password 'ubuntu'):"
        ssh -o StrictHostKeyChecking=no ubuntu@"$vm_ip"
    fi
}

stop_vm() {
    local vm_number="$1"
    local vm_dir="$VMS_SUBVOL/vm$vm_number"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    # Validate VM number
    if ! [[ "$vm_number" =~ ^[0-9]+$ ]]; then
        error "VM number must be a number"
    fi
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM number '$vm_number' not found"
    fi
    
    source "$vm_info_file"
    
    if [ -z "$VM_PID" ] || ! kill -0 "$VM_PID" 2>/dev/null; then
        warn "VM number '$vm_number' is not running"
        return
    fi
    
    log "Stopping VM number '$vm_number' (PID: $VM_PID)..."
    
    # Graceful shutdown first
    kill -TERM "$VM_PID" 2>/dev/null || true
    
    # Wait a bit for graceful shutdown
    sleep 3
    
    # Force kill if still running
    if kill -0 "$VM_PID" 2>/dev/null; then
        warn "Forcing shutdown..."
        kill -KILL "$VM_PID" 2>/dev/null || true
    fi
    
    # Cleanup network
    if [ -n "$TAP_NAME" ]; then
        ip link delete "$TAP_NAME" 2>/dev/null || true
    fi
    
    # Remove socket
    if [ -n "$VM_SOCKET" ]; then
        rm -f "$VM_SOCKET"
    fi
    
    log "VM number '$vm_number' stopped"
}

delete_vm() {
    local vm_number="$1"
    local vm_dir="$VMS_SUBVOL/vm$vm_number"
    
    # Validate VM number
    if ! [[ "$vm_number" =~ ^[0-9]+$ ]]; then
        error "VM number must be a number"
    fi
    
    if [ ! -d "$vm_dir" ]; then
        error "VM number '$vm_number' not found"
    fi
    
    # Stop VM first if running
    if [ -f "$vm_dir/vm-info.txt" ]; then
        source "$vm_dir/vm-info.txt"
        if [ -n "$VM_PID" ] && kill -0 "$VM_PID" 2>/dev/null; then
            log "Stopping VM before deletion..."
            stop_vm "$vm_number"
        fi
    fi
    
    # Confirm deletion
    echo -e "${RED}WARNING: This will permanently delete VM number '$vm_number' and all its data!${NC}"
    read -p "Are you sure? (yes/no): " confirm
    
    if [ "$confirm" = "yes" ]; then
        log "Deleting VM number '$vm_number'..."
        btrfs subvolume delete "$vm_dir"
        log "VM number '$vm_number' deleted successfully"
    else
        info "Deletion cancelled"
    fi
}

start_vm() {
    local vm_number="$1"
    local vm_dir="$VMS_SUBVOL/vm$vm_number"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    # Validate VM number
    if ! [[ "$vm_number" =~ ^[0-9]+$ ]]; then
        error "VM number must be a number"
    fi
    
    if [ ! -d "$vm_dir" ]; then
        error "VM number '$vm_number' not found"
    fi
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM info file not found for VM number '$vm_number'"
    fi
    
    source "$vm_info_file"
    
    # Check if VM is already running
    if [ -n "$VM_PID" ] && kill -0 "$VM_PID" 2>/dev/null; then
        warn "VM number '$vm_number' is already running (PID: $VM_PID)"
        return
    fi
    
    log "Starting VM number '$vm_number'..."
    
    # Create TAP interface
    if ! ip link show "$TAP_NAME" &>/dev/null; then
        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"
    fi
    
    # Remove existing socket if it exists
    rm -f "$VM_SOCKET"
    
    # Start Cloud Hypervisor in background with kernel boot
    KERNEL_PATH="/var/lib/vms/base/vmlinuz"
    INITRD_PATH="/var/lib/vms/base/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=52:54:00:$(printf '%02x:%02x:%02x' $((RANDOM%256)) $((RANDOM%256)) $((RANDOM%256)))" \
        --serial tty \
        --console off \
        --log-file /tmp/cloud-hypervisor-$VM_NAME.log &
    
    NEW_VM_PID=$!
    
    # Update VM info file with new PID
    sed -i "s/VM_PID=.*/VM_PID=$NEW_VM_PID/" "$vm_info_file"
    
    log "VM number '$vm_number' started with PID $NEW_VM_PID"
    log "VM socket: $VM_SOCKET"
    log "TAP interface: $TAP_NAME"
    log "To connect to console: ./ubuntu_vm_manage.sh console $vm_number"
}

# Main script logic
if [ $# -eq 0 ]; then
    show_usage
    exit 1
fi

COMMAND="$1"
VM_NUMBER="$2"

case "$COMMAND" in
    "list")
        list_vms
        ;;
    "status")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for status command"
        fi
        show_vm_status "$VM_NUMBER"
        ;;
    "console")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for console command"
        fi
        connect_console "$VM_NUMBER"
        ;;
    "ssh")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for ssh command"
        fi
        ssh_to_vm "$VM_NUMBER"
        ;;
    "stop")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for stop command"
        fi
        stop_vm "$VM_NUMBER"
        ;;
    "start")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for start command"
        fi
        start_vm "$VM_NUMBER"
        ;;
    "delete")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for delete command"
        fi
        delete_vm "$VM_NUMBER"
        ;;
    "ip")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for ip command"
        fi
        echo "VM $VM_NUMBER IP address: 192.168.100.$VM_NUMBER"
        ;;
    "logs")
        if [ -z "$VM_NUMBER" ]; then
            error "VM number required for logs command"
        fi
        log_file="/tmp/cloud-hypervisor-vm$VM_NUMBER.log"
        if [ -f "$log_file" ]; then
            log "Showing logs for VM number $VM_NUMBER:"
            tail -f "$log_file"
        else
            warn "Log file not found: $log_file"
        fi
        ;;
    *)
        error "Unknown command: $COMMAND"
        show_usage
        exit 1
        ;;
esac