#!/bin/bash
# enable-guest-agent-existing-vms.sh
# Enable QEMU guest agent on existing VMs via Proxmox API

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"

# Load environment
if [ -f "${PROJECT_ROOT}/.env" ]; then
    set -a
    source "${PROJECT_ROOT}/.env"
    set +a
fi

# Try API tokens first, fall back to password
PROXMOX_1_TOKEN="${PROXMOX_TOKEN_ML110_01:-}"
PROXMOX_2_TOKEN="${PROXMOX_TOKEN_R630_01:-}"
PROXMOX_PASS="${PROXMOX_ROOT_PASS:-L@kers2010}"
PROXMOX_1_URL="https://192.168.11.10:8006"
PROXMOX_2_URL="https://192.168.11.11:8006"

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

log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $*"
}

log_success() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] ✅${NC} $*"
}

log_error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ❌${NC} $*"
}

log_warning() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] ⚠️${NC} $*"
}

# Get auth - try API token first, fall back to password
get_auth() {
    local api_url=$1
    local api_token=$2
    local response
    
    # Try API token authentication first
    if [ -n "${api_token}" ]; then
        # Parse token format: root@pam!sankofa-instance-1-api-token=73c7e1a2-c969-409c-ae5b-68e83f012ee9
        # For Proxmox API tokens, we use the full token string in Authorization header
        response=$(curl -k -s -X GET \
            -H "Authorization: PVEAuthCookie=${api_token}" \
            "${api_url}/api2/json/version" 2>/dev/null)
        
        # If token auth works (we get version info), return token for direct use
        if echo "${response}" | grep -q "data\|version"; then
            echo "${api_token}|TOKEN"
            return 0
        fi
    fi
    
    # Fall back to password authentication
    response=$(curl -k -s -X POST \
        -d "username=root@pam&password=${PROXMOX_PASS}" \
        "${api_url}/api2/json/access/ticket" 2>/dev/null)
    
    if echo "${response}" | grep -q "authentication failure"; then
        echo ""
        return 1
    fi
    
    local ticket csrf
    if command -v jq &> /dev/null; then
        ticket=$(echo "${response}" | jq -r '.data.ticket // empty' 2>/dev/null)
        csrf=$(echo "${response}" | jq -r '.data.CSRFPreventionToken // empty' 2>/dev/null)
    else
        ticket=$(echo "${response}" | grep -o '"ticket":"[^"]*' | head -1 | cut -d'"' -f4)
        csrf=$(echo "${response}" | grep -o '"CSRFPreventionToken":"[^"]*' | head -1 | cut -d'"' -f4)
    fi
    
    if [ -z "${ticket}" ] || [ -z "${csrf}" ]; then
        echo ""
        return 1
    fi
    
    echo "${ticket}|${csrf}"
}

# List all nodes in the cluster
list_nodes() {
    local api_url=$1
    local auth_token=$2
    local auth_type=$3
    
    local response
    
    if [ "${auth_type}" = "TOKEN" ]; then
        response=$(curl -k -s -X GET \
            -H "Authorization: PVEAuthCookie=${auth_token}" \
            "${api_url}/api2/json/nodes" 2>/dev/null)
    else
        local ticket csrf
        IFS='|' read -r ticket csrf <<< "${auth_token}"
        response=$(curl -k -s -X GET \
            -H "CSRFPreventionToken: ${csrf}" \
            -b "PVEAuthCookie=${ticket}" \
            "${api_url}/api2/json/nodes" 2>/dev/null)
    fi
    
    # Extract node names from response
    if command -v jq &> /dev/null; then
        echo "${response}" | jq -r '.data[]?.node // empty' 2>/dev/null | grep -v '^$' | sort
    else
        # Fallback: extract node names using grep/sed
        echo "${response}" | grep -o '"node":"[^"]*' | cut -d'"' -f4 | sort | uniq
    fi
}

# List all VMs on a node
list_vms() {
    local api_url=$1
    local node=$2
    local auth_token=$3
    local auth_type=$4
    
    local response
    
    if [ "${auth_type}" = "TOKEN" ]; then
        # Use API token directly
        response=$(curl -k -s -X GET \
            -H "Authorization: PVEAuthCookie=${auth_token}" \
            "${api_url}/api2/json/nodes/${node}/qemu" 2>/dev/null)
    else
        # Use ticket and CSRF token
        local ticket csrf
        IFS='|' read -r ticket csrf <<< "${auth_token}"
        response=$(curl -k -s -X GET \
            -H "CSRFPreventionToken: ${csrf}" \
            -b "PVEAuthCookie=${ticket}" \
            "${api_url}/api2/json/nodes/${node}/qemu" 2>/dev/null)
    fi
    
    # Extract VMIDs from response
    if command -v jq &> /dev/null; then
        echo "${response}" | jq -r '.data[]?.vmid // empty' 2>/dev/null | grep -v '^$' | sort -n
    else
        # Fallback: extract VMIDs using grep/sed
        echo "${response}" | grep -o '"vmid":[0-9]*' | grep -o '[0-9]*' | sort -n | uniq
    fi
}

# Check if guest agent is already enabled
check_guest_agent() {
    local api_url=$1
    local node=$2
    local vmid=$3
    local auth_token=$4
    local auth_type=$5
    
    local response
    
    if [ "${auth_type}" = "TOKEN" ]; then
        response=$(curl -k -s -X GET \
            -H "Authorization: PVEAuthCookie=${auth_token}" \
            "${api_url}/api2/json/nodes/${node}/qemu/${vmid}/config" 2>/dev/null)
    else
        local ticket csrf
        IFS='|' read -r ticket csrf <<< "${auth_token}"
        response=$(curl -k -s -X GET \
            -H "CSRFPreventionToken: ${csrf}" \
            -b "PVEAuthCookie=${ticket}" \
            "${api_url}/api2/json/nodes/${node}/qemu/${vmid}/config" 2>/dev/null)
    fi
    
    # Check if agent is already enabled
    if echo "${response}" | grep -q '"agent"[[:space:]]*:[[:space:]]*"1"'; then
        return 0  # Already enabled
    fi
    return 1  # Not enabled
}

# Enable guest agent
enable_guest_agent() {
    local api_url=$1
    local node=$2
    local vmid=$3
    local auth_token=$4
    local auth_type=$5
    
    local response
    
    if [ "${auth_type}" = "TOKEN" ]; then
        # Use API token directly
        response=$(curl -k -s -X PUT \
            -H "Authorization: PVEAuthCookie=${auth_token}" \
            -d "agent=1" \
            "${api_url}/api2/json/nodes/${node}/qemu/${vmid}/config" 2>/dev/null)
    else
        # Use ticket and CSRF token
        local ticket csrf
        IFS='|' read -r ticket csrf <<< "${auth_token}"
        response=$(curl -k -s -X PUT \
            -H "CSRFPreventionToken: ${csrf}" \
            -b "PVEAuthCookie=${ticket}" \
            -d "agent=1" \
            "${api_url}/api2/json/nodes/${node}/qemu/${vmid}/config" 2>/dev/null)
    fi
    
    if echo "${response}" | grep -q '"data":null'; then
        return 0
    fi
    
    # Check if already enabled
    if echo "${response}" | grep -q "already"; then
        return 0
    fi
    
    return 1
}

process_node() {
    local api_url=$1
    local node=$2
    local auth_token=$3
    local auth_type=$4
    
    log "Processing node: ${node}"
    
    # Discover all VMs on this node
    local vmids
    vmids=$(list_vms "${api_url}" "${node}" "${auth_token}" "${auth_type}")
    
    if [ -z "${vmids}" ]; then
        log_warning "  No VMs found on ${node}"
        return 0
    fi
    
    local vm_count=0
    local enabled_count=0
    local skipped_count=0
    local failed_count=0
    
    while IFS= read -r vmid; do
        [ -z "${vmid}" ] && continue
        vm_count=$((vm_count + 1))
        
        # Check if already enabled
        if check_guest_agent "${api_url}" "${node}" "${vmid}" "${auth_token}" "${auth_type}"; then
            log "  VMID ${vmid}: guest agent already enabled"
            skipped_count=$((skipped_count + 1))
            continue
        fi
        
        log "  Enabling guest agent on VMID ${vmid}..."
        if enable_guest_agent "${api_url}" "${node}" "${vmid}" "${auth_token}" "${auth_type}"; then
            log_success "    VMID ${vmid} guest agent enabled"
            enabled_count=$((enabled_count + 1))
        else
            log_error "    Failed to enable guest agent on VMID ${vmid}"
            failed_count=$((failed_count + 1))
        fi
        sleep 0.3
    done <<< "${vmids}"
    
    log "  Summary for ${node}: ${vm_count} total, ${enabled_count} enabled, ${skipped_count} already enabled, ${failed_count} failed"
    
    # Return counts via global variables or echo
    echo "${vm_count}|${enabled_count}|${skipped_count}|${failed_count}"
}

process_site() {
    local api_url=$1
    local site_name=$2
    local auth_token=$3
    local auth_type=$4
    
    log "=========================================="
    log "Site: ${site_name}"
    log "=========================================="
    
    # Discover all nodes on this site
    local nodes
    nodes=$(list_nodes "${api_url}" "${auth_token}" "${auth_type}")
    
    if [ -z "${nodes}" ]; then
        log_error "Failed to discover nodes on ${site_name}"
        return 1
    fi
    
    log_success "Discovered nodes: $(echo "${nodes}" | tr '\n' ' ')"
    log ""
    
    local site_vm_count=0
    local site_enabled_count=0
    local site_skipped_count=0
    local site_failed_count=0
    
    # Process each node
    while IFS= read -r node; do
        [ -z "${node}" ] && continue
        
        local result
        result=$(process_node "${api_url}" "${node}" "${auth_token}" "${auth_type}")
        
        if [ -n "${result}" ]; then
            IFS='|' read -r vm_count enabled_count skipped_count failed_count <<< "${result}"
            site_vm_count=$((site_vm_count + vm_count))
            site_enabled_count=$((site_enabled_count + enabled_count))
            site_skipped_count=$((site_skipped_count + skipped_count))
            site_failed_count=$((site_failed_count + failed_count))
        fi
        
        log ""
    done <<< "${nodes}"
    
    log "Site Summary for ${site_name}:"
    log "  Total VMs: ${site_vm_count}"
    log "  Enabled: ${site_enabled_count}"
    log "  Already enabled: ${site_skipped_count}"
    log "  Failed: ${site_failed_count}"
    log ""
    
    echo "${site_vm_count}|${site_enabled_count}|${site_skipped_count}|${site_failed_count}"
}

main() {
    log "=========================================="
    log "Enable QEMU Guest Agent on All VMs"
    log "=========================================="
    log ""
    log "This script will:"
    log "1. Discover all nodes on each Proxmox site"
    log "2. Discover all VMs on each node"
    log "3. Check if guest agent is already enabled"
    log "4. Enable guest agent on VMs that need it"
    log ""
    
    local total_vm_count=0
    local total_enabled=0
    local total_skipped=0
    local total_failed=0
    
    # Site 1
    local auth1
    auth1=$(get_auth "${PROXMOX_1_URL}" "${PROXMOX_1_TOKEN}")
    
    if [ -z "${auth1}" ]; then
        log_error "Failed to authenticate to Site 1"
    else
        IFS='|' read -r auth_token1 auth_type1 <<< "${auth1}"
        log_success "Authenticated to Site 1"
        log ""
        
        local result1
        result1=$(process_site "${PROXMOX_1_URL}" "Site 1" "${auth_token1}" "${auth_type1}")
        
        if [ -n "${result1}" ]; then
            IFS='|' read -r vm_count enabled_count skipped_count failed_count <<< "${result1}"
            total_vm_count=$((total_vm_count + vm_count))
            total_enabled=$((total_enabled + enabled_count))
            total_skipped=$((total_skipped + skipped_count))
            total_failed=$((total_failed + failed_count))
        fi
    fi
    
    # Site 2
    local auth2
    auth2=$(get_auth "${PROXMOX_2_URL}" "${PROXMOX_2_TOKEN}")
    
    if [ -z "${auth2}" ]; then
        log_error "Failed to authenticate to Site 2"
    else
        IFS='|' read -r auth_token2 auth_type2 <<< "${auth2}"
        log_success "Authenticated to Site 2"
        log ""
        
        local result2
        result2=$(process_site "${PROXMOX_2_URL}" "Site 2" "${auth_token2}" "${auth_type2}")
        
        if [ -n "${result2}" ]; then
            IFS='|' read -r vm_count enabled_count skipped_count failed_count <<< "${result2}"
            total_vm_count=$((total_vm_count + vm_count))
            total_enabled=$((total_enabled + enabled_count))
            total_skipped=$((total_skipped + skipped_count))
            total_failed=$((total_failed + failed_count))
        fi
    fi
    
    log ""
    log "=========================================="
    log "Overall Summary"
    log "=========================================="
    log "Total VMs processed: ${total_vm_count}"
    log_success "Guest agent enabled: ${total_enabled}"
    log "Already enabled: ${total_skipped}"
    if [ "${total_failed}" -gt 0 ]; then
        log_error "Failed: ${total_failed}"
    else
        log_success "Failed: ${total_failed}"
    fi
    log ""
    log "=========================================="
    log_success "Guest agent enablement complete!"
    log "=========================================="
    log ""
    log_warning "IMPORTANT: Guest agent must also be installed in the OS."
    log ""
    log "For existing VMs, you need to:"
    log "1. Wait for VMs to get IP addresses"
    log "2. SSH into each VM: ssh admin@<vm-ip>"
    log "3. Install and enable guest agent:"
    log "   sudo apt-get update"
    log "   sudo apt-get install -y qemu-guest-agent"
    log "   sudo systemctl enable qemu-guest-agent"
    log "   sudo systemctl start qemu-guest-agent"
    log ""
    log "Note: New VMs created with the updated Crossplane provider will"
    log "automatically have guest agent enabled in Proxmox config."
    log ""
}

main "$@"