#!/usr/bin/env bash
## v1.00 (c) 2025 th@bogus.net
## vmfupg.sh - parallels helper tool; updates (linux) virtual machines based on their package manager
## ngl, had help from sourcegraph!

VMS=()
SELECTED_VMS=()

if [[ -t 1 ]] && command -v tput >/dev/null 2>&1 && tput colors >/dev/null 2>&1 && [[ $(tput colors) -ge 8 ]]; then
    COLOR_GREEN=$(tput setaf 2)
    COLOR_ORANGE=$(tput setaf 3)
    COLOR_RED=$(tput setaf 1)
    COLOR_RESET=$(tput sgr0)
else
    COLOR_GREEN=""
    COLOR_ORANGE=""
    COLOR_RED=""
    COLOR_RESET=""
fi

get_status_color() {
    local status="$1"
    case "$status" in
        running)
            echo "$COLOR_GREEN"
            ;;
        suspended)
            echo "$COLOR_ORANGE"
            ;;
        stopped)
            echo "$COLOR_RED"
            ;;
        *)
            echo ""
            ;;
    esac
}

get_vm_list() {
    prlctl list -a | tail -n +2 | while IFS= read -r line; do
        if [[ -n "$line" ]]; then
            uuid=$(echo "$line" | awk '{print $1}' | tr -d '{}')
            status=$(echo "$line" | awk '{print $2}')
            name=$(echo "$line" | awk '{for(i=4;i<=NF;i++) printf "%s ", $i; print ""}' | sed 's/[[:space:]]*$//')
            echo "$uuid|$status|$name"
        fi
    done
}

to_lower() {
    echo "$1" | tr '[:upper:]' '[:lower:]'
}

find_vms() {
    local vm_data
    vm_data=$(get_vm_list)
    
    if [[ -n "$VM_UUID" ]]; then
        echo "$vm_data" | grep "^$VM_UUID|"
    elif [[ -n "$VM_NAME" ]]; then
        echo "$vm_data" | grep -F "|$VM_NAME"
    elif [[ -n "$VM_GLOB" ]]; then
        local glob_lower
        glob_lower=$(to_lower "$VM_GLOB")
        echo "$vm_data" | while IFS='|' read -r uuid status name; do
            local name_lower
            name_lower=$(to_lower "$name")
            if [[ $name_lower == $glob_lower ]]; then
                echo "$uuid|$status|$name"
            fi
        done
    else
        echo "Error: No VM selection criteria provided. Use -n, -i, or -m option." >&2
        exit 1
    fi
}

detect_package_manager() {
    local vm_id="$1"
    
    if prlctl exec "$vm_id" 'command -v apt' &>/dev/null; then
        echo "apt"
    elif prlctl exec "$vm_id" 'command -v yum' &>/dev/null; then
        echo "yum"
    elif prlctl exec "$vm_id" 'command -v dnf' &>/dev/null; then
        echo "dnf"
    else
        echo "unknown"
    fi
}

update_vm() {
    local vm_id="$1"
    local pkg_mgr
    local exit_code
    
    echo "[+] Detecting package manager..."
    pkg_mgr=$(detect_package_manager "$vm_id")
    
    case "$pkg_mgr" in
        apt)
            echo "[+] Using apt package manager"
            prlctl exec "$vm_id" 'apt -y update && apt -y upgrade && apt -y autoclean && apt -y autoremove' 2>&1 | sed 's/^/   [>] /'
            exit_code=${PIPESTATUS[0]}
            ;;
        yum)
            echo "[+] Using yum package manager"
            prlctl exec "$vm_id" 'yum -y update && yum -y clean all' 2>&1 | sed 's/^/   [>] /'
            exit_code=${PIPESTATUS[0]}
            ;;
        dnf)
            echo "[+] Using dnf package manager"
            prlctl exec "$vm_id" 'dnf -y update && dnf -y clean all' 2>&1 | sed 's/^/   [>] /'
            exit_code=${PIPESTATUS[0]}
            ;;
        *)
            echo "[-] Warning: Unknown package manager for VM $vm_id. Skipping updates."
            return 1
            ;;
    esac
    
    return $exit_code
}

check_parallels_tools_installation() {
    local vm_id="$1"
    
    if prlctl exec "$vm_id" 'pgrep -f "zenity.*Parallels Tools Installation"' &>/dev/null; then
        return 0  
    else
        return 1 
    fi
}

check_parallels_tools_restart_prompt() {
    local vm_id="$1"
    
    if prlctl exec "$vm_id" 'pgrep -f "zenity.*--ok-label=Restart"' &>/dev/null; then
        return 0  
    else
        return 1 
    fi
}

wait_for_parallels_tools() {
    local vm_id="$1"
    local original_status="$2"
    
    if check_parallels_tools_installation "$vm_id"; then
        echo "[i] Parallels Agent Installation is requesting input"
        echo "[+] Waiting for Parallels Tools installation to complete..."
        
        while check_parallels_tools_installation "$vm_id"; do
            sleep 5
        done
        
        echo "[+] Parallels Tools installation completed"
    fi
    
    # Check if restart prompt is showing
    if check_parallels_tools_restart_prompt "$vm_id"; then
        echo "[i] Parallels Tools installation requires restart"
        echo "[+] Initiating restart..."
        
        # Kill the zenity process to trigger restart or handle it gracefully
        prlctl exec "$vm_id" 'pkill -f "zenity.*--ok-label=Restart"' &>/dev/null
        
        # Wait a moment for the restart to begin
        sleep 5
        
        # Wait for VM to shut down
        echo "[+] Waiting for VM to restart..."
        wait_for_state "$vm_id" "stopped"
        
        # Start VM again if original state was running, otherwise restore original state
        if [[ "$original_status" == "running" ]]; then
            echo "[+] Restarting VM..."
            if prlctl start "$vm_id" &>/dev/null; then
                echo "[+] VM restarted successfully"
                wait_for_vm_basic "$vm_id"  # Use basic wait to avoid recursion
            else
                echo "[-] Failed to restart VM"
                return 1
            fi
        else
            echo "[+] VM restarted, will restore to original state ($original_status) after updates"
        fi
    fi
}

wait_for_vm_basic() {
    local vm_id="$1"
    echo "[+] Waiting for VM to be ready..."
    local attempts=0
    local max_attempts=60
    
    until prlctl exec "$vm_id" uptime &>/dev/null; do
        sleep 5
        ((attempts++))
        if [[ $attempts -ge $max_attempts ]]; then
            echo "[-] Error: VM $vm_id failed to become ready after $((max_attempts * 5)) seconds"
            return 1
        fi
    done
    echo "[+] VM is ready!"
}

wait_for_vm() {
    local vm_id="$1"
    local original_status="$2"
    
    echo "[+] Waiting for VM to boot and be ready..."
    local attempts=0
    local max_attempts=60
    
    until prlctl exec "$vm_id" uptime &>/dev/null; do
        sleep 5
        ((attempts++))
        if [[ $attempts -ge $max_attempts ]]; then
            echo "[-] Error: VM $vm_id failed to become ready after $((max_attempts * 5)) seconds"
            return 1
        fi
    done
    echo "[+] VM is ready!"
    
    # Check for Parallels Tools installation after VM is ready
    wait_for_parallels_tools "$vm_id" "$original_status"
}

wait_for_state() {
    local vm_id="$1"
    local target_state="$2"
    
    echo "[+] Waiting for VM $vm_id to reach state: $target_state..."
    while true; do
        local current_state
        current_state=$(prlctl list --all 2>/dev/null | awk -v id="{$vm_id}" '$1 == id {print $2}')
        if [[ "$current_state" == "$target_state" ]]; then
            break
        fi
        sleep 2
    done
    echo "[+] VM reached target state: $target_state"
}

while getopts ":i:n:m:h" opt; do
    case $opt in
        n)
            VM_NAME="$OPTARG"
            ;;
        i)
            VM_UUID="$OPTARG"
            ;;
        m)
            VM_GLOB="$OPTARG"
            ;;
        h)
            echo "Usage: $0 [-n <VM Name>] [-i <VM UUID>] [-m <VM Glob Pattern>] [-h]"
            echo "  -n <VM Name>        Specify VM by exact name"
            echo "  -i <VM UUID>        Specify VM by UUID" 
            echo "  -m <VM Glob>        Specify VMs by glob pattern (case insensitive)"
            echo "  -h                  Show this help message"
            exit 0
            ;;
        \?)
            echo "Invalid option: -$OPTARG" >&2
            echo "Usage: $0 [-n <VM Name>] [-i <VM UUID>] [-m <VM Glob Pattern>] [-h]"
            exit 1
            ;;
        :)
            echo "Option -$OPTARG requires an argument." >&2
            exit 1
            ;;
    esac
done

shift $((OPTIND -1))

echo "[+] Searching for matching VMs..."
vm_matches=$(find_vms)

if [[ -z "$vm_matches" ]]; then
    echo "[-] No matching VMs found."
    exit 1
fi

echo "[+] Found matching VMs:"
echo "$vm_matches" | while IFS='|' read -r uuid status name; do
    status_color=$(get_status_color "$status")
    echo "   [*] $name ($uuid) - ${status_color}${status}${COLOR_RESET}"
done

vm_array=()
while IFS= read -r line; do
    if [[ -n "$line" ]]; then
        vm_array+=("$line")
    fi
done <<< "$vm_matches"

for vm_line in "${vm_array[@]}"; do
    IFS='|' read -r vm_uuid original_status vm_name <<< "$vm_line"
    
    echo ""
    echo "[+] Processing VM: $vm_name ($vm_uuid)"
    status_color=$(get_status_color "$original_status")
    echo "[+] Original status: ${status_color}${original_status}${COLOR_RESET}"
    
    # Start VM if not running
    if [[ "$original_status" != "running" ]]; then
        echo "[+] Starting VM..."
        if prlctl start "$vm_uuid" &>/dev/null; then
            echo "[+] VM started successfully"
        else
            echo "[-] Failed to start VM"
            continue
        fi
        wait_for_vm "$vm_uuid" "$original_status" || continue
    else
        echo "[+] VM is already running"
        wait_for_vm "$vm_uuid" "$original_status" || continue
    fi
    
    # Update the VM
    echo "[+] Updating VM: $vm_name"
    if update_vm "$vm_uuid"; then
        echo "[+] Updates completed successfully for $vm_name"
        
        # Only restore original state if updates were successful
        if [[ "$original_status" == "stopped" ]]; then
            echo "[+] Shutting down VM to restore original state..."
            if prlctl exec "$vm_uuid" 'shutdown -h now' &>/dev/null || prlctl stop "$vm_uuid" &>/dev/null; then
                echo "[+] Shutdown command sent"
            fi
            wait_for_state "$vm_uuid" "stopped"
        elif [[ "$original_status" == "suspended" ]]; then
            echo "[+] Suspending VM to restore original state..."
            if prlctl suspend "$vm_uuid" &>/dev/null; then
                echo "[+] Suspend command sent"
            fi
            wait_for_state "$vm_uuid" "suspended"
        fi
        
        echo "[+] VM $vm_name processing complete!"
    else
        echo "[!] Updates failed for $vm_name"
        echo "[-] Skipping state restoration due to update failure"
        echo "[-] VM $vm_name left in current state for troubleshooting"
    fi
done

echo ""
echo "[+] All VMs processed!"