#!/bin/bash

# ======================================================
# 🚀 Automated VirtualBox VM Management Script (Updated)
# ======================================================

# --- Configuration ---
BASE_VM_NAME="DB_Image"         # Name of the template/base VM to clone from
VM_USER="asf_user"              # User inside the VM (for IP detection and final output)
VM_PASSWORD="ASF"               # Password for the VM user (for IP detection)

# Ports Configuration
BASE_SSH_PORT=6002              # Starting SSH NAT port
BASE_VNC_PORT=5002              # Starting VNC (VRDE) port
PORT_DB_FILE="$HOME/.vms_ports_db.json" # JSON file for port persistence

# VM Resources (50 GB is 51200 MB)
VM_MEMORY=51200                 
VM_VRAM=128
VM_CPUS=2

# --- Utility Functions (All functions remain the same) ---

usage() {
    echo "Usage:"
    echo "  $0 create <new_vm_name>"
    echo "  $0 status <vm_name>"
    echo "  $0 activate <vm_name>"
    echo "  $0 delete <vm_name>"
    exit 1
}

# Function to remove VM port entry from JSON
remove_port_entry() {
    local vm_name="$1"
    init_port_db # Ensure DB exists
    
    if [[ -f "$PORT_DB_FILE" ]]; then
        if grep -q "\"$vm_name\"" "$PORT_DB_FILE"; then
            # Use grep and sed to remove the entry, handling list structure complexities
            grep -v "\"$VM_NAME\"" "$PORT_DB_FILE" | sed 's/}, /}/' | sed 's/, }/}/' > "${PORT_DB_FILE}.tmp"
            
            if grep -q "^{.*}$" "${PORT_DB_FILE}.tmp"; then
                mv "${PORT_DB_FILE}.tmp" "$PORT_DB_FILE"
            else
                echo "{}" > "$PORT_DB_FILE"
            fi
            # NOTE: Removed the success echo to keep the main flow cleaner
        fi
    fi
}

# Function to safely create/initialize the JSON port database
init_port_db() {
    if [[ ! -f "$PORT_DB_FILE" ]]; then
        echo "{}" > "$PORT_DB_FILE"
    fi
}

# Function to get VM status
get_vm_status() {
    local vm_name="$1"
    local status
    status=$(VBoxManage showvminfo "$vm_name" --machinereadable 2>/dev/null | grep -i "^VMState=" | cut -d'"' -f2)

    if [[ "$status" == "running" ]]; then
        echo "active"
    else
        echo "not active"
    fi
}

# Function to start VM if not running
activate_vm() {
    local vm_name="$1"
    local status
    status=$(get_vm_status "$vm_name")
    if [[ "$status" == "active" ]]; then
        echo "✅ VM '$vm_name' is already active."
        return 0
    fi

    echo "🚀 Starting VM '$vm_name'..."
    VBoxManage startvm "$vm_name" --type headless
    if [[ $? -eq 0 ]]; then
        echo "✅ VM '$vm_name' started successfully."
    else
        echo "❌ Failed to start VM '$vm_name'."
        return 1
    fi
}

# Function to read ports from JSON
get_vm_ports() {
    local vm_name="$1"
    local ports
    
    # Check if the VM entry exists in the JSON file
    ports=$(cat "$PORT_DB_FILE" | grep "\"$vm_name\"" | awk -F': ' '{print $2}' | tr -d '",}{')
    
    if [[ -z "$ports" ]]; then
        echo "" # Return empty if not found
    else
        # Assuming the format in the JSON is {"vm_name": "SSH_PORT VNC_PORT"}
        echo "$ports"
    fi
}

# Function to assign unique ports for SSH/VNC per VM using JSON
assign_ports() {
    local vm_name="$1"
    init_port_db

    # 1. Check if ports already assigned in JSON
    local assigned_ports
    assigned_ports=$(get_vm_ports "$vm_name")

    if [[ -n "$assigned_ports" ]]; then
        echo "$assigned_ports"
        return 0
    fi

    # 2. Collect ALL currently used ports from the JSON file
    local used_ports
    used_ports=$(grep -oE "[0-9]+" "$PORT_DB_FILE" 2>/dev/null)
    
    local ssh_port=$BASE_SSH_PORT
    local vnc_port=$BASE_VNC_PORT

    # Find unique SSH port
    while grep -qw "$ssh_port" <<<"$used_ports" || netstat -tln 2>/dev/null | grep -q ":$ssh_port"; do
        ssh_port=$((ssh_port + 1))
    done

    # Find unique VNC port
    while grep -qw "$vnc_port" <<<"$used_ports" || netstat -tln 2>/dev/null | grep -q ":$vnc_port"; do
        vnc_port=$((vnc_port + 1))
    done

    # 3. Save new ports to JSON using a simple tool for JSON manipulation
    local new_entry="\"$vm_name\": \"$ssh_port $vnc_port\""
    
    local current_content
    current_content=$(cat "$PORT_DB_FILE")

    if [[ "$current_content" == "{}" ]]; then
        echo "{$new_entry}" > "$PORT_DB_FILE"
    else
        sed '$s/}$/, '"$new_entry"'}' "$PORT_DB_FILE" > "${PORT_DB_FILE}.tmp" && mv "${PORT_DB_FILE}.tmp" "$PORT_DB_FILE"
    fi

    echo "$ssh_port $vnc_port"
}

# --- Main Logic ---

if [[ $# -lt 2 && "$1" != "create" ]]; then
    usage
fi

COMMAND="$1"
ARG="$2"

case "$COMMAND" in
    status)
        get_vm_status "$ARG"
        ;;

    activate)
        activate_vm "$ARG"
        ;;

    create)
        NEW_VM_NAME="$ARG"
        init_port_db # Ensure DB exists

        if VBoxManage showvminfo "$NEW_VM_NAME" &>/dev/null; then
            echo "VM '$NEW_VM_NAME' already exists. Please delete it or choose another name."
            exit 1
        fi

        # Assign ports and store them in the JSON
        read -r SSH_PORT VNC_PORT < <(assign_ports "$NEW_VM_NAME")
        
        echo "Cloning VM '$BASE_VM_NAME' to '$NEW_VM_NAME' (SSH:$SSH_PORT, VNC:$VNC_PORT)..."
        VBoxManage clonevm "$BASE_VM_NAME" --name "$NEW_VM_NAME" --register --mode all
        
        if [[ $? -ne 0 ]]; then
            echo "❌ Failed to clone VM. Aborting."
            remove_port_entry "$NEW_VM_NAME"
            exit 1
        fi

        # Configure VM resources
        echo "Configuring VM resources (Memory: $VM_MEMORY MB, CPUs: $VM_CPUS)..."
        VBoxManage modifyvm "$NEW_VM_NAME" --memory "$VM_MEMORY" --vram "$VM_VRAM" --cpus "$VM_CPUS"
        VBoxManage modifyvm "$NEW_VM_NAME" --nic1 nat
        VBoxManage modifyvm "$NEW_VM_NAME" --cableconnected1 on
        
        # --- PHASE 1: Initial Boot and Shutdown ---
        
        echo "Starting VM '$NEW_VM_NAME' for initial boot (1 minute)..."
        VBoxManage startvm "$NEW_VM_NAME" --type headless
        
        # Wait for 1 minute (60 seconds) for first boot/OS initialization
        sleep 60
        echo "Initial boot complete. Shutting down VM."
        
        VBoxManage controlvm "$NEW_VM_NAME" poweroff 
        
        # --- PHASE 2: Setting Network Ports ---
        echo "Setting up network port forwarding (SSH:$SSH_PORT, VNC:$VNC_PORT)..."
        
        # Remove old (potentially conflicting) rules and add the new one
        VBoxManage modifyvm "$NEW_VM_NAME" --natpf1 delete ssh 2>/dev/null || true
        
        # Add the new NAT port forwarding rules
        VBoxManage modifyvm "$NEW_VM_NAME" --natpf1 "ssh,tcp,,${SSH_PORT},,22"
        VBoxManage modifyvm "$NEW_VM_NAME" --vrde on --vrdeport "$VNC_PORT"

        # --- PHASE 3: Final Start and Wait ---
        echo "Starting VM '$NEW_VM_NAME' headless for final usage (Waiting 5 minutes)..."
        VBoxManage startvm "$NEW_VM_NAME" --type headless

        # Wait for 5 minutes (300 seconds) for the system and services to fully initialize
        sleep 300
        echo "Initialization complete. VM should be ready."

        # --- PHASE 4: Final Output ---
        
        echo "--- ✅ VM $NEW_VM_NAME Created and Ready ---"
        echo "Host: 127.0.0.1 (use your external DDNS for remote access)"
        echo "User: $VM_USER"
        echo "SSH Port: $SSH_PORT"
        echo "VNC Port: $VNC_PORT"
        
        ;;

    delete)
        VM_NAME="$ARG"
        init_port_db # Ensure DB exists
        
        echo "🛑 Powering off VM '$VM_NAME'..."
        VBoxManage controlvm "$VM_NAME" poweroff 2>/dev/null || true

        echo "🗑️ Unregistering and deleting VM '$VM_NAME'..."
        VBoxManage unregistervm "$VM_NAME" --delete
        
        # Remove entry from port DB
        if [[ -f "$PORT_DB_FILE" ]]; then
            
            if grep -q "\"$VM_NAME\"" "$PORT_DB_FILE"; then
                # Remove the entire line containing the VM entry
                grep -v "\"$VM_NAME\"" "$PORT_DB_FILE" | sed 's/}, /}/' | sed 's/, }/}/' > "${PORT_DB_FILE}.tmp"
                
                # Check if the file is now just "{}" and fix it
                if grep -q "^{.*}$" "${PORT_DB_FILE}.tmp"; then
                    mv "${PORT_DB_FILE}.tmp" "$PORT_DB_FILE"
                else
                    echo "{}" > "$PORT_DB_FILE" 
                fi
                
                echo "✅ Removed '$VM_NAME' entry from $PORT_DB_FILE"
            fi
        fi
        ;;

    *)
        usage
        ;;
esac