#!/bin/bash

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

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_name]"
    echo ""
    echo "Commands:"
    echo "  list                    - List all VMs and their status"
    echo "  status <vm_name>        - Show detailed status of a specific VM"
    echo "  console <vm_name>       - Connect to VM console (serial)"
    echo "  ssh <vm_name>           - SSH to VM (requires network setup)"
    echo "  stop <vm_name>          - Stop a running VM"
    echo "  start <vm_name>         - Start a stopped VM"
    echo "  delete <vm_name>        - Delete a VM completely"
    echo "  ip <vm_name>            - Show VM IP address"
    echo "  logs <vm_name>          - Show VM logs"
    echo ""
    echo "Examples:"
    echo "  $0 list"
    echo "  $0 status test-vm"
    echo "  $0 console test-vm"
    echo "  $0 ssh test-vm"
}

list_vms() {
    log "Listing all VMs..."
    echo ""
    printf "%-15s %-10s %-8s %-15s %-20s\n" "VM NAME" "STATUS" "PID" "MEMORY" "STARTED"
    printf "%-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"/*; do
        if [ -d "$vm_dir" ]; then
            vm_name=$(basename "$vm_dir")
            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
                    
                    printf "%-15s %-18s %-8s %-15s %-20s\n" "$vm_name" "$(printf "%b" "$status")" "$pid" "${MEMORY_MB}MB" "$STARTED"
                else
                    printf "%-15s %-18s %-8s %-15s %-20s\n" "$vm_name" "$(printf "%b" "${RED}ERROR${NC}")" "N/A" "N/A" "Config Error"
                fi
            else
                printf "%-15s %-18s %-8s %-15s %-20s\n" "$vm_name" "$(printf "%b" "${YELLOW}UNKNOWN${NC}")" "N/A" "N/A" "N/A"
            fi
        fi
    done
    echo ""
}

show_vm_status() {
    local vm_name="$1"
    local vm_dir="$VMS_SUBVOL/$vm_name"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    if [ ! -d "$vm_dir" ]; then
        error "VM '$vm_name' not found"
    fi
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM info file not found for '$vm_name'"
    fi
    
    source "$vm_info_file"
    
    log "VM Status for: $vm_name"
    echo ""
    echo "VM Name:        $VM_NAME"
    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 "Status:         ${GREEN}RUNNING${NC} (PID: $VM_PID)"
        
        # Show network info
        if ip link show "$TAP_NAME" &>/dev/null; then
            echo "Network:        ${GREEN}TAP interface active${NC}"
        else
            echo "Network:        ${RED}TAP interface not found${NC}"
        fi
        
        # Show socket info
        if [ -S "$VM_SOCKET" ]; then
            echo "API Socket:     ${GREEN}Available${NC}"
        else
            echo "API Socket:     ${RED}Not available${NC}"
        fi
    else
        echo "Status:         ${RED}STOPPED${NC}"
    fi
    echo ""
}

connect_console() {
    local vm_name="$1"
    local vm_dir="$VMS_SUBVOL/$vm_name"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM '$vm_name' not found"
    fi
    
    source "$vm_info_file"
    
    if [ -z "$VM_PID" ] || ! kill -0 "$VM_PID" 2>/dev/null; then
        error "VM '$vm_name' is not running"
    fi
    
    info "Connecting to console for VM '$vm_name'"
    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_name="$1"
    
    warn "SSH connection requires:"
    warn "1. VM to be fully booted"
    warn "2. Network bridge configured with IP"
    warn "3. VM to have received IP via DHCP"
    echo ""
    info "To set up networking:"
    info "1. Configure bridge IP: ip addr add 192.168.100.1/24 dev br0"
    info "2. Enable IP forwarding: echo 1 > /proc/sys/net/ipv4/ip_forward"
    info "3. Set up NAT: iptables -t nat -A POSTROUTING -s 192.168.100.0/24 -j MASQUERADE"
    echo ""
    info "Default login: ubuntu / ubuntu"
    info "SSH command: ssh ubuntu@<vm-ip>"
}

stop_vm() {
    local vm_name="$1"
    local vm_dir="$VMS_SUBVOL/$vm_name"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM '$vm_name' not found"
    fi
    
    source "$vm_info_file"
    
    if [ -z "$VM_PID" ] || ! kill -0 "$VM_PID" 2>/dev/null; then
        warn "VM '$vm_name' is not running"
        return
    fi
    
    log "Stopping VM '$vm_name' (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 '$vm_name' stopped"
}

delete_vm() {
    local vm_name="$1"
    local vm_dir="$VMS_SUBVOL/$vm_name"
    
    if [ ! -d "$vm_dir" ]; then
        error "VM '$vm_name' 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_name"
        fi
    fi
    
    # Confirm deletion
    echo -e "${RED}WARNING: This will permanently delete VM '$vm_name' and all its data!${NC}"
    read -p "Are you sure? (yes/no): " confirm
    
    if [ "$confirm" = "yes" ]; then
        log "Deleting VM '$vm_name'..."
        btrfs subvolume delete "$vm_dir"
        log "VM '$vm_name' deleted successfully"
    else
        info "Deletion cancelled"
    fi
}

start_vm() {
    local vm_name="$1"
    local vm_dir="$VMS_SUBVOL/$vm_name"
    local vm_info_file="$vm_dir/vm-info.txt"
    
    if [ ! -d "$vm_dir" ]; then
        error "VM '$vm_name' not found"
    fi
    
    if [ ! -f "$vm_info_file" ]; then
        error "VM info file not found for '$vm_name'"
    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 '$vm_name' is already running (PID: $VM_PID)"
        return
    fi
    
    log "Starting VM '$vm_name'..."
    
    # 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
    cloud-hypervisor \
        --api-socket "$VM_SOCKET" \
        --memory "size=${MEMORY_MB}M" \
        --cpus "boot=$CPU_CORES" \
        --kernel "/var/lib/vms/base/hypervisor-fw" \
        --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 '$vm_name' 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_name"
}

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

COMMAND="$1"
VM_NAME="$2"

case "$COMMAND" in
    "list")
        list_vms
        ;;
    "status")
        if [ -z "$VM_NAME" ]; then
            error "VM name required for status command"
        fi
        show_vm_status "$VM_NAME"
        ;;
    "console")
        if [ -z "$VM_NAME" ]; then
            error "VM name required for console command"
        fi
        connect_console "$VM_NAME"
        ;;
    "ssh")
        if [ -z "$VM_NAME" ]; then
            error "VM name required for ssh command"
        fi
        ssh_to_vm "$VM_NAME"
        ;;
    "stop")
        if [ -z "$VM_NAME" ]; then
            error "VM name required for stop command"
        fi
        stop_vm "$VM_NAME"
        ;;
    "start")
        if [ -z "$VM_NAME" ]; then
            error "VM name required for start command"
        fi
        start_vm "$VM_NAME"
        ;;
    "delete")
        if [ -z "$VM_NAME" ]; then
            error "VM name required for delete command"
        fi
        delete_vm "$VM_NAME"
        ;;
    "ip"|"logs")
        warn "Command '$COMMAND' not yet implemented"
        ;;
    *)
        error "Unknown command: $COMMAND"
        show_usage
        exit 1
        ;;
esac