#!/bin/bash

#########################
# Helper Functions      #
#########################

# Logging Functions
log_info() { echo -e "\033[1;34m[INFO] $1\033[0m"; }     # Blue for info
log_success() { echo -e "\033[1;32m[SUCCESS] $1\033[0m"; } # Green for success
log_error() { echo -e "\033[1;31m[ERROR] $1\033[0m"; }   # Red for errors
log_warn() { echo -e "\033[1;33m[WARN] $1\033[0m"; }     # Yellow for warnings
log_verbose() {
    if [ "$VERBOSE" = true ]; then
        echo -e "\033[1;36m[VERBOSE] $1\033[0m" # Cyan for verbose/debug
    fi
}

# Display help
show_help() {
    log_info "Usage: eon-azure.sh [options]"
    echo ""
    echo "Options:"
    echo "  --tenant-id <tenant_id>          The Azure Entra ID (optional)."
    echo "  --project-id <project_id>        The project ID required for API request (required)."
    echo "  --connect-id <connect_id>        The connection ID required for API request (required)."
    echo "  --account-type <Source|Restore>  The type of account to onboard (required)."
    echo "  --location <location>            Azure region for resource deployment (required for Restore)."
    echo "  --subscription-id <subscription_id>  (Source only) Use this subscription non-interactively."
    echo "  --mg-id <management_group_id>       (Source only) Use this management group non-interactively."
    echo "  --backup-app-id <app_id>            Optional pre-created Scanning App ID."
    echo "  --verbose                        Enable verbose output."
    echo "  --help                           Show this help message and exit."
    echo ""
    log_info "Example usage:"
    echo "./eon-azure.sh --account-type Source --project-id <id> --connect-id <id> --subscription-id <sub-id> [--backup-app-id <app-id>]"
    echo "./eon-azure.sh --account-type Source --project-id <id> --connect-id <id> --mg-id <mg-id>"
    echo "./eon-azure.sh --account-type Restore --project-id <id> --connect-id <id> --location eastus"
}

# Parse command-line arguments
parse_arguments() {
    while [[ "$#" -gt 0 ]]; do
        case "$1" in
            --tenant-id)
                if [[ -n "$2" && ! "$2" =~ ^-- ]]; then
                    TENANT_ID="$2"
                    shift 2
                else
                    log_error "Error: --tenant-id requires a value."
                    show_help
                    exit 1
                fi
                ;;
            --project-id)
                if [[ -n "$2" && ! "$2" =~ ^-- ]]; then
                    PROJECT_ID="$2"
                    shift 2
                else
                    log_error "Error: --project-id requires a value."
                    show_help
                    exit 1
                fi
                ;;
            --connect-id)
                if [[ -n "$2" && ! "$2" =~ ^-- ]]; then
                    CONNECT_ID="$2"
                    shift 2
                else
                    log_error "Error: --connect-id requires a value."
                    show_help
                    exit 1
                fi
                ;;
            --account-type)
                if [[ -n "$2" && "$2" != --* ]]; then
                    ACCOUNT_TYPE=$(echo "$2" | tr '[:upper:]' '[:lower:]')  # Normalize to lowercase
                    if [[ "$ACCOUNT_TYPE" != "source" && "$ACCOUNT_TYPE" != "restore" ]]; then
                        log_error "Error: --account-type must be 'Source' or 'Restore'."
                        show_help
                        exit 1
                    fi
                    shift 2
                else
                    log_error "Error: --account-type requires a value (Source/Restore)."
                    show_help
                    exit 1
                fi
                ;;
            --location)
                if [[ -n "$2" && ! "$2" =~ ^-- ]]; then
                    LOCATION="$2"
                    shift 2
                else
                    log_error "Error: --location requires a value."
                    show_help
                    exit 1
                fi
                ;;
            --subscription-id)
                if [[ -n "$2" && ! "$2" =~ ^-- ]]; then
                    SUBSCRIPTION_ID="$2"
                    shift 2
                else
                    log_error "Error: --subscription-id requires a value."
                    show_help
                    exit 1
                fi
                ;;
            --mg-id)
                if [[ -n "$2" && ! "$2" =~ ^-- ]]; then
                    MG_ID="$2"
                    shift 2
                else
                    log_error "Error: --mg-id requires a value."
                    show_help
                    exit 1
                fi
                ;;
            --backup-app-id)
                if [[ -n "$2" && ! "$2" =~ ^-- ]]; then
                    SCANNING_APP_ID="$2"
                    shift 2
                else
                    log_error "Error: --backup-app-id requires a value."
                    show_help; exit 1
                fi
                ;;
            --verbose)
                VERBOSE=true
                shift
                ;;
            --help)
                HELP=true
                shift
                ;;
            *)
                log_error "Unknown parameter: $1"
                show_help
                exit 1
                ;;
        esac
    done
}

# Input validation function
validate_input() {
    # Convert account type to lowercase for case-insensitive validation
    local account_type_lower=$(echo "$ACCOUNT_TYPE" | tr '[:upper:]' '[:lower:]')

    if [[ "$account_type_lower" == "source" ]]; then
        ACCOUNT_TYPE="Source"
    elif [[ "$account_type_lower" == "restore" ]]; then
        ACCOUNT_TYPE="Restore"
    else
        log_error "Error: --account-type must be either 'Source' or 'Restore'."
        show_help
        exit 1
    fi

    # Validate common required parameters
    if [ -z "$PROJECT_ID" ]; then
        log_error "Error: --project-id is required."
        show_help
        exit 1
    fi

    if [ -z "$CONNECT_ID" ]; then
        log_error "Error: --connect-id is required."
        show_help
        exit 1
    fi

    # Validate Tenant ID format only if provided
    if [ -n "$TENANT_ID" ] && ! [[ "$TENANT_ID" =~ ^[0-9a-fA-F-]{36}$ ]]; then
        log_error "Error: Invalid Tenant ID format. It should be a valid GUID."
        exit 1
    fi

    # Validate Scanning App ID format if provided
    if [ -n "$SCANNING_APP_ID" ] && ! [[ "$SCANNING_APP_ID" =~ ^[0-9a-fA-F-]{36}$ ]]; then
        log_error "Error: --backup-app-id must be a valid GUID."
        exit 1
    fi

    # Validate subscription and management group IDs if provided
    if [ -n "$SUBSCRIPTION_ID" ] && ! [[ "$SUBSCRIPTION_ID" =~ ^[0-9a-fA-F-]{36}$ ]]; then
        log_error "Error: --subscription-id must be a valid GUID."
        exit 1
    fi

    # Prohibit specifying both subscription-id and mg-id together
    if [ -n "$SUBSCRIPTION_ID" ] && [ -n "$MG_ID" ]; then
        log_error "Error: Specify only one of --subscription-id or --mg-id."
        exit 1
    fi

    # Additional validation for Restore account type
    if [ "$ACCOUNT_TYPE" == "Restore" ]; then
        validate_restore_account
    fi
}

# Function to validate input specifically for Restore account type
validate_restore_account() {
    if [ -z "$LOCATION" ]; then
        log_error "Error: --location is required for Restore account type."
        show_help
        exit 1
    fi

    # Validate the location format (for example, basic validation to avoid empty spaces)
    if [[ ! "$LOCATION" =~ ^[a-zA-Z0-9]+$ ]]; then
        log_error "Error: Invalid location format."
        exit 1
    fi
}

# Check if Azure CLI is installed
check_az_cli_installed() {
    if command -v az &> /dev/null; then
        log_success "✔ Azure CLI is already installed."
    else
        log_warn "⚠ Azure CLI is not installed. Installing now..."

        if [[ "$OSTYPE" == "linux-gnu"* ]]; then
            curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash
        elif [[ "$OSTYPE" == "darwin"* ]]; then
            brew update && brew install azure-cli
        else
            log_error "Unsupported OS. Please install Azure CLI manually."
            exit 1
        fi

        if command -v az &> /dev/null; then
            log_success "✔ Azure CLI installation completed successfully."
        else
            log_error "❌ Azure CLI installation failed. Please install manually."
            exit 1
        fi
    fi
}

login_to_azure() {
    # Check for and save the current value of core.login_experience_v2 (if it exists)
    original_login_experience=$(az config get --output json --only-show-errors 2>/dev/null | jq -r '.core[] | select(.name=="login_experience_v2") | .value' 2>/dev/null)

    if [ -z "$original_login_experience" ] || [ "$original_login_experience" == "null" ]; then
        login_experience_not_set=true
    else
        login_experience_not_set=false
    fi

    # Set the property to off so that the subscription options aren't shown
    az config set core.login_experience_v2=off

    # Ensure cleanup on exit (restore configuration)
    restore_login_experience() {
        log_info "Restoring Azure CLI login experience configuration to its original state: $original_login_experience"
        if [ "$login_experience_not_set" = true ]; then
            az config unset core.login_experience_v2
        else
            az config set core.login_experience_v2="$original_login_experience"
        fi
    }
    trap restore_login_experience EXIT

    # Login
    if [ -n "$TENANT_ID" ]; then
        log_info "Logging into Azure Entra tenant: $TENANT_ID"
        
        # Attempt to login specifically to the given tenant
        if ! az login --tenant "$TENANT_ID" --output none 2>/dev/null; then
            log_error "❌ Azure login failed for tenant ID: $TENANT_ID. Please check your credentials and tenant ID."
            exit 1
        fi
        
        log_success "✔ Successfully logged into Azure tenant: $TENANT_ID"
        
    else
        log_info "Logging into Azure. You'll be prompted if necessary."
        if ! az login --output none 2>/dev/null; then
            log_error "❌ Azure login failed. Please check your credentials."
            exit 1
        fi
        
        log_success "✔ Successfully logged into Azure."
    fi

    # If TENANT_ID was not provided, determine it from available subscriptions
    if [ -z "$TENANT_ID" ]; then
        tenant_ids=$(az account list --query "[].tenantId" -o tsv 2>/dev/null | sort | uniq)
        tenant_array=()
        while IFS= read -r line; do
            tenant_array+=("$line")
        done <<< "$tenant_ids"

        tenant_count=${#tenant_array[@]}

        if [ "$tenant_count" -eq 0 ]; then
            log_error "❌ No tenants found in your Azure account."
            exit 1
        elif [ "$tenant_count" -eq 1 ]; then
            TENANT_ID="${tenant_array[0]}"
            log_info "Automatically selected the only available tenant: $TENANT_ID"
        else
            log_info "Multiple tenants found. Please select one to proceed."

            # Display a table of tenants
            table_data="No\tTenant ID"
            table_data+="\n------\t------------------------------------"
            for i in "${!tenant_array[@]}"; do
                table_data+="\n[$((i+1))]\t${tenant_array[$i]}"
            done

            echo -e "$table_data" | column -t -s $'\t'
            echo ""

            # Prompt the user to select a tenant
            while true; do
                read -p "Enter the tenant index to proceed: " user_input
                if [[ "$user_input" =~ ^[0-9]+$ ]] && [ "$user_input" -ge 1 ] && [ "$user_input" -le "$tenant_count" ]; then
                    TENANT_ID="${tenant_array[$((user_input-1))]}"
                    break
                else
                    log_warn "⚠ Invalid input. Please enter a number between 1 and $tenant_count."
                fi
            done

            log_success "✔ Selected tenant ID: $TENANT_ID"
            
            # Explicitly login to the selected tenant
            log_info "Logging into the selected tenant: $TENANT_ID"
            if ! az login --tenant "$TENANT_ID" --output none 2>/dev/null; then
                log_error "❌ Azure login failed for tenant ID: $TENANT_ID. Please check your credentials and tenant ID."
                exit 1
            fi
            log_success "✔ Successfully logged into Azure tenant: $TENANT_ID"
        fi
    fi

    trap - EXIT
    restore_login_experience
}

# Set the onboarding scope based on the provided subscription or management group ID
set_onboarding_scope() {
    if [ -n "$SUBSCRIPTION_ID" ]; then
        SCOPE="/subscriptions/$SUBSCRIPTION_ID"
        SCOPE_NAME="$SUBSCRIPTION_NAME"
    elif [ -n "$MG_ID" ]; then
        SCOPE="/providers/Microsoft.Management/managementGroups/$MG_ID"
        SCOPE_NAME="$MG_NAME"
    else
        log_error "❌ Neither subscription ID nor management group ID is set."
        exit 1
    fi
}

# Helper: Non-interactive subscription setup
noninteractive_subscription() {
    log_info "Non-interactive: using subscription $SUBSCRIPTION_ID"
    if ! az account set --subscription "$SUBSCRIPTION_ID" --output none; then
        log_error "Failed to set subscription: $SUBSCRIPTION_ID"
        exit 1
    fi

    # Retrieve subscription details and set name
    subscription_info=$(az account show --output json)
    SUBSCRIPTION_NAME=$(echo "$subscription_info" | jq -r '.name')
    log_success "✔ Successfully set subscription: $SUBSCRIPTION_NAME ($SUBSCRIPTION_ID)"

    # Set global scope variables
    SCOPE="/subscriptions/$SUBSCRIPTION_ID"
    SCOPE_NAME="$SUBSCRIPTION_NAME"
}

# Helper: Non-interactive management group validation
noninteractive_management_group() {
    log_info "Non-interactive: validating management group $MG_ID"
    if ! az account management-group show --name "$MG_ID" --output none; then
        log_error "Invalid management group ID: $MG_ID"
        exit 1
    fi
    MG_NAME=$(az account management-group show --name "$MG_ID" --query displayName -o tsv)
    log_success "✔ Validated management group: $MG_NAME ($MG_ID)"
}

select_subscription() {
    log_info "Fetching available subscriptions..."
    
    # Fetch the list of subscriptions
    subscriptions_json=$(az account list --query "[].{name:name, id:id, tenant:tenantId}" -o json 2>/dev/null)

    # Check if the JSON is empty or invalid
    if [ -z "$subscriptions_json" ] || [ "$subscriptions_json" == "[]" ]; then
        log_error "❌ No subscriptions found for this tenant."
        exit 1
    fi

    # Filter subscriptions by tenant ID if provided
    if [ -n "$TENANT_ID" ]; then
        subscriptions_json=$(echo "$subscriptions_json" | jq --arg tenant "$TENANT_ID" '[.[] | select(.tenant == $tenant)]')
        
        # Check if any subscriptions remain after filtering
        if [ -z "$subscriptions_json" ] || [ "$subscriptions_json" == "[]" ]; then
            log_error "❌ No subscriptions found for tenant ID: $TENANT_ID"
            exit 1
        fi
    fi

    # Convert JSON into an indexed array manually
    subscriptions=()
    while IFS= read -r line; do
        subscriptions+=("$line")
    done < <(echo "$subscriptions_json" | jq -c '.[]')

    # Auto-select if there's only one subscription
    if [ ${#subscriptions[@]} -eq 1 ]; then
        selected_subscription="${subscriptions[0]}"
        SUBSCRIPTION_NAME=$(echo "$selected_subscription" | jq -r '.name')
        SUBSCRIPTION_ID=$(echo "$selected_subscription" | jq -r '.id')
        log_info "Automatically selected the only available subscription: $SUBSCRIPTION_NAME ($SUBSCRIPTION_ID)"
        set_subscription
        return
    fi

    # Display subscriptions in a formatted table
    log_info "Available subscriptions:"
    # Prepare table data
    table_data="No\tSubscription Name\tSubscription ID\tTenant"
    table_data+="\n------\t------------------------------\t------------------------------------\t-------------------"
    index=1
    for sub in "${subscriptions[@]}"; do
        sub_name=$(echo "$sub" | jq -r '.name')
        sub_id=$(echo "$sub" | jq -r '.id')
        sub_tenant=$(echo "$sub" | jq -r '.tenant')

        table_data+="\n[$index]\t$sub_name\t$sub_id\t$sub_tenant"
        index=$((index + 1))
    done

    # Display subscriptions in a formatted table
    echo -e "$table_data" | column -t -s $'\t'

    echo ""

    # Loop until a valid selection is made
    while true; do
        read -p "Enter the subscription index or ID: " user_input

        # Check if input is an index
        if [[ "$user_input" =~ ^[0-9]+$ ]] && [ "$user_input" -ge 1 ] && [ "$user_input" -le "${#subscriptions[@]}" ]; then
            selected_subscription="${subscriptions[$((user_input-1))]}"
            SUBSCRIPTION_NAME=$(echo "$selected_subscription" | jq -r '.name')
            SUBSCRIPTION_ID=$(echo "$selected_subscription" | jq -r '.id')
            break
        fi

        # Check if input is a valid Subscription ID
        selected_subscription=$(echo "$subscriptions_json" | jq -c --arg sub_id "$user_input" '.[] | select(.id == $sub_id)')
        if [ -n "$selected_subscription" ]; then
            SUBSCRIPTION_NAME=$(echo "$selected_subscription" | jq -r '.name')
            SUBSCRIPTION_ID=$(echo "$selected_subscription" | jq -r '.id')
            break
        fi

        log_warn "⚠ Invalid selection. Please enter a valid index or subscription ID."
    done

    log_success "✔ Selected subscription: $SUBSCRIPTION_NAME ($SUBSCRIPTION_ID)"
    set_subscription
}

set_subscription() {
    log_info "Setting subscription to: $SUBSCRIPTION_ID"
    if ! az account set --subscription "$SUBSCRIPTION_ID" --output none 2>/dev/null; then
        log_error "❌ Failed to set subscription."
        exit 1
    fi
    log_success "✔ Successfully set subscription."
}

select_management_group() {
    log_info "Fetching available management groups..."

    # Fetch the list of management groups, including tenantId
    mg_json=$(az account management-group list \
        --query "[].{name:name, id:id, displayName:displayName, tenantId:tenantId}" \
        -o json 2>/dev/null)

    # Check if the JSON is empty or invalid
    if [ -z "$mg_json" ] || [ "$mg_json" == "[]" ]; then
        log_error "❌ No management groups found."
        exit 1
    fi

    # Convert JSON into an indexed array manually
    management_groups=()
    while IFS= read -r line; do
        management_groups+=("$line")
    done < <(echo "$mg_json" | jq -c '.[]')

    # Auto-select if there's only one management group
    if [ ${#management_groups[@]} -eq 1 ]; then
        selected_mg="${management_groups[0]}"
        MG_NAME=$(echo "$selected_mg" | jq -r '.displayName')
        MG_ID=$(echo "$selected_mg" | jq -r '.name')  # Extract just the GUID
        MG_SCOPE="/providers/Microsoft.Management/managementGroups/$MG_ID"  # Scope for role assignments
        MG_TENANT=$(echo "$selected_mg" | jq -r '.tenantId')
        log_info "Automatically selected the only available management group: $MG_NAME ($MG_ID), Tenant: $MG_TENANT"
        return
    fi

    # Prepare table data with a separator line (4 columns)
    table_data="No\tManagement Group Name\tManagement Group ID (GUID)\tTenant ID"
    table_data+="\n------\t------------------------------------\t------------------------------------\t------------------------------------"

    index=1
    for mg in "${management_groups[@]}"; do
        mg_name=$(echo "$mg" | jq -r '.displayName')
        mg_id=$(echo "$mg" | jq -r '.name')
        mg_tenant=$(echo "$mg" | jq -r '.tenantId')

        table_data+="\n[$index]\t$mg_name\t$mg_id\t$mg_tenant"
        index=$((index + 1))
    done

    # Display management groups in a formatted table
    echo -e "$table_data" | column -t -s $'\t'
    echo ""

    # Loop until a valid selection is made
    while true; do
        read -p "Enter the management group index or ID: " user_input

        # Check if input is an index
        if [[ "$user_input" =~ ^[0-9]+$ ]] && [ "$user_input" -ge 1 ] && [ "$user_input" -le "${#management_groups[@]}" ]; then
            selected_mg="${management_groups[$((user_input-1))]}"
            MG_NAME=$(echo "$selected_mg" | jq -r '.displayName')
            MG_ID=$(echo "$selected_mg" | jq -r '.name')
            MG_SCOPE="/providers/Microsoft.Management/managementGroups/$MG_ID"  # Scope for role assignments
            MG_TENANT=$(echo "$selected_mg" | jq -r '.tenantId')
            break
        fi

        # Check if input is a valid Management Group GUID
        selected_mg=$(echo "$mg_json" | jq -c --arg mg_id "$user_input" '.[] | select(.name == $mg_id)')
        if [ -n "$selected_mg" ]; then
            MG_NAME=$(echo "$selected_mg" | jq -r '.displayName')
            MG_ID=$(echo "$selected_mg" | jq -r '.name')
            MG_SCOPE="/providers/Microsoft.Management/managementGroups/$MG_ID"  
            MG_TENANT=$(echo "$selected_mg" | jq -r '.tenantId')
            break
        fi

        log_warn "⚠ Invalid selection. Please enter a valid index or management group GUID."
    done

    log_success "✔ Selected management group: $MG_NAME ($MG_ID), Tenant: $MG_TENANT"
}

handle_account_type_selection() {
    log_info "Determining selection process based on account type: $ACCOUNT_TYPE"

    if [ "$ACCOUNT_TYPE" == "Restore" ]; then
        log_info "Restore account type detected."
        
        # Ask the user if they want to select a subscription or a resource group
        while true; do
            echo "Would you like to select a Subscription or a Resource Group?"
            echo "1) Subscription"
            echo "2) Resource Group"
            read -p "Enter your choice (1 or 2): " selection

            case "$selection" in
                1)
                    log_info "You have chosen to select a subscription."
                    select_subscription
                    RESTORE_SCOPE="subscription"
                    break
                    ;;
                2)
                    log_info "You have chosen to select a resource group."
                    select_subscription
                    select_resource_group
                    RESTORE_SCOPE="resource_group"
                    break
                    ;;
                *)
                    log_warn "⚠ Invalid choice. Please enter 1 for Subscription or 2 for Resource Group."
                    ;;
            esac
        done
        return
    fi

    if [ "$ACCOUNT_TYPE" == "Source" ]; then
        log_info "Source account type detected."

        # Ask the user if they want to select a subscription or a management group
        while true; do
            echo "Would you like to select a Subscription or a Management Group?"
            echo "1) Subscription"
            echo "2) Management Group"
            read -p "Enter your choice (1 or 2): " selection

            case "$selection" in
                1)
                    log_info "You have chosen to select a subscription."
                    select_subscription
                    break
                    ;;
                2)
                    log_info "You have chosen to select a management group."
                    select_management_group
                    break
                    ;;
                *)
                    log_warn "⚠ Invalid choice. Please enter 1 for Subscription or 2 for Management Group."
                    ;;
            esac
        done
    else
        log_error "❌ Invalid ACCOUNT_TYPE: $ACCOUNT_TYPE. Expected 'Restore' or 'Source'."
        exit 1
    fi
}

confirm_selection() {
    local entity_type=""
    local entity_name=""
    local entity_id=""

    # Determine if we're confirming a subscription or management group
    if [ -n "$RESOURCE_GROUP_NAME" ]; then
        entity_type="resource group"
        entity_name="$RESOURCE_GROUP_NAME"
        entity_id="$RESOURCE_GROUP_NAME"
    elif [ -n "$SUBSCRIPTION_ID" ]; then
        entity_type="subscription"
        entity_name="$SUBSCRIPTION_NAME"
        entity_id="$SUBSCRIPTION_ID"
    elif [ -n "$MG_ID" ]; then
        entity_type="management group"
        entity_name="$MG_NAME"
        entity_id="$MG_ID"
    else
        log_error "❌ No subscription, management group, or resource group selected. Cannot proceed."
        exit 1
    fi

    # Confirmation prompt with retry on invalid input
    while true; do
        read -p "Do you want to proceed with onboarding $entity_type '$entity_name' (ID: $entity_id) as a $ACCOUNT_TYPE account? (yes/no): " confirmation
        confirmation=$(echo "$confirmation" | tr '[:upper:]' '[:lower:]')  # Convert input to lowercase

        case "$confirmation" in
            yes|y)
                log_success "✔ Proceeding with onboarding for $entity_type: $entity_name ($entity_id)"
                break
                ;;
            no|n)
                log_error "❌ Onboarding aborted by the user."
                exit 0
                ;;
            *)
                log_warn "⚠ Invalid input. Please enter 'yes' or 'no'."
                ;;
        esac
    done
}

check_service_principal_exists() {
    local app_id="$1"
    log_info "Checking if service principal exists for app ID: $app_id in tenant $TENANT_ID"

    # Run the az command and store the result
    output=$(az ad sp list --all --filter "appId eq '$app_id'" --output json 2>/dev/null)
    
    if [ $? -ne 0 ]; then
        log_error "❌ Failed to check service principal existence."
        exit 1
    fi

    # Count the number of SPs found
    local count
    count=$(echo "$output" | jq '. | length')

    if [[ "$count" -gt 0 ]]; then
        # Extract the first service principal ID
        service_principal_id=$(echo "$output" | jq -r '.[0].id')
        log_success "✔ Service principal exists with ID: $service_principal_id"
        return 0  # SP exists
    else
        log_info "Service principal not found for app ID: $app_id"
        return 1  # SP does not exist
    fi
}

create_service_principal() {
    local app_id="$1"
    log_info "Creating service principal for app ID: $app_id"

    output=$(az ad sp create --id "$app_id" --output json 2>&1)
    if [ $? -ne 0 ]; then
        log_error "❌ Failed to create service principal. Error: $output"
        exit 1
    fi

    service_principal_id=$(echo "$output" | jq -r '.id')

    if [ -z "$service_principal_id" ]; then
        log_error "❌ Failed to extract service principal ID from the response."
        exit 1
    fi

    log_success "✔ Successfully created service principal with ID: $service_principal_id"
}

create_service_principal_if_not_exists() {
    local app_id="$1"

    if check_service_principal_exists "$app_id"; then
        log_info "Service principal already exists. Skipping creation."
    else
        create_service_principal "$app_id"
    fi
}

assign_role_to_service_principal_in_scope() {
    local principal_id="$1"
    local role_name="$2"
    local scope="$3"
    local max_retries=4
    local retry_delay=15
    local attempt=1

    log_info "Assigning '$role_name' role to Service Principal (ID: $principal_id) at scope: $scope"

    while [ $attempt -le $max_retries ]; do
        log_info "Attempt $attempt/$max_retries: Assigning '$role_name' role to Service Principal (ID: $principal_id) at scope: $scope"
        
        az role assignment create \
            --assignee-object-id "$principal_id" \
            --assignee-principal-type ServicePrincipal \
            --role "$role_name" \
            --scope "$scope" \
            --output none 2>/dev/null

        if [ $? -eq 0 ]; then
            log_success "✔ Successfully assigned '$role_name' role to Service Principal (ID: $principal_id) at scope: $scope"
            return 0
        fi

        if [ $attempt -lt $max_retries ]; then
            log_warn "⚠️  Attempt $attempt failed. Retrying in ${retry_delay} seconds..."
            sleep $retry_delay
        fi

        attempt=$((attempt + 1))
    done

    log_error "❌ Failed to assign '$role_name' role to Service Principal after $max_retries attempts."
    exit 1
}

assign_subscription_reader_role_to_eon_sp() {
    log_info "Assigning 'Reader' role to Eon service principal..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Reader" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_subscription_reader_role_to_customer_scanning_sp() {
    log_info "Assigning 'Reader' role to Customer Scanning service principal..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Reader" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_subscription_contributor_role_to_customer_restore_sp() {
    log_info "Assigning 'Contributor' role to Customer Restore service principal..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Contributor" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_subscription_storage_blob_data_contributor_to_customer_scanning_sp() {
    log_info "Assigning 'Storage Blob Data Contributor' role to Customer Scanning service principal..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Storage Blob Data Contributor" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_subscription_storage_blob_data_contributor_to_customer_restore_sp() {
    log_info "Assigning 'Storage Blob Data Contributor' role to Customer Restore service principal..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Storage Blob Data Contributor" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_subscription_contributor_role_to_customer_scanning_sp() {
    log_info "Assigning 'Contributor' role to Customer Scanning service principal..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Contributor" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_subscription_storage_blob_data_reader_role_to_eon_sp() {
    log_info "Assigning 'Storage Blob Data Reader' role to Eon service principal..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Storage Blob Data Reader" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_resource_group_storage_blob_data_contributor_to_customer_restore_sp() {
    log_info "Assigning 'Storage Blob Data Contributor' role to Customer Restore service principal for resource group: $RESOURCE_GROUP_NAME..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Storage Blob Data Contributor" "/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESOURCE_GROUP_NAME"
}

assign_resource_group_storage_blob_data_contributor_to_customer_scanning_sp() {
    log_info "Assigning 'Storage Blob Data Contributor' role to Customer Restore service principal for resource group: $RESOURCE_GROUP_NAME..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Storage Blob Data Contributor" "/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESOURCE_GROUP_NAME"
}

assign_resource_group_contributor_to_customer_restore_sp() {
    log_info "Assigning 'Contributor' role to Customer Restore service principal for resource group: $RESOURCE_GROUP_NAME..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Contributor" "/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESOURCE_GROUP_NAME"
}

assign_resource_group_contributor_to_customer_scanning_sp() {
    log_info "Assigning 'Contributor' role to Customer Scanning service principal for resource group: $RESOURCE_GROUP_NAME..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Contributor" "/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESOURCE_GROUP_NAME"
}

assign_restore_resource_group_contributor_to_customer_scanning_sp() {
    log_info "Assigning 'Contributor' role to Customer Scanning service principal for restore resource group: $RESTORE_RESOURCE_GROUP_NAME..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Contributor" "/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESTORE_RESOURCE_GROUP_NAME"
}

print_role_json() {
    local name="$1"
    local description="$2"
    local assignable_scopes="$3"
    local actions="$4"
    local data_actions="$5"
    echo "{
        \"Name\": \"$name\",
        \"Description\": \"$description\",
        \"Actions\": $actions,
        \"DataActions\": $data_actions,
        \"AssignableScopes\": $assignable_scopes
    }"
}

check_custom_role_exists() {
    local role_name="$1"
    local scope="$2"

    log_info "Checking if the custom role '$role_name' exists..."

    # Check if the role exists
    existing_role=$(az role definition list --custom-role-only true --scope "$scope" --name "$role_name" --output json 2>/dev/null)

    if [ -z "$existing_role" ] || [ "$existing_role" == "[]" ]; then
        log_info "Custom role '$role_name' does not exist."
        return 1  # Role does not exist
    else
        log_success "✔ Custom role '$role_name' already exists."
        return 0  # Role exists
    fi
}

create_custom_role() {
    local name="$1"
    local description="$2"
    local assignable_scopes="$3"
    local actions="$4"
    local data_actions="$5"

    log_info "Creating custom role '$name'..."

    # Generate JSON and save to file
    print_role_json "$name" "$description" "$assignable_scopes" "$actions" "$data_actions" > role.json

    output=$(az role definition create --role-definition @role.json 2>&1)

    if [ $? -ne 0 ]; then
        log_error "❌ Failed to create custom role '$name'."
        log_error "Error: $output"
        rm -f role.json
        exit 1
    fi

    log_verbose "$output"

    log_success "✔ Successfully created custom role '$name'."
    rm -f role.json
}


create_or_update_custom_role() {
    local name="$1"
    local description="$2"
    local scope="$3"
    local actions_json="$4"
    local data_actions_json="$5"

    if check_custom_role_exists "$name" "$scope"; then
        log_info "Updating existing custom role '$name' on scope '$scope'..."

        print_role_json "$name" "$description" "[\"$scope\"]" "$actions_json" "$data_actions_json" > role.json

        output=$(az role definition update --role-definition @role.json 2>&1)

        if [ $? -ne 0 ]; then
            log_error "❌ Failed to update custom role '$name'."
            log_error "Error: $output"
            rm -f role.json
            exit 1
        fi

        log_verbose "$output"

        log_success "✔ Successfully updated custom role '$name'."
        rm -f role.json
    else
        create_custom_role "$name" "$description" "[\"$scope\"]" "$actions_json" "$data_actions_json"
    fi
}

create_eon_source_discovery_role() {
    log_info "Creating custom role for Eon Source Discovery if it does not exist"
    
    local scope=""
    local role_id=""

    if [ -n "$SUBSCRIPTION_ID" ]; then
        scope="/subscriptions/$SUBSCRIPTION_ID"
        role_id="$SUBSCRIPTION_ID"
    elif [ -n "$MG_ID" ]; then
        scope="/providers/Microsoft.Management/managementGroups/$MG_ID"
        role_id="$MG_ID"
    else
        log_error "❌ Neither subscription ID nor management group ID is set."
        exit 1
    fi

    eon_source_discovery_role_name="EonSourceDiscoveryRole-$role_id"

    EON_DISCOVERY_ROLE_DESCRIPTION="Custom role for Eon Source Discovery"

    # Define common discovery actions
    COMMON_DISCOVERY_ACTIONS=$(jq -nc '[
        "Microsoft.Resources/subscriptions/resources/read",
        "Microsoft.Compute/virtualMachines/read",
        "Microsoft.Compute/virtualMachines/instanceView/read",
        "Microsoft.Compute/disks/read",
        "Microsoft.Storage/storageAccounts/read",
        "Microsoft.Storage/storageAccounts/blobServices/containers/read",
        "Microsoft.Storage/storageAccounts/fileServices/shares/read",
        "Microsoft.Sql/servers/databases/read",
        "Microsoft.Sql/managedInstances/read",
        "Microsoft.DBforPostgreSQL/flexibleServers/read",
        "Microsoft.DBforMySQL/flexibleServers/read",
        "Microsoft.SqlVirtualMachine/sqlVirtualMachines/read",
        "Microsoft.Compute/snapshots/read",
        "Microsoft.RecoveryServices/vaults/read",
        "Microsoft.RecoveryServices/vaults/backupProtectedItems/read",
        "Microsoft.RecoveryServices/vaults/backupFabrics/protectionContainers/protectedItems/read",
        "Microsoft.RecoveryServices/vaults/backupFabrics/protectionContainers/protectedItems/recoveryPoints/read",
        "Microsoft.DataProtection/backupVaults/read",
        "Microsoft.DataProtection/backupVaults/backupInstances/read",
        "Microsoft.DataProtection/backupVaults/backupInstances/recoveryPoints/read"
    ]')
    if [ -n "$MG_ID" ]; then
        EON_DISCOVERY_ROLE_ACTIONS_JSON=$(echo "$COMMON_DISCOVERY_ACTIONS" | jq '. + ["Microsoft.Management/managementGroups/read"]')
    else
        EON_DISCOVERY_ROLE_ACTIONS_JSON="$COMMON_DISCOVERY_ACTIONS"
    fi

    create_or_update_custom_role "$eon_source_discovery_role_name" "$EON_DISCOVERY_ROLE_DESCRIPTION" "$scope" "$EON_DISCOVERY_ROLE_ACTIONS_JSON" "[]"
}

create_eon_source_backup_role() {
    log_info "Creating custom role for Eon Source Backup if it does not exist"
    
    local scope=""
    local role_id=""

    if [ -n "$SUBSCRIPTION_ID" ]; then
        scope="/subscriptions/$SUBSCRIPTION_ID"
        role_id="$SUBSCRIPTION_ID"
    elif [ -n "$MG_ID" ]; then
        scope="/providers/Microsoft.Management/managementGroups/$MG_ID"
        role_id="$MG_ID"
    else
        log_error "❌ Neither subscription ID nor management group ID is set."
        exit 1
    fi

    eon_source_backup_role_name="EonSourceBackupRole-$role_id"

    EON_BACKUP_ROLE_DESCRIPTION="Custom role for Eon Source Backup"

    # Define backup actions
    EON_BACKUP_ROLE_ACTIONS_JSON=$(jq -nc '[
        "Microsoft.Compute/virtualMachines/read",
        "Microsoft.Compute/virtualMachines/write",
        "Microsoft.Compute/virtualMachines/instanceView/read",
        "Microsoft.Compute/disks/read",
        "Microsoft.Compute/disks/write",
        "Microsoft.Compute/disks/delete",
        "Microsoft.Compute/disks/beginGetAccess/action",
        "Microsoft.Compute/disks/endGetAccess/action",
        "Microsoft.Compute/restorePointCollections/read",
        "Microsoft.Compute/restorePointCollections/write",
        "Microsoft.Compute/restorePointCollections/delete",
        "Microsoft.Compute/restorePointCollections/restorePoints/read",
        "Microsoft.Compute/restorePointCollections/restorePoints/write",
        "Microsoft.Compute/restorePointCollections/restorePoints/delete",
        "Microsoft.Sql/servers/read",
        "Microsoft.Sql/servers/databases/read",
        "Microsoft.Sql/servers/databases/write",
        "Microsoft.Sql/managedInstances/read",
        "Microsoft.Sql/managedInstances/databases/read",
        "Microsoft.Sql/managedInstances/databases/write",
        "Microsoft.DBforPostgreSQL/flexibleServers/read",
        "Microsoft.DBforPostgreSQL/flexibleServers/databases/read",
        "Microsoft.DBforPostgreSQL/flexibleServers/databases/write",
        "Microsoft.SqlVirtualMachine/sqlVirtualMachines/read",
        "Microsoft.SqlVirtualMachine/sqlVirtualMachines/write",
        "Microsoft.DBforMySQL/flexibleServers/read",
        "Microsoft.DBforMySQL/flexibleServers/databases/read",
        "Microsoft.DBforMySQL/flexibleServers/databases/write",
        "Microsoft.DocumentDB/databaseAccounts/read",
        "Microsoft.DocumentDB/databaseAccounts/write",
        "Microsoft.Network/virtualNetworks/read",
        "Microsoft.Network/virtualNetworks/subnets/read",
        "Microsoft.Network/networkInterfaces/read",
        "Microsoft.Network/networkInterfaces/write",
        "Microsoft.Network/privateEndpoints/read",
        "Microsoft.Network/privateEndpoints/write",
        "Microsoft.Storage/storageAccounts/privateEndpointConnections/read",
        "Microsoft.Storage/storageAccounts/privateEndpointConnections/write",
        "Microsoft.Storage/storageAccounts/blobServices/containers/read",
        "Microsoft.Storage/storageAccounts/fileServices/shares/read"
    ]')
    # Define backup data actions
    EON_BACKUP_ROLE_DATA_ACTIONS_JSON=$(jq -nc '[
        "Microsoft.Storage/storageAccounts/blobServices/containers/blobs/read",
        "Microsoft.Storage/storageAccounts/blobServices/containers/blobs/write",
        "Microsoft.Storage/storageAccounts/blobServices/containers/blobs/add/action",
        "Microsoft.Storage/storageAccounts/blobServices/containers/blobs/delete",
        "Microsoft.Storage/storageAccounts/fileServices/fileshares/files/read"
    ]')

    create_or_update_custom_role "$eon_source_backup_role_name" \
        "$EON_BACKUP_ROLE_DESCRIPTION" \
        "$scope" \
        "$EON_BACKUP_ROLE_ACTIONS_JSON" \
        "$EON_BACKUP_ROLE_DATA_ACTIONS_JSON"
}

# Generic assign functions for custom roles
assign_eon_source_discovery_role_to_service_principal_in_scope() {
    assign_role_to_service_principal_in_scope "$service_principal_id" "$eon_source_discovery_role_name" "$SCOPE"
}

assign_eon_source_backup_role_to_service_principal_in_scope() {
    assign_role_to_service_principal_in_scope "$service_principal_id" "$eon_source_backup_role_name" "$SCOPE"
}

send_source_api_request() {
    log_info "Sending API request to source account onboarding endpoint..."
    
    if [ -n "$MG_ID" ]; then
        send_source_mg_api_request
    else
        send_source_subscription_api_request
    fi
}

assign_eon_source_role_to_customer_scanning_sp_on_subscription() {
    log_info "Assigning Eon Source Role to Service Principal (ID: $service_principal_id) at scope: /subscriptions/$SUBSCRIPTION_ID"

    # Validate required variables
    if [ -z "$service_principal_id" ] || [ -z "$SUBSCRIPTION_ID" ]; then
        log_error "❌ Service Principal ID or Subscription ID is missing."
        exit 1
    fi

    assign_role_to_service_principal_in_scope "$service_principal_id" "$eon_source_role_name" "/subscriptions/$SUBSCRIPTION_ID"
}

assign_eon_source_role_to_customer_scanning_sp_on_mg() {
    log_info "Assigning Eon Source Role to Service Principal (ID: $service_principal_id) at scope: $MG_SCOPE"

    # Validate required variables
    if [ -z "$service_principal_id" ] || [ -z "$MG_SCOPE" ]; then
        log_error "❌ Service Principal ID or Management Group ID is missing."
        exit 1
    fi

    assign_role_to_service_principal_in_scope "$service_principal_id" "$eon_source_role_name" "$MG_SCOPE"
}

assign_management_group_reader_role_to_eon_sp() {
    log_info "Assigning 'Reader' role to Eon service principal for Management Group: $MG_NAME ($MG_ID)..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Reader" "$MG_SCOPE"
}

assign_management_group_mg_reader_role_to_eon_sp() {
    log_info "Assigning 'Management Group Reader' role to Eon service principal for Management Group: $MG_NAME ($MG_ID)..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Management Group Reader" "$MG_SCOPE"
}

assign_management_group_reader_role_to_customer_scanning_sp() {
    log_info "Assigning 'Reader' role to Customer Scanning service principal for Management Group: $MG_NAME ($MG_ID)..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Reader" "$MG_SCOPE"
}

assign_management_group_mg_reader_role_to_customer_scanning_sp() {
    log_info "Assigning 'Management Group Reader' role to Customer Scanning service principal for Management Group: $MG_NAME ($MG_ID)..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Management Group Reader" "$MG_SCOPE"
}

assign_management_group_blob_data_contributor_to_customer_scanning_sp() {
    log_info "Assigning 'Storage Blob Data Contributor' role to Customer Scanning service principal for Management Group: $MG_NAME ($MG_ID)..."

    assign_role_to_service_principal_in_scope "$service_principal_id" "Storage Blob Data Contributor" "$MG_SCOPE"
}

check_resource_group_exists() {
    local rg_name="$1"

    log_info "Checking if resource group '$rg_name' exists..."

    local output
    output=$(az group exists --name "$rg_name" --output tsv 2>/dev/null)

    if [ "$output" == "true" ]; then
        log_success "✔ Resource group '$rg_name' already exists."
        return 0  # Resource group exists
    else
        log_info "Resource group '$rg_name' does not exist."
        return 1  # Resource group does not exist
    fi
}

create_resource_group() {
    local rg_name="$1"
    local location="${2}"

    log_info "Creating resource group '$rg_name' in location '$location'..."

    az group create --name "$rg_name" --location "$location" --output none 2>/dev/null

    if [ $? -ne 0 ]; then
        log_error "❌ Failed to create resource group '$rg_name'."
        exit 1
    fi

    log_success "✔ Successfully created resource group '$rg_name' in location '$location'."
}

# Function to create a resource group if it does not exist
create_resource_group_if_not_exists() {
    local rg_name="$1"
    local location="${2}"  # Default location to 'eastus' if not provided

    if check_resource_group_exists "$rg_name"; then
        log_info "Skipping creation. Resource group '$rg_name' already exists."
    else
        create_resource_group "$rg_name" "$location"
    fi
}

create_restore_resource_group() {
    # Generate the restore resource group name based on the PROJECT_ID
    export RESTORE_RESOURCE_GROUP_NAME="eon-restore-${PROJECT_ID:0:8}"

    log_info "Creating restore resource group if it does not exist: $RESTORE_RESOURCE_GROUP_NAME"

    # Call the function to create the resource group if it does not exist
    create_resource_group_if_not_exists "$RESTORE_RESOURCE_GROUP_NAME" "$LOCATION"
}

# Function to check if a managed identity exists
check_managed_identity_exists() {
    local identity_name="$1"
    local resource_group="$2"

    log_info "Checking if managed identity '$identity_name' exists in resource group '$resource_group'..."

    output=$(az identity show --name "$identity_name" --resource-group "$resource_group" --output json 2>/dev/null)

    if [ -n "$output" ]; then
        log_success "✔ Managed identity '$identity_name' exists."
        return 0  # Managed identity exists
    else
        log_info "Managed identity '$identity_name' does not exist."
        return 1  # Managed identity does not exist
    fi
}

# Function to create a managed identity
create_managed_identity() {
    local identity_name="$1"
    local resource_group="$2"
    local location="$3"
    
    log_info "Creating managed identity '$identity_name' in resource group '$resource_group' in location '$location'..."

    output=$(az identity create --name "$identity_name" --resource-group "$resource_group" --location "$location" --output json 2>/dev/null)

    if [ $? -ne 0 ]; then
        log_error "❌ Failed to create managed identity '$identity_name'."
        exit 1
    fi

    log_success "✔ Successfully created managed identity '$identity_name'."

    log_verbose "$output"
}

# Function to create a managed identity if it does not exist
create_managed_identity_if_not_exists() {
    local identity_name="$1"
    local resource_group="$2"
    local location="$3"

    if check_managed_identity_exists "$identity_name" "$resource_group"; then
        log_info "Managed identity '$identity_name' already exists. Skipping creation."
    else
        create_managed_identity "$identity_name" "$resource_group" "$location"
    fi
}

# Function to create the restore managed identity
create_restore_managed_identity() {
    # Generate the managed identity name using the first 8 characters of the project ID
    RESTORE_MANAGED_IDENTITY_NAME="eon-restore-identity-${PROJECT_ID:0:8}"

    log_info "Ensuring the Restore managed identity '$RESTORE_MANAGED_IDENTITY_NAME' exists in resource group '$RESTORE_RESOURCE_GROUP_NAME'..."

    create_managed_identity_if_not_exists "$RESTORE_MANAGED_IDENTITY_NAME" "$RESTORE_RESOURCE_GROUP_NAME" "$LOCATION"

    # Extract the required properties from the global output variable
    RESTORE_MI_PRINCIPAL_ID=$(echo "$output" | jq -r '.principalId')
    RESTORE_MI_CLIENT_ID=$(echo "$output" | jq -r '.clientId')

    if [ -z "$RESTORE_MI_PRINCIPAL_ID" ] || [ -z "$RESTORE_MI_CLIENT_ID" ]; then
        log_error "❌ Failed to retrieve managed identity details after creation."
        exit 1
    fi

    log_success "✔ Restore managed identity ensured. Principal ID: $RESTORE_MI_PRINCIPAL_ID, Client ID: $RESTORE_MI_CLIENT_ID"
}

check_app_registration_exists() {
    local app_name="$1"
    log_info "Checking if app registration '$app_name' exists..."

    output=$(az ad app list --display-name "$app_name" -o json 2>/dev/null)

    if [ -z "$output" ] || [ "$output" == "[]" ]; then
        log_info "App registration '$app_name' does not exist."
        return 1  # App does not exist
    else
        log_success "✔ App registration '$app_name' already exists."
        return 0  # App exists
    fi
}

create_app_registration() {
    local app_name="$1"
    local is_multi_tenant="$2"

    log_info "Creating app registration '$app_name'..."

    local sign_in_audience="AzureADMyOrg"
    if [ "$is_multi_tenant" == "true" ]; then
        sign_in_audience="AzureADMultipleOrgs"
    fi

    output=$(az ad app create --display-name "$app_name" --sign-in-audience "$sign_in_audience" --output json 2>/dev/null)

    if [ $? -ne 0 ]; then
        log_error "❌ Failed to create app registration '$app_name'."
        exit 1
    fi

    log_success "✔ Successfully created app registration '$app_name'."
}

create_app_registration_if_not_exists() {
    local app_name="$1"
    local is_multi_tenant="$2"

    if check_app_registration_exists "$app_name"; then
        log_info "App registration '$app_name' already exists. Skipping creation."
    else
        create_app_registration "$app_name" "$is_multi_tenant"
    fi
}

create_restore_app_registration() {
    local app_name="eon-restore-app-${PROJECT_ID:0:8}"
    log_info "Ensuring restore app registration exists with name '$app_name'..."

    create_app_registration_if_not_exists "$app_name" "true"

    # Check if the output is an array (list command)
    if echo "$output" | jq -e '.[0]' &>/dev/null; then
        RESTORE_APP_ID=$(echo "$output" | jq -r '.[0].appId')
    else
        RESTORE_APP_ID=$(echo "$output" | jq -r '.appId')
    fi

    if [ -z "$RESTORE_APP_ID" ] || [ "$RESTORE_APP_ID" == "null" ]; then
        log_error "❌ Failed to retrieve application ID for restore app registration."
        exit 1
    fi

    log_success "✔ Restore app registration created with Application ID: $RESTORE_APP_ID"
}

create_restore_app_service_principal() {
    log_info "Ensuring service principal exists for restore app with Application ID: $RESTORE_APP_ID"

    create_service_principal_if_not_exists "$RESTORE_APP_ID"

    if [ -z "$service_principal_id" ]; then
        log_error "❌ Failed to retrieve Service Principal ID for restore app."
        exit 1
    fi

    log_success "✔ Service principal created with ID: $service_principal_id for restore app."
}

create_customer_scanning_app_service_principal() {
    log_info "Ensuring service principal exists for Scanning app with Application ID: $SCANNING_APP_ID"

    create_service_principal_if_not_exists "$SCANNING_APP_ID"

    if [ -z "$service_principal_id" ]; then
        log_error "❌ Failed to retrieve Service Principal ID for Scanning app."
        exit 1
    fi

    log_success "✔ Service principal created with ID: $service_principal_id for Scanning app."
}

check_federated_identity_exists() {
    local app_id="$1"
    local federated_credential_name="$2"

    log_info "Checking if federated identity credential '$federated_credential_name' exists for app ID: $app_id"

    output=$(az ad app federated-credential list --id "$app_id" --output json 2>/dev/null)

    if [ $? -ne 0 ]; then
        log_error "❌ Failed to check federated identity credential existence."
        exit 1
    fi

    local count
    count=$(echo "$output" | jq --arg name "$federated_credential_name" '[.[] | select(.name == $name)] | length')

    if [[ "$count" -gt 0 ]]; then
        log_success "✔ Federated identity credential '$federated_credential_name' already exists."
        return 0  # Credential exists
    else
        log_info "Federated identity credential '$federated_credential_name' does not exist."
        return 1  # Credential does not exist
    fi
}

create_federated_identity() {
    local app_id="$1"
    local federated_credential_name="$2"
    local subject="$3"
    local description="$4"

    log_info "Creating federated identity credential '$federated_credential_name' for app ID: $app_id"

    output=$(az ad app federated-credential create --id "$app_id" --parameters "{
        \"name\": \"$federated_credential_name\",
        \"issuer\": \"https://login.microsoftonline.com/$TENANT_ID/v2.0\",
        \"subject\": \"$subject\",
        \"audiences\": [\"api://AzureADTokenExchange\"],
        \"description\": \"$description\"
    }" --output json 2>&1)

    if [ $? -ne 0 ]; then
        log_error "❌ Failed to create federated identity credential. Error: $output"
        exit 1
    fi

    log_success "✔ Successfully created federated identity credential '$federated_credential_name'"
}

create_federated_identity_if_not_exists() {
    local app_id="$1"
    local federated_credential_name="$2"
    local subject="$3"
    local description="$4"

    if check_federated_identity_exists "$app_id" "$federated_credential_name"; then
        log_info "Federated identity credential '$federated_credential_name' already exists. Skipping creation."
    else
        create_federated_identity "$app_id" "$federated_credential_name" "$subject" "$description"
    fi
}

federate_eon_restore_identity() {
    # Generate the federated credential name
    local federated_credential_name="eon-restore-credential-$RESTORE_MI_PRINCIPAL_ID"

    log_info "Federating managed identity (Principal ID: $RESTORE_MI_PRINCIPAL_ID) for Eon with app ID: $RESTORE_APP_ID"

    create_federated_identity_if_not_exists "$RESTORE_APP_ID" "$federated_credential_name" "$RESTORE_MI_PRINCIPAL_ID" "Federation for Eon restore process"

    log_success "✔ Successfully federated managed identity (Principal ID: $RESTORE_MI_PRINCIPAL_ID) for Eon with app ID: $RESTORE_APP_ID"
}

send_create_customer_scanning_multi_tenant_app_request() {
    local url="$BASE_API_URL/projects/$PROJECT_ID/azure/account/connect/$CONNECT_ID/createscanningapp"
    
    log_info "Sending request to create scanning multi-tenant app at: $url"

    # Make API request (POST with no body)
    response=$(curl -s -X POST "$url" \
        -H "Content-Type: application/json" \
        --fail 2>/dev/null)

    # Check if request was successful
    if [ $? -ne 0 ] || [ -z "$response" ]; then
        log_error "❌ Failed to create scanning multi-tenant app. No response or request error."
        exit 1
    fi

    # Extract values from response into meaningful parameters
    SCANNING_APP_ID=$(echo "$response" | jq -r '.appId')

    # Validate extracted values
    if [ -z "$SCANNING_APP_ID" ] || [ "$SCANNING_APP_ID" == "null" ]; then
        log_error "❌ API response did not contain valid appId."
        exit 1
    fi

    log_success "✔ Successfully created scanning multi-tenant app. Scanning App ID: $SCANNING_APP_ID."
}

send_connect_api_request() {
    local request_payload="$1"
    local api_url="$BASE_API_URL/projects/$PROJECT_ID/azure/account/connect/$CONNECT_ID"

    log_info "Sending API request to: $api_url"

    # Print the payload for debugging if verbose mode is enabled
    log_verbose "Request URL: $api_url"
    log_verbose "Request Payload: $request_payload"

    # Send the HTTP request using curl and capture response and HTTP status
    response=$(curl -s -w "\n%{http_code}" -X POST "$api_url" \
        -H "Content-Type: application/json" \
        -d "$request_payload")

    # Extract response body and HTTP status code
    http_status=$(echo "$response" | tail -n 1)
    response_body=$(echo "$response" | sed '$d')  # Remove last line (status code)

    # Print response if verbose mode is enabled
    log_verbose "Response Body: $response_body"
    log_verbose "HTTP Status Code: $http_status"

    # Validate the response status
    if [[ "$http_status" == "200" || "$http_status" == "201" ]]; then
        log_success "✔ API request successful."
        echo "Response: $response_body"
    else
        log_error "❌ API request failed with status code $http_status"
        log_error "Response: $response_body"
        exit 1
    fi
}

send_source_subscription_api_request() {
    local request_payload=$(jq -n --arg accountType "Source" \
                                 --arg tenantId "$TENANT_ID" \
                                 --arg subscriptionId "$SUBSCRIPTION_ID" \
                                 '{
                                      accountType: $accountType,
                                      tenantId: $tenantId,
                                      subscriptionId: $subscriptionId
                                  }')

    log_info "Sending Connect Source Account API request..."
    send_connect_api_request "$request_payload"
}

send_source_mg_api_request() {
    local request_payload=$(jq -n --arg accountType "Source" \
                                 --arg tenantId "$TENANT_ID" \
                                 --arg mgId "$MG_ID" \
                                 '{
                                      accountType: $accountType,
                                      tenantId: $tenantId,
                                      managementGroupId: $mgId
                                  }')

    log_info "Sending Connect Source Account API request..."
    send_connect_api_request "$request_payload"
}

send_restore_api_request() {
    local request_payload=$(jq -n --arg accountType "Restore" \
                                 --arg tenantId "$TENANT_ID" \
                                 --arg subscriptionId "$SUBSCRIPTION_ID" \
                                 --arg applicationId "$RESTORE_APP_ID" \
                                 --arg location "$LOCATION" \
                                 --arg resourceGroupName "$RESOURCE_GROUP_NAME" \
                                 '{
                                      accountType: $accountType,
                                      tenantId: $tenantId,
                                      location: $location,
                                      subscriptionId: $subscriptionId,
                                      applicationId: $applicationId,
                                      resourceGroupName: $resourceGroupName
                                  }')

    log_info "Sending Connect Restore Account API request..."
    send_connect_api_request "$request_payload"
}

send_get_scanning_app_request() {
    local url="$BASE_API_URL/projects/$PROJECT_ID/azure/account/connect/$CONNECT_ID/getscanningapp"
    
    log_info "Sending request to get scanning app at: $url"

    # Make API request (GET with no body)
    response=$(curl -s -X GET "$url" \
        -H "Content-Type: application/json" \
        --fail 2>/dev/null)

    # Check if request was successful
    if [ $? -ne 0 ] || [ -z "$response" ]; then
        log_error "❌ Failed to get scanning app. No response or request error."
        exit 1
    fi

    # Extract values from response into meaningful parameters
    SCANNING_APP_ID=$(echo "$response" | jq -r '.appId')

    # Validate extracted values
    if [ -z "$SCANNING_APP_ID" ] || [ "$SCANNING_APP_ID" == "null" ]; then
        log_error "❌ API response did not contain valid appId."
        exit 1
    fi

    log_success "✔ Successfully retrieved scanning app. Scanning App ID: $SCANNING_APP_ID."
}

select_resource_group() {
    log_info "Fetching resource groups in subscription: $SUBSCRIPTION_NAME ($SUBSCRIPTION_ID)..."
    
    # Fetch the list of resource groups
    rgs_json=$(az group list --subscription "$SUBSCRIPTION_ID" --query "[].{name:name, location:location}" -o json 2>/dev/null)

    # Check if the JSON is empty or invalid
    if [ -z "$rgs_json" ] || [ "$rgs_json" == "[]" ]; then
        log_error "❌ No resource groups found in this subscription."
        exit 1
    fi

    # Convert JSON into an indexed array manually
    resource_groups=()
    while IFS= read -r line; do
        resource_groups+=("$line")
    done < <(echo "$rgs_json" | jq -c '.[]')

    # Display resource groups in a formatted table
    log_info "Available resource groups:"
    table_data="No\tResource Group Name\tLocation"
    table_data+="\n------\t------------------------------\t------------------------------"
    index=1
    for rg in "${resource_groups[@]}"; do
        rg_name=$(echo "$rg" | jq -r '.name')
        rg_location=$(echo "$rg" | jq -r '.location')
        table_data+="\n[$index]\t$rg_name\t$rg_location"
        index=$((index + 1))
    done

    echo -e "$table_data" | column -t -s $'\t'
    echo ""

    # Loop until a valid selection is made
    while true; do
        read -p "Enter the resource group index or name: " user_input

        # Check if input is an index
        if [[ "$user_input" =~ ^[0-9]+$ ]] && [ "$user_input" -ge 1 ] && [ "$user_input" -le "${#resource_groups[@]}" ]; then
            selected_rg="${resource_groups[$((user_input-1))]}"
            export RESOURCE_GROUP_NAME=$(echo "$selected_rg" | jq -r '.name')
            export RESOURCE_GROUP_LOCATION=$(echo "$selected_rg" | jq -r '.location')
            break
        fi

        # Check if input is a valid Resource Group name
        selected_rg=$(echo "$rgs_json" | jq -c --arg rg_name "$user_input" '.[] | select(.name == $rg_name)')
        if [ -n "$selected_rg" ]; then
            export RESOURCE_GROUP_NAME=$(echo "$selected_rg" | jq -r '.name')
            export RESOURCE_GROUP_LOCATION=$(echo "$selected_rg" | jq -r '.location')
            break
        fi

        log_warn "⚠ Invalid selection. Please enter a valid index or resource group name."
    done

    log_success "✔ Selected resource group: $RESOURCE_GROUP_NAME in location: $RESOURCE_GROUP_LOCATION"
}

# Handle source onboarding for a given scope (subscription or management group)
handle_source_scope() {
    log_info "Executing Source account onboarding for: $SCOPE_NAME ($SCOPE)"
    
    if [ -n "$SCANNING_APP_ID" ]; then
        log_info "Using provided scanning app ID: $SCANNING_APP_ID"
    else
        send_create_customer_scanning_multi_tenant_app_request
    fi

    create_customer_scanning_app_service_principal

    # Assign custom roles
    create_eon_source_discovery_role
    assign_eon_source_discovery_role_to_service_principal_in_scope
    create_eon_source_backup_role
    assign_eon_source_backup_role_to_service_principal_in_scope

    send_source_api_request
}

handle_restore_common_steps() {
    create_restore_resource_group
    create_restore_managed_identity
    create_restore_app_registration
    create_restore_app_service_principal
    federate_eon_restore_identity
}

handle_restore_subscription() {
    log_info "Executing Restore account onboarding flow for subscription scope: $SUBSCRIPTION_NAME ($SUBSCRIPTION_ID)"
    
    handle_restore_common_steps

    assign_subscription_contributor_role_to_customer_restore_sp 
    assign_subscription_storage_blob_data_contributor_to_customer_restore_sp

    send_get_scanning_app_request
    create_customer_scanning_app_service_principal
    assign_subscription_contributor_role_to_customer_scanning_sp
    assign_subscription_storage_blob_data_contributor_to_customer_scanning_sp

    send_restore_api_request
}

handle_restore_resource_group() {
    log_info "Executing Restore account onboarding flow for resource group scope: $RESOURCE_GROUP_NAME in subscription: $SUBSCRIPTION_NAME ($SUBSCRIPTION_ID)"
    
    handle_restore_common_steps

    assign_resource_group_contributor_to_customer_restore_sp
    assign_resource_group_storage_blob_data_contributor_to_customer_restore_sp

    send_get_scanning_app_request
    create_customer_scanning_app_service_principal
    assign_resource_group_contributor_to_customer_scanning_sp
    assign_restore_resource_group_contributor_to_customer_scanning_sp
    assign_resource_group_storage_blob_data_contributor_to_customer_scanning_sp

    send_restore_api_request
}

#########################
# Main Script Flow      #
#########################

main() {
    # Default values
    PROJECT_ID="${PROJECT_ID:-}"
    CONNECT_ID="${CONNECT_ID:-}"
    TENANT_ID="${TENANT_ID:-}"
    SUBSCRIPTION_ID="${SUBSCRIPTION_ID:-}"
    SUBSCRIPTION_NAME="${SUBSCRIPTION_NAME:-}"
    MG_ID="${MG_ID:-}"
    MG_NAME="${MG_NAME:-}"
    MG_SCOPE="${MG_SCOPE:-}"
    ACCOUNT_TYPE="${ACCOUNT_TYPE:-}"
    SCANNING_APP_ID="${SCANNING_APP_ID:-}"
    VERBOSE="${VERBOSE:-false}"
    HELP="${HELP:-false}"
    LOCATION="${LOCATION:-}"
    RESTORE_RESOURCE_GROUP_NAME="${RESTORE_RESOURCE_GROUP_NAME:-}"
    RESTORE_SCOPE="${RESTORE_SCOPE:-}"
    BASE_API_URL="${BASE_API_URL:-"https://console.eon.io/api"}"

    # Parse command-line arguments
    parse_arguments "$@"

    # Show help if requested
    if [ "$HELP" = true ]; then
        show_help
        exit 0
    fi

    # Validate required input
    validate_input

    # Start main script operations
    if [ -n "$TENANT_ID" ]; then
        log_info "Starting onboarding process for Tenant ID: $TENANT_ID"
    else
        log_info "Starting onboarding process"
    fi

    if [ -n "$SUBSCRIPTION_ID" ]; then
        noninteractive_subscription
    elif [ -n "$MG_ID" ]; then
        noninteractive_management_group
    else
        check_az_cli_installed
        login_to_azure
        handle_account_type_selection
    fi

    confirm_selection

    case "$ACCOUNT_TYPE" in
        "Source")
            set_onboarding_scope
            handle_source_scope
            ;;
        "Restore")
            if [ "$RESTORE_SCOPE" == "subscription" ] || [ -z "$RESTORE_SCOPE" ]; then
                handle_restore_subscription
            else
                handle_restore_resource_group
            fi
            ;;
        *)
            log_error "❌ Invalid account type: $ACCOUNT_TYPE. Use 'Source' or 'Restore'."
            exit 1
            ;;
    esac
}

# Run the main script only if it is not being sourced
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi