#!/usr/bin/env bash
# Check Azure Prerequisites: Regions, Resource Providers, and Quotas
# Ensures all required providers are registered and quotas are available

set -euo pipefail

# Colors for output

# Required Resource Providers
REQUIRED_PROVIDERS=(
    "Microsoft.ContainerService"      # AKS
    "Microsoft.KeyVault"              # Key Vault
    "Microsoft.Storage"                # Storage Accounts
    "Microsoft.Network"                # VNet, Application Gateway, Load Balancer
    "Microsoft.Compute"                # VMs, VM Scale Sets
    "Microsoft.Insights"                # Azure Monitor
    "Microsoft.ManagedIdentity"        # Managed Identities
    "Microsoft.Authorization"          # RBAC
    "Microsoft.Resources"               # Resource Groups
)

# Required Preview Features (if any)
REQUIRED_PREVIEW_FEATURES=(
    # Add any preview features needed
)

# VM Sizes needed for quota checking
VM_SIZES=(
    "Standard_D2s_v3"   # System nodes
    "Standard_D4s_v3"   # Validators and Sentries
    "Standard_D8s_v3"   # RPC nodes
)

# Node counts per VM size
declare -A NODE_COUNTS=(
    ["Standard_D2s_v3"]=3
    ["Standard_D4s_v3"]=7
    ["Standard_D8s_v3"]=3
)

log() {
    log_success "[✓] $1"
}

error() {
    log_error "[✗] $1"
}

warn() {
    log_warn "[!] $1"
}

info() {
    log_info "[i] $1"
}

section() {
    echo
    log_info "=== $1 ==="
}

# Check Azure CLI
check_azure_cli() {
    if ! command -v az &> /dev/null; then
        error "Azure CLI is not installed"
        exit 1
    fi
    
    if ! az account show &> /dev/null; then
        error "Not logged in to Azure. Run 'az login' first"
        exit 1
    fi
    
    log "Azure CLI is installed and authenticated"
}

# Get Commercial Regions (excluding US and Government)
get_commercial_regions() {
    section "Available Commercial Regions (Non-US, Non-Government)"
    
    # Get all locations, filter out US and Government regions
    az account list-locations \
        --query "[?metadata.regionType=='Physical' && !contains(name, 'us') && !contains(name, 'gov') && !contains(name, 'dod')].{Name:name, DisplayName:displayName, Geography:metadata.geographyGroup}" \
        -o table
    
    # Get as array for processing
    COMMERCIAL_REGIONS=$(az account list-locations \
        --query "[?metadata.regionType=='Physical' && !contains(name, 'us') && !contains(name, 'gov') && !contains(name, 'dod')].name" \
        -o tsv)
    
    echo
    info "Total Commercial Regions (Non-US): $(echo "$COMMERCIAL_REGIONS" | wc -l)"
}

# Check Resource Provider Registration
check_resource_providers() {
    section "Required Resource Providers"
    
    local all_registered=true
    local unregistered=()
    
    for provider in "${REQUIRED_PROVIDERS[@]}"; do
        local status=$(az provider show --namespace "$provider" --query "registrationState" -o tsv 2>/dev/null || echo "NotRegistered")
        
        if [ "$status" == "Registered" ]; then
            log "$provider: Registered"
        else
            error "$provider: $status"
            unregistered+=("$provider")
            all_registered=false
        fi
    done
    
    echo
    
    if [ "$all_registered" = false ]; then
        warn "Some resource providers are not registered"
        info "Registering unregistered providers..."
        
        for provider in "${unregistered[@]}"; do
            info "Registering $provider..."
            az provider register --namespace "$provider" --wait || warn "Failed to register $provider"
        done
        
        # Wait for registration to complete
        info "Waiting for provider registrations to complete..."
        sleep 10
        
        # Verify registration
        for provider in "${unregistered[@]}"; do
            local status=$(az provider show --namespace "$provider" --query "registrationState" -o tsv 2>/dev/null || echo "NotRegistered")
            if [ "$status" == "Registered" ]; then
                log "$provider: Now registered"
            else
                warn "$provider: Still $status (may take a few minutes)"
            fi
        done
    else
        log "All required resource providers are registered"
    fi
}

# Check Preview Features
check_preview_features() {
    if [ ${#REQUIRED_PREVIEW_FEATURES[@]} -eq 0 ]; then
        info "No preview features required"
        return
    fi
    
    section "Required Preview Features"
    
    for feature in "${REQUIRED_PREVIEW_FEATURES[@]}"; do
        # Check if feature is registered
        local registered=$(az feature show --name "$feature" --namespace "Microsoft.ContainerService" --query "properties.state" -o tsv 2>/dev/null || echo "Unknown")
        
        if [ "$registered" == "Registered" ]; then
            log "$feature: Registered"
        else
            warn "$feature: $registered"
            info "Registering preview feature: $feature"
            az feature register --name "$feature" --namespace "Microsoft.ContainerService" || warn "Failed to register $feature"
        fi
    done
}

# Check Quotas for a specific region
check_region_quotas() {
    local region=$1
    
    section "Quotas for Region: $region"
    
    # Get subscription ID
    local sub_id=$(az account show --query id -o tsv)
    
    # Check VM family quotas
    info "Checking VM quotas..."
    
    for vm_size in "${VM_SIZES[@]}"; do
        # Extract VM family (e.g., StandardD2sv3Family from Standard_D2s_v3)
        local family=$(echo "$vm_size" | sed 's/_/-/g' | sed 's/Standard-/Standard/' | sed 's/$/Family/')
        
        # Get quota for this VM family
        local quota=$(az vm list-usage \
            --location "$region" \
            --query "[?name.value=='$family'].{Name:name.value, Current:currentValue, Limit:limit}" \
            -o tsv 2>/dev/null || echo "")
        
        if [ -n "$quota" ]; then
            local current=$(echo "$quota" | awk '{print $2}')
            local limit=$(echo "$quota" | awk '{print $3}')
            local needed=${NODE_COUNTS[$vm_size]}
            
            if [ -z "$current" ] || [ -z "$limit" ]; then
                warn "$vm_size: Could not retrieve quota"
            else
                local available=$((limit - current))
                if [ $available -ge $needed ]; then
                    log "$vm_size: $current/$limit used, $available available (need $needed) ✓"
                else
                    error "$vm_size: $current/$limit used, only $available available (need $needed) ✗"
                fi
            fi
        else
            warn "$vm_size: Quota information not available"
        fi
    done
    
    # Check AKS quota
    info "Checking AKS cluster quota..."
    local aks_quota=$(az aks list --query "length(@)" -o tsv 2>/dev/null || echo "0")
    log "Current AKS clusters: $aks_quota (typically 50 per subscription)"
    
    # Check Public IP quota
    info "Checking Public IP quota..."
    local pip_quota=$(az network list-usages \
        --location "$region" \
        --query "[?name.value=='PublicIPAddresses'].{Current:currentValue, Limit:limit}" \
        -o tsv 2>/dev/null || echo "")
    
    if [ -n "$pip_quota" ]; then
        local current=$(echo "$pip_quota" | awk '{print $1}')
        local limit=$(echo "$pip_quota" | awk '{print $2}')
        log "Public IPs: $current/$limit used"
    fi
    
    # Check Load Balancer quota
    info "Checking Load Balancer quota..."
    local lb_quota=$(az network list-usages \
        --location "$region" \
        --query "[?name.value=='LoadBalancers'].{Current:currentValue, Limit:limit}" \
        -o tsv 2>/dev/null || echo "")
    
    if [ -n "$lb_quota" ]; then
        local current=$(echo "$lb_quota" | awk '{print $1}')
        local limit=$(echo "$lb_quota" | awk '{print $2}')
        log "Load Balancers: $current/$limit used"
    fi
}

# Check quotas for all commercial regions
check_all_region_quotas() {
    section "Quota Check for All Commercial Regions"
    
    local default_region="westeurope"
    local regions_to_check=("$default_region")
    
    # Optionally check other key regions
    local key_regions=("northeurope" "uksouth" "francecentral" "germanywestcentral")
    
    for region in "${regions_to_check[@]}" "${key_regions[@]}"; do
        # Verify region exists in commercial regions list
        if echo "$COMMERCIAL_REGIONS" | grep -q "^${region}$"; then
            check_region_quotas "$region"
        fi
    done
}

# Main function
main() {
    log_info "Azure Prerequisites Check"
    log_info "========================="
    echo
    
    check_azure_cli
    
    get_commercial_regions
    
    check_resource_providers
    
    check_preview_features
    
    check_all_region_quotas
    
    echo
    section "Summary"
    log "Default region set to: westeurope"
    log "All commercial regions (non-US, non-Government) listed above"
    log "Resource providers checked and registered if needed"
    log "Quotas checked for key regions"
    
    echo
    info "Next steps:"
    info "1. Review quota availability in your preferred region"
    info "2. Request quota increases if needed: az vm list-usage --location westeurope"
    info "3. Update .env file with: AZURE_LOCATION=westeurope"
    info "4. Update terraform.tfvars with: location = \"westeurope\""
}

# Run main function
main "$@"