#!/usr/bin/env bash
# Fix Storage Configuration for Proxmox Migrations
# Ensures target nodes have compatible storage for migrations

set -euo pipefail

# Configuration
PROXMOX_HOST_PVE="192.168.11.11"
PROXMOX_HOST_PVE2="192.168.11.12"
PVE_PASS="password"
PVE2_PASS="password"

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

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[✓]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_header() { echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"; }

# SSH helper
ssh_node() {
    local host=$1
    local pass=$2
    shift 2
    sshpass -p "$pass" ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new root@"$host" "$@" 2>&1
}

# Check if storage exists and is active
check_storage() {
    local host=$1
    local pass=$2
    local storage=$3
    
    local status=$(ssh_node "$host" "$pass" "pvesm status 2>/dev/null | grep '^$storage' | awk '{print \$3}'" || echo "")
    
    if [ -z "$status" ]; then
        return 1  # Not found
    fi
    
    if echo "$status" | grep -qi "active\|enabled"; then
        return 0  # Active
    else
        return 2  # Exists but inactive
    fi
}

# Ensure thin1 storage is available
ensure_thin1_storage() {
    local host=$1
    local pass=$2
    local node_name=$3
    
    log_info "Ensuring thin1 storage is available on $node_name..."
    
    # Check if thin1 exists
    check_storage "$host" "$pass" "thin1"
    local check_result=$?
    
    if [ $check_result -eq 0 ]; then
        log_success "thin1 storage is already active on $node_name"
        return 0
    elif [ $check_result -eq 2 ]; then
        log_warn "thin1 exists but is inactive, attempting to enable..."
        # Try to enable it
        ssh_node "$host" "$pass" "pvesm set thin1 --content images,rootdir" || true
        sleep 2
        check_storage "$host" "$pass" "thin1"
        if [ $? -eq 0 ]; then
            log_success "thin1 storage enabled on $node_name"
            return 0
        fi
    fi
    
    # thin1 doesn't exist, check if we can create it
    log_info "thin1 storage not found, checking for available volume groups..."
    
    # Get available VGs
    local vgs=$(ssh_node "$host" "$pass" "vgs --noheadings -o vg_name 2>/dev/null | head -1" | tr -d ' ' || echo "")
    
    if [ -z "$vgs" ]; then
        log_error "No volume groups found on $node_name"
        log_info "Available storage:"
        ssh_node "$host" "$pass" "pvesm status" || true
        return 1
    fi
    
    log_info "Found volume group: $vgs"
    
    # Check if thin pool exists
    local thin_pool=$(ssh_node "$host" "$pass" "lvs $vgs --noheadings -o lv_name 2>/dev/null | grep -E '(thin1|data)' | head -1" | tr -d ' ' || echo "")
    
    if [ -z "$thin_pool" ]; then
        log_warn "No suitable thin pool found in $vgs"
        log_info "Checking available space..."
        
        local vg_free=$(ssh_node "$host" "$pass" "vgs -o vg_free --noheadings --units g $vgs 2>/dev/null | awk '{print int(\$1)}'" || echo "0")
        
        if [ "$vg_free" -lt 10 ]; then
            log_error "Not enough free space in $vgs (${vg_free}G available)"
            return 1
        fi
        
        log_info "Creating thin pool 'thin1' in $vgs..."
        local pool_size=$((vg_free * 80 / 100))
        
        ssh_node "$host" "$pass" <<EOF
lvcreate -L ${pool_size}G -n thin1 ${vgs} 2>&1
lvconvert --type thin-pool ${vgs}/thin1 2>&1
EOF
        
        thin_pool="thin1"
    else
        log_info "Found existing thin pool: $vgs/$thin_pool"
    fi
    
    # Add storage to Proxmox
    log_info "Adding thin1 storage to Proxmox configuration..."
    
    ssh_node "$host" "$pass" <<EOF
pvesm add lvmthin thin1 \
    --thinpool ${thin_pool} \
    --vgname ${vgs} \
    --content images,rootdir \
    --nodes ${node_name} 2>&1 || true
EOF
    
    sleep 2
    
    # Verify
    check_storage "$host" "$pass" "thin1"
    if [ $? -eq 0 ]; then
        log_success "thin1 storage is now active on $node_name"
        return 0
    else
        log_error "Failed to create/enable thin1 storage on $node_name"
        return 1
    fi
}

# Fix storage on a node
fix_node_storage() {
    local host=$1
    local pass=$2
    local node_name=$3
    
    log_header
    log_info "Fixing storage on $node_name ($host)"
    log_header
    echo ""
    
    # Check current storage
    log_info "Current storage status:"
    ssh_node "$host" "$pass" "pvesm status" || true
    echo ""
    
    # Ensure thin1 is available
    if ensure_thin1_storage "$host" "$pass" "$node_name"; then
        log_success "Storage configuration complete for $node_name"
        echo ""
        
        # Show final status
        log_info "Final storage status:"
        ssh_node "$host" "$pass" "pvesm status" || true
        echo ""
        return 0
    else
        log_error "Failed to configure storage on $node_name"
        echo ""
        return 1
    fi
}

# Main execution
main() {
    echo ""
    log_header
    log_info "Proxmox Storage Fix Tool"
    log_header
    echo ""
    
    log_info "This script will ensure target nodes have compatible storage for migrations"
    echo ""
    
    # Check for non-interactive mode
    if [[ "${NON_INTERACTIVE:-}" == "1" ]] || [[ ! -t 0 ]]; then
        log_info "Non-interactive mode: proceeding automatically"
    else
        read -p "Continue? (y/N): " -n 1 -r
        echo ""
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "Operation cancelled"
            exit 0
        fi
    fi
    
    echo ""
    
    local failed=0
    
    # Fix pve
    if sshpass -p "$PVE_PASS" ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new root@"$PROXMOX_HOST_PVE" "echo 'connected'" 2>/dev/null; then
        if ! fix_node_storage "$PROXMOX_HOST_PVE" "$PVE_PASS" "pve"; then
            failed=$((failed + 1))
        fi
    else
        log_error "Cannot connect to pve"
        failed=$((failed + 1))
    fi
    
    # Fix pve2
    if sshpass -p "$PVE2_PASS" ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=accept-new root@"$PROXMOX_HOST_PVE2" "echo 'connected'" 2>/dev/null; then
        if ! fix_node_storage "$PROXMOX_HOST_PVE2" "$PVE2_PASS" "pve2"; then
            failed=$((failed + 1))
        fi
    else
        log_warn "Cannot connect to pve2, skipping..."
    fi
    
    echo ""
    log_header
    log_info "Summary"
    log_header
    echo ""
    
    if [ $failed -eq 0 ]; then
        log_success "Storage configuration complete on all nodes!"
        log_info "You can now run the migration script:"
        log_info "  ./scripts/diagnose-and-fix-migration-storage.sh"
        echo ""
    else
        log_warn "Storage configuration completed with some failures"
        log_info "Review the output above for details"
        echo ""
    fi
}

main "$@"