#!/usr/bin/env bash
# Check AI, ML, OpenAI, and related Resource Providers and Preview Features
# Ensures all AI/ML related providers are registered

set -euo pipefail

# Colors for output

# AI/ML/OpenAI Related Resource Providers
AI_ML_PROVIDERS=(
    "Microsoft.CognitiveServices"      # Azure Cognitive Services (OpenAI, etc.)
    "Microsoft.MachineLearningServices" # Azure Machine Learning
    "Microsoft.MachineLearning"        # Legacy ML services
    "Microsoft.AI"                     # Azure AI services
    "Microsoft.OpenAI"                # OpenAI service (if exists)
    "Microsoft.AzureOpenAI"            # Azure OpenAI Service
    "Microsoft.Cognitive"             # Cognitive Services (legacy)
    "Microsoft.BotService"             # Bot Framework (AI-powered)
    "Microsoft.Search"                 # Azure AI Search (formerly Cognitive Search)
    "Microsoft.OperationalInsights"    # Log Analytics (used by AI/ML)
    "Microsoft.Synapse"                # Azure Synapse Analytics (ML integration)
    "Microsoft.Databricks"              # Azure Databricks (ML platform)
    "Microsoft.DataFactory"            # Azure Data Factory (ML pipelines)
)

# AI/ML Related Preview Features
AI_ML_PREVIEW_FEATURES=(
    # Add any preview features related to AI/ML
    # Example: "AksWindows2019Preview"
)

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"
}

# Check Resource Provider Registration
check_ai_ml_providers() {
    section "AI/ML/OpenAI Resource Providers"
    
    local all_registered=true
    local unregistered=()
    local not_found=()
    
    for provider in "${AI_ML_PROVIDERS[@]}"; do
        # Check if provider exists
        local provider_info=$(az provider show --namespace "$provider" 2>/dev/null || echo "")
        
        if [ -z "$provider_info" ]; then
            # Provider namespace doesn't exist (might be wrong name or not available)
            not_found+=("$provider")
            warn "$provider: Namespace not found (may not exist or not available in this subscription)"
            continue
        fi
        
        local status=$(echo "$provider_info" | jq -r '.registrationState' 2>/dev/null || echo "Unknown")
        
        if [ "$status" == "Registered" ]; then
            log "$provider: Registered"
        else
            error "$provider: $status"
            unregistered+=("$provider")
            all_registered=false
        fi
    done
    
    echo
    
    if [ ${#not_found[@]} -gt 0 ]; then
        info "Note: Some provider namespaces were not found. This may be normal if:"
        info "  - The namespace name has changed"
        info "  - The service is not available in your subscription"
        info "  - The service requires special access"
    fi
    
    if [ "$all_registered" = false ] && [ ${#unregistered[@]} -gt 0 ]; then
        warn "Some AI/ML resource providers are not registered"
        info "Registering unregistered providers..."
        
        for provider in "${unregistered[@]}"; do
            info "Registering $provider..."
            if az provider register --namespace "$provider" --wait 2>/dev/null; then
                log "$provider: Registration initiated"
            else
                warn "Failed to register $provider (may require special permissions or not available)"
            fi
        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 "Unknown")
            if [ "$status" == "Registered" ]; then
                log "$provider: Now registered"
            else
                warn "$provider: Still $status (may take a few minutes or require approval)"
            fi
        done
    elif [ "$all_registered" = true ]; then
        log "All available AI/ML resource providers are registered"
    fi
}

# List all AI/ML related providers (discover available ones)
discover_ai_ml_providers() {
    section "Discovering AI/ML Related Providers"
    
    info "Searching for AI/ML related providers in subscription..."
    
    # Get all providers and filter for AI/ML related
    local all_providers=$(az provider list --query "[].namespace" -o tsv)
    
    local ai_ml_keywords=("cognitive" "machinelearning" "ml" "ai" "openai" "bot" "search" "synapse" "databricks" "datafactory")
    
    echo "Found AI/ML related providers:"
    for provider in $all_providers; do
        provider_lower=$(echo "$provider" | tr '[:upper:]' '[:lower:]')
        for keyword in "${ai_ml_keywords[@]}"; do
            if [[ "$provider_lower" == *"$keyword"* ]]; then
                local status=$(az provider show --namespace "$provider" --query "registrationState" -o tsv 2>/dev/null || echo "Unknown")
                if [ "$status" == "Registered" ]; then
                    log "$provider: Registered"
                else
                    warn "$provider: $status"
                fi
                break
            fi
        done
    done
}

# Check Preview Features
check_preview_features() {
    if [ ${#AI_ML_PREVIEW_FEATURES[@]} -eq 0 ]; then
        info "No specific AI/ML preview features configured"
        return
    fi
    
    section "AI/ML Preview Features"
    
    for feature in "${AI_ML_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 Azure OpenAI Service availability
check_openai_service() {
    section "Azure OpenAI Service Check"
    
    # Check if Cognitive Services provider is registered (required for OpenAI)
    local cognitive_status=$(az provider show --namespace "Microsoft.CognitiveServices" --query "registrationState" -o tsv 2>/dev/null || echo "NotRegistered")
    
    if [ "$cognitive_status" == "Registered" ]; then
        log "Microsoft.CognitiveServices is registered (required for Azure OpenAI)"
        
        # Try to list OpenAI resources (requires OpenAI access)
        info "Checking Azure OpenAI resource access..."
        local openai_resources=$(az cognitiveservices account list --query "[?kind=='OpenAI'].{Name:name, ResourceGroup:resourceGroup, Kind:kind}" -o json 2>/dev/null || echo "[]")
        
        if [ "$openai_resources" != "[]" ] && [ -n "$openai_resources" ]; then
            local count=$(echo "$openai_resources" | jq 'length')
            log "Found $count Azure OpenAI resource(s)"
            echo "$openai_resources" | jq -r '.[] | "  - \(.Name) in \(.ResourceGroup)"'
        else
            info "No Azure OpenAI resources found (may require access request)"
            info "To request Azure OpenAI access: https://aka.ms/oai/access"
        fi
    else
        warn "Microsoft.CognitiveServices is not registered (required for Azure OpenAI)"
        info "Registering Microsoft.CognitiveServices..."
        az provider register --namespace "Microsoft.CognitiveServices" --wait || warn "Registration may require approval"
    fi
}

# Check Azure Machine Learning
check_ml_service() {
    section "Azure Machine Learning Service Check"
    
    local ml_status=$(az provider show --namespace "Microsoft.MachineLearningServices" --query "registrationState" -o tsv 2>/dev/null || echo "NotRegistered")
    
    if [ "$ml_status" == "Registered" ]; then
        log "Microsoft.MachineLearningServices is registered"
        
        # Try to list ML workspaces
        info "Checking Azure ML workspace access..."
        local ml_workspaces=$(az ml workspace list --query "[].{Name:name, ResourceGroup:resourceGroup}" -o json 2>/dev/null || echo "[]")
        
        if [ "$ml_workspaces" != "[]" ] && [ -n "$ml_workspaces" ]; then
            local count=$(echo "$ml_workspaces" | jq 'length')
            log "Found $count Azure ML workspace(s)"
            echo "$ml_workspaces" | jq -r '.[] | "  - \(.Name) in \(.ResourceGroup)"'
        else
            info "No Azure ML workspaces found"
        fi
    else
        warn "Microsoft.MachineLearningServices is not registered"
        info "Registering Microsoft.MachineLearningServices..."
        az provider register --namespace "Microsoft.MachineLearningServices" --wait || warn "Registration may require approval"
    fi
}

# Main function
main() {
    log_info "AI/ML/OpenAI Resource Providers Check"
    log_info "======================================"
    echo
    
    check_azure_cli
    
    discover_ai_ml_providers
    
    check_ai_ml_providers
    
    check_openai_service
    
    check_ml_service
    
    check_preview_features
    
    echo
    section "Summary"
    log "AI/ML/OpenAI resource providers checked"
    log "Available providers discovered and verified"
    info "Note: Some services (like Azure OpenAI) may require additional access requests"
    info "Azure OpenAI access request: https://aka.ms/oai/access"
}

# Run main function
main "$@"