#!/bin/bash
source ~/.bashrc
# Install and Configure QEMU Guest Agent on All VMs
# This script installs qemu-guest-agent on Ubuntu VMs and enables it in Proxmox

set -e

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

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

# Load environment variables
if [ -f .env ]; then
    set -a
    source <(grep -v '^#' .env | grep -v '^$' | sed 's/#.*$//' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//' | grep '=')
    set +a
fi

PVE_USERNAME="${PVE_USERNAME:-root@pam}"
PVE_PASSWORD="${PVE_ROOT_PASS:-}"
PROXMOX_URL="https://192.168.1.206:8006"
PROXMOX_NODE="pve"

# VM Configuration
declare -A VMS=(
    [100]="cloudflare-tunnel:192.168.1.60"
    [101]="k3s-master:192.168.1.188"
    [102]="git-server:192.168.1.121"
    [103]="observability:192.168.1.82"
)

SSH_USER="${SSH_USER:-ubuntu}"
SSH_KEY="${SSH_KEY:-~/.ssh/id_rsa}"

# Get authentication ticket
get_ticket() {
    local response=$(curl -k -s -d "username=$PVE_USERNAME&password=$PVE_PASSWORD" \
        "$PROXMOX_URL/api2/json/access/ticket")
    
    local ticket=$(echo "$response" | grep -o '"ticket":"[^"]*' | cut -d'"' -f4)
    local csrf=$(echo "$response" | grep -o '"CSRFPreventionToken":"[^"]*' | cut -d'"' -f4)
    
    if [ -z "$ticket" ] || [ -z "$csrf" ]; then
        log_error "Failed to authenticate with Proxmox"
        return 1
    fi
    
    echo "$ticket|$csrf"
}

# Check if VM is reachable
check_vm_reachable() {
    local ip=$1
    if ping -c 1 -W 2 "$ip" > /dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Check SSH connectivity
check_ssh() {
    local ip=$1
    local user=$2
    
    if ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no -i "$SSH_KEY" "${user}@${ip}" "echo 'SSH OK'" > /dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Install guest agent on VM
install_guest_agent_on_vm() {
    local vmid=$1
    local name=$2
    local ip=$3
    
    log_step "Installing QEMU Guest Agent on VM $vmid: $name"
    
    # Check if VM is reachable
    if ! check_vm_reachable "$ip"; then
        log_error "VM at $ip is not reachable, skipping..."
        return 1
    fi
    
    # Check SSH
    if ! check_ssh "$ip" "$SSH_USER"; then
        log_error "SSH not available on $ip, skipping..."
        return 1
    fi
    
    log_info "Installing qemu-guest-agent via SSH..."
    
    # Install qemu-guest-agent
    ssh -o StrictHostKeyChecking=no -i "$SSH_KEY" "${SSH_USER}@${ip}" <<EOF
sudo apt-get update
sudo apt-get install -y qemu-guest-agent
sudo systemctl enable qemu-guest-agent
sudo systemctl start qemu-guest-agent
sudo systemctl status qemu-guest-agent --no-pager | head -5
EOF
    
    if [ $? -eq 0 ]; then
        log_info "✓ Guest agent installed and started on VM $vmid"
        return 0
    else
        log_error "✗ Failed to install guest agent on VM $vmid"
        return 1
    fi
}

# Enable guest agent in Proxmox VM configuration
enable_guest_agent_in_proxmox() {
    local auth=$1
    local vmid=$2
    local name=$3
    
    local ticket=$(echo "$auth" | cut -d'|' -f1)
    local csrf=$(echo "$auth" | cut -d'|' -f2)
    
    log_step "Enabling guest agent in Proxmox for VM $vmid: $name"
    
    # Enable agent in VM config
    local response=$(curl -k -s -X PUT \
        -H "Cookie: PVEAuthCookie=$ticket" \
        -H "CSRFPreventionToken: $csrf" \
        -d "agent=1" \
        "$PROXMOX_URL/api2/json/nodes/$PROXMOX_NODE/qemu/$vmid/config" 2>&1)
    
    if echo "$response" | grep -q '"errors"'; then
        log_error "Failed to enable agent: $response"
        return 1
    fi
    
    log_info "✓ Guest agent enabled in Proxmox for VM $vmid"
    return 0
}

# Verify guest agent is working
verify_guest_agent() {
    local auth=$1
    local vmid=$2
    local name=$3
    
    local ticket=$(echo "$auth" | cut -d'|' -f1)
    local csrf=$(echo "$auth" | cut -d'|' -f2)
    
    log_step "Verifying guest agent for VM $vmid: $name"
    
    # Check agent status via Proxmox API
    local response=$(curl -k -s \
        -H "Cookie: PVEAuthCookie=$ticket" \
        -H "CSRFPreventionToken: $csrf" \
        "$PROXMOX_URL/api2/json/nodes/$PROXMOX_NODE/qemu/$vmid/agent/get-fsinfo" 2>&1)
    
    if echo "$response" | grep -q '"result"'; then
        log_info "✓ Guest agent is responding"
        return 0
    else
        log_warn "⚠ Guest agent may not be fully ready yet"
        log_info "  This is normal if VM was just configured"
        log_info "  Agent may take a few minutes to initialize"
        return 1
    fi
}

main() {
    echo "========================================="
    echo "Setup QEMU Guest Agent on All VMs"
    echo "========================================="
    echo ""
    
    if [ -z "$PVE_PASSWORD" ]; then
        log_error "PVE_ROOT_PASS not set in .env"
        exit 1
    fi
    
    if [ ! -f "$SSH_KEY" ]; then
        log_error "SSH key not found: $SSH_KEY"
        log_info "Set SSH_KEY environment variable or create key pair"
        exit 1
    fi
    
    log_info "Using SSH key: $SSH_KEY"
    log_info "SSH user: $SSH_USER"
    echo ""
    
    # Authenticate with Proxmox
    auth=$(get_ticket)
    if [ $? -ne 0 ]; then
        exit 1
    fi
    
    # Process each VM
    for vmid in 100 101 102 103; do
        IFS=':' read -r name ip <<< "${VMS[$vmid]}"
        
        echo "----------------------------------------"
        log_step "Processing VM $vmid: $name"
        echo ""
        
        # Step 1: Install guest agent on VM
        if install_guest_agent_on_vm "$vmid" "$name" "$ip"; then
            log_info "✓ Guest agent installed on VM"
        else
            log_error "✗ Failed to install guest agent"
            echo ""
            continue
        fi
        
        # Step 2: Enable agent in Proxmox
        if enable_guest_agent_in_proxmox "$auth" "$vmid" "$name"; then
            log_info "✓ Agent enabled in Proxmox"
        else
            log_error "✗ Failed to enable agent in Proxmox"
        fi
        
        # Step 3: Verify (optional, may take time)
        sleep 2
        verify_guest_agent "$auth" "$vmid" "$name" || true
        
        echo ""
    done
    
    log_info "========================================="
    log_info "Guest Agent Setup Complete"
    log_info "========================================="
    echo ""
    log_info "Benefits of QEMU Guest Agent:"
    echo "  • Proper VM shutdown/reboot from Proxmox"
    echo "  • Automatic IP address detection"
    echo "  • Better VM status reporting"
    echo "  • File system information"
    echo ""
    log_warn "Note: Guest agent may take a few minutes to fully initialize"
    log_info "You can verify in Proxmox Web UI:"
    echo "  VM → Monitor → QEMU Guest Agent"
}

main "$@"