#!/usr/bin/env bash
# Create volume groups thin1-thin3 on pve using available disks
# This script sets up VGs and thin pools for pve similar to pve2

set -euo pipefail

PROXMOX_HOST="192.168.11.11"
PVE_PASS="password"

# Map storage name to disk
declare -A STORAGE_DISKS=(
    ["thin1"]="sdc"
    ["thin2"]="sdd"
    ["thin3"]="sde"
)

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
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"; }

# Check if disk is available
check_disk() {
    local disk=$1
    local vg_name=$2
    
    log_info "Checking disk /dev/$disk..."
    
    # Check if disk exists
    if ! sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
        "lsblk /dev/$disk 2>/dev/null" > /dev/null; then
        log_error "Disk /dev/$disk does not exist"
        return 1
    fi
    
    # Check if disk has partitions
    local partitions=$(sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
        "lsblk -n /dev/$disk 2>/dev/null | grep -c part || echo 0" | tr -d '\n' || echo "0")
    
    if [ "${partitions:-0}" -gt 0 ] 2>/dev/null; then
        log_warn "Disk /dev/$disk has partitions. Checking if already in use..."
        # Check if already part of a VG
        local pvs=$(sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
            "pvs 2>/dev/null | grep $disk" || echo "")
        if [ -n "$pvs" ]; then
            log_warn "Disk /dev/$disk is already in use by a volume group"
            return 1
        fi
    fi
    
    # Check if VG already exists
    if sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
        "vgs $vg_name 2>/dev/null | grep -q '$vg_name'" 2>/dev/null; then
        log_warn "Volume group $vg_name already exists"
        return 2
    fi
    
    log_success "Disk /dev/$disk is available"
    return 0
}

# Create physical volume
create_pv() {
    local disk=$1
    
    log_info "Creating physical volume on /dev/$disk..."
    
    sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} <<EOF
# Create PV directly on disk (LVM can work on whole disk)
pvcreate -y /dev/$disk 2>&1
EOF
    
    if [ $? -eq 0 ]; then
        log_success "Physical volume created on /dev/$disk"
        return 0
    else
        log_error "Failed to create physical volume"
        return 1
    fi
}

# Create volume group and thin pool
create_vg_and_thinpool() {
    local vg_name=$1
    local disk=$2
    local pool_name=$3
    
    log_info "Creating volume group $vg_name and thin pool..."
    
    # Use the whole disk
    local pv_path="/dev/$disk"
    
    sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} <<EOF
# Create VG
vgcreate -y $vg_name $pv_path 2>&1

# Get VG size and create thin pool (use 95% of available space)
VG_SIZE=\$(vgs -o vg_size --noheadings --units g $vg_name | awk '{print int(\$1)}')
POOL_SIZE=\$((VG_SIZE * 95 / 100))

# Create thin pool (use --yes to avoid interactive prompt)
lvcreate -L \${POOL_SIZE}G -n $pool_name $vg_name 2>&1
lvconvert --type thin-pool --yes $vg_name/$pool_name 2>&1
EOF
    
    if [ $? -eq 0 ]; then
        log_success "Volume group $vg_name and thin pool created"
        return 0
    else
        log_error "Failed to create VG and thin pool"
        return 1
    fi
}

# Add storage to Proxmox
add_storage() {
    local storage=$1
    local vg_name=$2
    local pool_name=$3
    
    log_info "Adding storage $storage to Proxmox..."
    
    sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} <<EOF
pvesm add lvmthin $storage \
    --thinpool $pool_name \
    --vgname $vg_name \
    --content images,rootdir \
    --nodes pve 2>&1
EOF
    
    if [ $? -eq 0 ]; then
        log_success "Storage $storage added to Proxmox"
        return 0
    else
        log_warn "Storage add returned non-zero, checking if it exists..."
        if sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
            "pvesm status 2>/dev/null | grep -q '$storage'" 2>/dev/null; then
            log_success "Storage $storage exists"
            return 0
        fi
        return 1
    fi
}

# Process a storage
process_storage() {
    local storage=$1
    local disk=${STORAGE_DISKS[$storage]}
    local vg_name=$storage
    local pool_name=$storage
    
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log_info "Processing: $storage (disk: $disk, VG: $vg_name)"
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    
    # Check disk
    local check_result
    check_disk "$disk" "$vg_name"
    check_result=$?
    
    if [ $check_result -eq 1 ]; then
        log_error "Cannot proceed with $storage"
        return 1
    elif [ $check_result -eq 2 ]; then
        log_info "VG exists, checking thin pool and storage..."
        # Check if thin pool exists and is properly configured
        local lv_attr=$(sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
            "lvs $vg_name/$pool_name -o lv_attr --noheadings 2>/dev/null" | tr -d ' ' || echo "")
        
        if [[ "$lv_attr" == *"t"* ]] && [[ "$lv_attr" == *"i"* ]]; then
            log_info "Thin pool exists and is properly configured, adding storage..."
            if ! add_storage "$storage" "$vg_name" "$pool_name"; then
                log_error "Failed to add storage $storage"
                return 1
            fi
        else
            log_info "VG exists but thin pool needs to be created/converted..."
            # Check if LV exists
            if sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
                "lvs $vg_name/$pool_name 2>/dev/null" > /dev/null 2>&1; then
                log_info "LV exists, converting to thin pool..."
                sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} \
                    "lvconvert --type thin-pool --yes $vg_name/$pool_name 2>&1" || {
                    log_error "Failed to convert LV to thin pool"
                    return 1
                }
            else
                log_info "Creating new thin pool..."
                sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} <<EOF
VG_SIZE=\$(vgs -o vg_size --noheadings --units g $vg_name | awk '{print int(\$1)}')
POOL_SIZE=\$((VG_SIZE * 95 / 100))
lvcreate -L \${POOL_SIZE}G -n $pool_name $vg_name 2>&1
lvconvert --type thin-pool --yes $vg_name/$pool_name 2>&1
EOF
            fi
            if ! add_storage "$storage" "$vg_name" "$pool_name"; then
                log_error "Failed to add storage $storage"
                return 1
            fi
        fi
        return 0
    fi
    
    # Create PV
    if ! create_pv "$disk"; then
        log_error "Failed to create physical volume"
        return 1
    fi
    
    # Create VG and thin pool
    if ! create_vg_and_thinpool "$vg_name" "$disk" "$pool_name"; then
        log_error "Failed to create VG and thin pool"
        return 1
    fi
    
    # Add storage to Proxmox
    if ! add_storage "$storage" "$vg_name" "$pool_name"; then
        log_error "Failed to add storage to Proxmox"
        return 1
    fi
    
    log_success "$storage setup complete"
    echo ""
    return 0
}

# Main execution
main() {
    echo ""
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log_info "Create Volume Groups for pve (thin1-thin3)"
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    
    log_warn "This script will:"
    log_warn "  1. Create physical volumes on sdc, sdd, sde"
    log_warn "  2. Create volume groups thin1, thin2, thin3"
    log_warn "  3. Create thin pools"
    log_warn "  4. Add storage to Proxmox"
    echo ""
    log_warn "⚠️  This will use the disks and create partitions!"
    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
    
    # Process each storage
    for storage in thin1 thin2 thin3; do
        process_storage "$storage"
    done
    
    # Show final status
    echo ""
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log_info "Final Status"
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    log_info "Volume groups:"
    sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} "vgs thin1 thin2 thin3 2>/dev/null" || true
    echo ""
    log_info "Storage status:"
    sshpass -p "$PVE_PASS" ssh -o StrictHostKeyChecking=accept-new root@${PROXMOX_HOST} "pvesm status | grep -E '(thin1|thin2|thin3)'" || true
    echo ""
}

main "$@"