#!/bin/bash
# Pre-deployment quota check for VM deployments
# Checks quota before applying VM manifests

set -euo pipefail

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

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

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

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

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

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

# Load environment
load_env() {
    if [ -f "${PROJECT_ROOT}/.env" ]; then
        source "${PROJECT_ROOT}/.env"
    fi
    
    SANKOFA_API_URL="${SANKOFA_API_URL:-http://localhost:4000/graphql}"
    SANKOFA_API_TOKEN="${SANKOFA_API_TOKEN:-}"
}

# Extract VM resources from YAML
extract_vm_resources() {
    local file=$1
    
    local cpu=$(grep "cpu:" "$file" | head -1 | sed 's/.*cpu: *\([0-9]*\).*/\1/')
    local memory=$(grep "memory:" "$file" | head -1 | sed 's/.*memory: *"\(.*\)".*/\1/')
    local disk=$(grep "disk:" "$file" | head -1 | sed 's/.*disk: *"\(.*\)".*/\1/')
    local tenant=$(grep "tenant.sankofa.nexus/id:" "$file" | head -1 | sed 's/.*tenant.sankofa.nexus\/id: *"\(.*\)".*/\1/' || echo "")
    
    echo "${cpu}|${memory}|${disk}|${tenant}"
}

# Convert memory to GB
memory_to_gb() {
    local memory=$1
    if [[ "$memory" =~ Gi$ ]]; then
        echo "${memory%Gi}"
    elif [[ "$memory" =~ Mi$ ]]; then
        echo "scale=2; ${memory%Mi} / 1024" | bc
    else
        echo "$memory"
    fi
}

# Convert disk to GB
disk_to_gb() {
    local disk=$1
    if [[ "$disk" =~ Gi$ ]]; then
        echo "${disk%Gi}"
    elif [[ "$disk" =~ Ti$ ]]; then
        echo "scale=2; ${disk%Ti} * 1024" | bc
    else
        echo "$disk"
    fi
}

# Check quota via API
check_quota_api() {
    local tenant_id=$1
    local cpu=$2
    local memory_gb=$3
    local disk_gb=$4
    
    if [[ -z "$SANKOFA_API_TOKEN" ]]; then
        log_warning "SANKOFA_API_TOKEN not set, skipping API quota check"
        return 0
    fi
    
    local query=$(cat <<EOF
mutation {
  checkQuota(
    tenantId: "$tenant_id"
    resourceRequest: {
      compute: {
        vcpu: $cpu
        memory: $memory_gb
        instances: 1
      }
      storage: {
        size: $disk_gb
      }
    }
  ) {
    allowed
    exceeded
    current {
      compute {
        instances
        vcpu
        memory
      }
      storage {
        total
      }
    }
    limits {
      compute {
        instances
        vcpu
        memory
      }
      storage {
        total
      }
    }
  }
}
EOF
)
    
    local response=$(curl -s -X POST "$SANKOFA_API_URL" \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $SANKOFA_API_TOKEN" \
        -d "{\"query\": \"$query\"}")
    
    local allowed=$(echo "$response" | jq -r '.data.checkQuota.allowed // false')
    local exceeded=$(echo "$response" | jq -r '.data.checkQuota.exceeded // []' | jq -r '.[]' | tr '\n' ',' | sed 's/,$//')
    
    if [[ "$allowed" == "true" ]]; then
        log_success "Quota check passed for tenant: $tenant_id"
        return 0
    else
        log_error "Quota check failed for tenant: $tenant_id"
        if [[ -n "$exceeded" ]]; then
            log_error "Exceeded resources: $exceeded"
        fi
        return 1
    fi
}

# Check Proxmox resources
check_proxmox_resources() {
    local file=$1
    local cpu=$2
    local memory_gb=$3
    local disk_gb=$4
    
    log "Checking Proxmox resources..."
    
    # Run the quota check script
    if [[ -f "${SCRIPT_DIR}/check-proxmox-quota-ssh.sh" ]]; then
        "${SCRIPT_DIR}/check-proxmox-quota-ssh.sh" > /tmp/quota-check-output.txt 2>&1 || true
        
        # Parse output (simplified - would need more robust parsing)
        log_info "Proxmox resource check completed (see output above)"
    else
        log_warning "Proxmox quota check script not found"
    fi
}

# Check single VM file
check_vm_file() {
    local file=$1
    
    log "Checking: $file"
    
    local resources=$(extract_vm_resources "$file")
    IFS='|' read -r cpu memory disk tenant <<< "$resources"
    
    if [[ -z "$cpu" || -z "$memory" || -z "$disk" ]]; then
        log_error "$file: Missing resource specifications"
        return 1
    fi
    
    local memory_gb=$(memory_to_gb "$memory")
    local disk_gb=$(disk_to_gb "$disk")
    
    log "Resources: CPU=$cpu, Memory=${memory_gb}GB, Disk=${disk_gb}GB, Tenant=${tenant:-none}"
    
    # Check quota if tenant is specified
    if [[ -n "$tenant" && "$tenant" != "infrastructure" ]]; then
        check_quota_api "$tenant" "$cpu" "$memory_gb" "$disk_gb" || return 1
    else
        log_info "Infrastructure VM - skipping tenant quota check"
        check_proxmox_resources "$file" "$cpu" "$memory_gb" "$disk_gb"
    fi
    
    return 0
}

# Main function
main() {
    local files_to_check=("$@")
    
    load_env
    
    log "Pre-deployment quota check"
    echo
    
    if [[ ${#files_to_check[@]} -eq 0 ]]; then
        log "No files specified. Checking all VM files..."
        while IFS= read -r -d '' file; do
            files_to_check+=("$file")
        done < <(find "$VM_DIR" -name "*.yaml" -type f -print0)
    fi
    
    local total_files=${#files_to_check[@]}
    local passed=0
    local failed=0
    
    for file in "${files_to_check[@]}"; do
        if check_vm_file "$file"; then
            passed=$((passed + 1))
        else
            failed=$((failed + 1))
        fi
        echo
    done
    
    # Summary
    echo
    log "=== Quota Check Summary ==="
    echo "Files checked: $total_files"
    echo "Passed: $passed"
    echo "Failed: $failed"
    echo
    
    if [[ $failed -eq 0 ]]; then
        log_success "All quota checks passed!"
        return 0
    else
        log_error "$failed quota check(s) failed. Do not proceed with deployment."
        return 1
    fi
}

main "$@"