#!/bin/bash

dry_run=false
verbose=false

info() {
  if [ "$dry_run" == "false" ]; then
    echo -e "\033[1;32m$1\033[0m"
  else
    echo -e "\033[1;32m[DRY RUN] $1\033[0m"
  fi
}

verbose() {
  if [ "$verbose" == "true" ]; then
    echo -e "\033[1;36m$1\033[0m"
  fi
}

helm() {
  echo -e "\033[1;34m$1\033[0m"
}

error() {
  if [ "$dry_run" == "false" ]; then
    echo -e "\033[1;31m$1\033[0m" >&2
  else
    echo -e "\033[1;31m[DRY RUN] $1\033[0m" >&2
  fi
}

# Common verbose run function used by both commands
verbose_run() {
  if [ "$verbose" == "true" ]; then
    echo -e "\033[1;36m[VERBOSE] Running: $*\033[0m" >&2
  fi
  
  local exit_code
  local stderr_output
  local temp_stderr=$(mktemp)
  
  "$@" 2>"$temp_stderr"
  exit_code=$?
  stderr_output=$(cat "$temp_stderr")
  rm -f "$temp_stderr"
  
  # Only exit if command actually failed (non-zero exit code)
  if [ $exit_code -ne 0 ]; then
    error "Command failed: $*"
    if [ -n "$stderr_output" ]; then
      error "$stderr_output"
    fi
    
    # Check for known errors
    # Multi-tenant app error - app is from another tenant
    if echo "$stderr_output" | grep -q "backing application.*must in the local tenant"; then
      info "The application is registered in a different Azure AD tenant."
      info "Please use the \"--create-application\" flag."
      info "Docs: https://docs.scaleops.com/cloud-billing-integration/azure/advanced-setup#azure-cost-integration-setup"
    fi

    exit $exit_code
  fi
  
  # Show stderr output for successful commands if verbose
  if [ "$verbose" == "true" ] && [ -n "$stderr_output" ]; then
    echo -e "\033[1;33m[STDERR] $stderr_output\033[0m" >&2
  fi
  
  return $exit_code
}


# Common capture run function
capture_run() {
  if $verbose; then
    echo -e "\033[1;36m[VERBOSE] Running: $*\033[0m" >&2
  fi
  
  local output
  local exit_code
  local stderr_output
  local temp_stderr=$(mktemp)
  
  output=$("$@" 2>"$temp_stderr")
  exit_code=$?
  stderr_output=$(cat "$temp_stderr")
  rm -f "$temp_stderr"
  
  # Only exit if command actually failed (non-zero exit code)
  if [ $exit_code -ne 0 ]; then
    error "Command failed: $*"
    if [ -n "$stderr_output" ]; then
      error "$stderr_output"
    fi
    exit $exit_code
  fi
  
  # Show stderr output for successful commands if verbose
  if [ "$verbose" == "true" ] && [ -n "$stderr_output" ]; then
    echo -e "\033[1;33m[STDERR] $stderr_output\033[0m" >&2
  fi
  
  echo "$output"
  return $exit_code
}

# Common silent run function
silent_run() {
  if $verbose; then
    echo -e "\033[1;36m[VERBOSE] Running: $*\033[0m" >&2
    "$@"
  else
    # Capture stderr to temp file, suppress both stdout and stderr
    local temp_stderr=$(mktemp)
    if "$@" >/dev/null 2>"$temp_stderr"; then
      # Success: remove temp file
      rm -f "$temp_stderr"
      return 0
    else
      # Failure: show captured stderr and return failure
      local exit_code=$?
      if [ -s "$temp_stderr" ]; then
        cat "$temp_stderr" >&2
      fi
      rm -f "$temp_stderr"
      return $exit_code
    fi
  fi
}

# Common Azure CLI check function
check_azure_cli() {
  # Skip Azure CLI checks in dry-run mode
  if [ "$dry_run" == "true" ]; then
    verbose "Skipping Azure CLI checks in dry-run mode"
    return 0
  fi

  if ! command -v az >/dev/null 2>&1; then
    error 'az not found'
    info "Please install az from https://github.com/Azure/azure-cli"
    exit 2
  fi

  verbose "Checking if the Azure CLI is logged in..."
  if [ "$verbose" == "true" ]; then
    if ! verbose_run az account show; then
      error "You are not logged in to Azure CLI."
      error "Please run 'az login' first and then retry."
      exit 2
    fi
  else
    if ! az account show >/dev/null 2>&1; then
      error "You are not logged in to Azure CLI."
      error "Please run 'az login' first and then retry."
      exit 2
    fi
  fi
  
  # Set the subscription context if provided
  if [ -n "$subscription" ]; then
    verbose "Setting Azure CLI context to subscription: $subscription"
    if [ "$verbose" == "true" ]; then
      if ! verbose_run az account set --subscription "$subscription"; then
        error "Failed to set subscription context to: $subscription"
        error "Please check that you have access to this subscription."
        exit 2
      fi
    else
      if ! az account set --subscription "$subscription" >/dev/null 2>&1; then
        error "Failed to set subscription context to: $subscription"
        error "Please check that you have access to this subscription."
        exit 2
      fi
    fi
    info "Azure CLI context set to subscription: $subscription"
  fi
}

setup_workload_identity() {
  if [ "$use_workload_identity" == "true" ]; then
    info "Setting up Azure Workload Identity"
    
    # Validate required parameters for workload identity
    if [ -z "$cluster_name" ]; then
      error "Workload identity requires --cluster-name parameter"
      info "Usage: --use-workload-identity --cluster-name <cluster>"
      exit 1
    fi
    
    # Discover cluster resource group automatically
    info "Discovering AKS cluster resource group..."
    cluster_resource_group=$(capture_run az aks list --query "[?name=='${cluster_name}'].resourceGroup" -o tsv --subscription "$subscription")
    
    if [ -z "$cluster_resource_group" ]; then
      error "AKS cluster '${cluster_name}' not found in subscription"
      info "Make sure:"
      info "  1. The cluster name is correct"
      info "  2. You have access to the cluster"
      info "  3. You're logged in with 'az login'"
      exit 1
    fi
    
    info "Found cluster '${cluster_name}' in resource group '${cluster_resource_group}'"
    
    # Get the AKS cluster OIDC issuer URL
    info "Getting AKS cluster OIDC issuer URL..."
    oidc_issuer_url=$(capture_run az aks show --resource-group "${cluster_resource_group}" --name "${cluster_name}" --query oidcIssuerProfile.issuerUrl -o tsv --subscription "$subscription")
    
    if [ -z "$oidc_issuer_url" ] || [ "$oidc_issuer_url" == "null" ]; then
      error "Failed to get OIDC issuer URL from AKS cluster '${cluster_name}' in resource group '${cluster_resource_group}'"
      info "Make sure:"
      info "  1. The cluster exists and you have access to it"
      info "  2. OIDC issuer is enabled on the cluster"
      info "  3. You're logged in with 'az login'"
      exit 1
    fi
    
    info "Found OIDC issuer URL: ${oidc_issuer_url}"
    
    # For Azure workload identity, we need to set up federated identity credentials
    # This connects the Kubernetes service accounts to the Azure AD application
    SERVICE_ACCOUNTS=("scaleops-dashboards" "scaleops-agent" "scaleops-recommender" "scaleops-updater")
    
    for SA in "${SERVICE_ACCOUNTS[@]}"; do
      credential_name="k8s-${installation_namespace}-${SA}"
      subject="system:serviceaccount:${installation_namespace}:${SA}"
      
      # Check if federated credential already exists
      existing_credential=$(capture_run az ad app federated-credential list --id "${application}" --query "[?name=='${credential_name}'].name" -o tsv  2>/dev/null || echo "")
      
      if [ -n "$existing_credential" ]; then
        info "Federated identity credential '${credential_name}' already exists - skipping"
        continue
      fi
      
      # Create the federated identity credential
      if ! silent_run az ad app federated-credential create \
        --id "${application}" \
        --parameters "{\"name\":\"${credential_name}\",\"issuer\":\"${oidc_issuer_url}\",\"subject\":\"${subject}\",\"description\":\"Federated identity for ScaleOps ${SA} service account\",\"audiences\":[\"api://AzureADTokenExchange\"]}"; then
        error "Failed to create federated identity credential for ${SA}"
        exit 1
      fi
      
      info "Successfully created federated identity credential for ${SA}"
    done
    
    info "Azure Workload Identity setup completed successfully"
  fi
}

# Cost integration specific functions
check_cost_export_exists() {
    verbose "Checking if cost export exists via REST API"
    
    az rest --method get \
      --subscription "$subscription" \
      --url "${export_scope}/providers/Microsoft.CostManagement/exports/${cost_export_name}?api-version=2025-03-01" \
      >/dev/null 2>&1
}

create_cost_export() {
    local from_date="$1"
    local payload=$(cat <<EOF
{
  "properties": {
    "schedule": {
      "status": "Active",
      "recurrence": "Daily",
      "recurrencePeriod": {
        "from": "${from_date}",
        "to":   "${end_date}"
      }
    },
    "format": "Csv",
    "deliveryInfo": {
      "destination": {
        "type": "AzureBlob",
        "resourceId": "${expected_storage_account_id}",
        "container":  "${container_name}",
        "rootFolderPath": "${export_directory_name}"
      }
    },
    "compressionMode": "GZip",
    "definition": {
      "type": "${export_type}",
      "timeframe": "MonthToDate",
      "dataSet": {
        "granularity": "Daily",
        "configuration": { 
          "dataVersion": "2021-10-01"
        }
      }
    }
  }
}
EOF
)
    
  if $verbose; then
    echo "$payload"
  fi
      
  verbose_run az rest --method put \
    --url "${export_scope}/providers/Microsoft.CostManagement/exports/${cost_export_name}?api-version=2025-03-01" \
    --body "$payload" \
    --subscription "$subscription"
}

# Check and create cost export role assignment
assign_cost_export_role() {
  local should_create_assignment=true
  

  sp_object_id=$(capture_run az ad sp show --id "$application" --query id -o tsv)
  if [ -z "$sp_object_id" ]; then
    error "Failed to get service principal object ID for application ${application}"
    exit 1
  fi
  
  # Check if role assignment already exists
  existing_assignment=$(capture_run timeout 10s az role assignment list --assignee "$sp_object_id" --role "${cost_role_definition_name}" --scope "/subscriptions/${subscription}" --query '[0].principalId' -o tsv || echo "")
  if [ -n "$existing_assignment" ] && [ "$existing_assignment" != "null" ]; then
    should_create_assignment=false
  fi

  if [ "$should_create_assignment" == "true" ]; then
      info "Assigning cost export role ${cost_role_definition_name} to service principal ${sp_object_id} (app: ${application}) (service_principal: ${sp_object_id})..."
      if ! verbose_run az role assignment create --assignee "$sp_object_id" --role "${cost_role_definition_name}" --scope "/subscriptions/${subscription}" --subscription "$subscription"; then
        error "Failed to assign cost export role ${cost_role_definition_name} to service principal ${sp_object_id}"
        exit 1
      fi
    info "Cost export role assignment created successfully"
  else
    info "Cost export role assignment already exists for service principal ${sp_object_id} - skipping creation"
  fi
}

  # Create cost export role definition
create_cost_export_role() {
    local actions=$(printf '"%s",' "${COST_EXPORT_PERMISSIONS[@]}")
    actions="[${actions%,}]"

    local cost_role_definition=$(cat <<EOF
{
  "Name":"${cost_role_definition_name}",
  "IsCustom":true,
  "Description":"Execute and manage cost exports for ScaleOps",
  "Actions":${actions},
  "NotActions":[],
  "AssignableScopes":["/subscriptions/${subscription}"]
}
EOF
)

    verbose "Checking cost export role definition ${cost_role_definition_name}..."
    existing_role_json=$(capture_run az role definition list --name "${cost_role_definition_name}" --output json --subscription "$subscription")
    if [ "$existing_role_json" == "[]" ]; then
      # Role doesn't exist, create it
      info "Creating cost export role definition: ${cost_role_definition_name} with permissions:"
      for permission in "${COST_EXPORT_PERMISSIONS[@]}"; do
        info "\t${permission}"
      done
      info "Creating cost export role definition ${cost_role_definition_name}..."
      if ! verbose_run az role definition create --role-definition "$cost_role_definition" --subscription "$subscription"; then
        error "Failed to create cost export role definition ${cost_role_definition_name}"
        exit 1
      fi
      sleep 10
      info "Cost export role definition ${cost_role_definition_name} created"
    else
      verbose "Validation role definition ${cost_role_definition_name} permissions..."
      for permission in "${COST_EXPORT_PERMISSIONS[@]}"; do
        verbose "\t${permission}"
        if [ "$(verbose_run az role definition list --name "${cost_role_definition_name}" --query "[?permissions[?contains(actions, '${permission}')]]" --output tsv --subscription "$subscription")" == "" ]; then
            verbose "\t\tPermission ${permission} is missing, attempting to add it..."
            
            # Get current actions as a comma-separated list
            current_actions=$(verbose_run az role definition list --name "${cost_role_definition_name}" --query "[0].permissions[0].actions" --output tsv --subscription "$subscription" | tr '\n' ',' | sed 's/,$//')
            if [ -z "$current_actions" ]; then
              error "\t\tFailed to retrieve current actions for role definition ${cost_role_definition_name}"
              exit 1
            fi
            
            # Check if permission already exists (double-check)
            if [[ "$current_actions" != *"$permission"* ]]; then
              # Build updated actions array - add the missing permission
              updated_actions="${current_actions},${permission}"
              
              # Build the updated role definition JSON
              role_id=$(verbose_run az role definition list --name "${cost_role_definition_name}" --query "[0].id" --output tsv --subscription "$subscription")
              
              # Create all actions as a properly formatted JSON array
              actions_json="["
              IFS=',' read -ra ACTION_ARRAY <<< "$updated_actions"
              for i in "${!ACTION_ARRAY[@]}"; do
                if [ $i -gt 0 ]; then
                  actions_json+=","
                fi
                actions_json+="\"${ACTION_ARRAY[$i]}\""
              done
              actions_json+="]"
              
              # Create the updated role definition
              updated_role_json="{
                \"roleName\":\"${cost_role_definition_name}\",
                \"Id\":\"${role_id}\",
                \"IsCustom\":true,
                \"Description\":\"Execute and manage cost exports for ScaleOps\",
                \"Actions\":${actions_json},
                \"NotActions\":[],
                \"AssignableScopes\":[\"/subscriptions/${subscription}\"]
              }"
              
              # Try to update the role definition
              if ! echo "$updated_role_json" | verbose_run az role definition update --role-definition @- --subscription "$subscription"; then
                error "\t\tRole definition ${cost_role_definition_name} does not have the required permissions ${permission}"
                error "\t\tFailed to add missing permission ${permission} to role definition"
                exit 1
              else
                verbose "\t\tSuccessfully added permission ${permission} to role definition ${cost_role_definition_name}"
                sleep 5  # Wait for Azure to propagate the changes
              fi
            fi
          fi
  done
    fi
}

create_cloud_integration_role() {
  actions=$(printf '"%s",' "${CLOUD_INTEGRATION_PERMISSIONS[@]}")
  actions="[${actions%,}]"

  ROLE_DEFINITION=$(cat <<EOF
{
  "Name":"${role_definition_name}",
  "IsCustom":true,
  "Description":"Read AKS clusters and manage their node pools",
  "Actions":${actions},
  "NotActions":[],
  "AssignableScopes":["/subscriptions/${subscription}"]
}
EOF
)
  if [ "$(verbose_run az role definition list --name "${role_definition_name}" --subscription "$subscription")" == "[]" ]; then
      info "Creating role definition: ${role_definition_name} with permissions:"
      for permission in "${CLOUD_INTEGRATION_PERMISSIONS[@]}"; do
        info "\t${permission}"
      done
      info "Creating role definition ${role_definition_name}..."
      if ! verbose_run az role definition create --role-definition "$ROLE_DEFINITION" --subscription "$subscription"; then
        error "Failed to create role definition ${role_definition_name}"
        exit 1
      fi
      sleep 10
      info "Role definition ${role_definition_name} created"
  else
    info "Role definition ${role_definition_name} already exists - skipping creation"
    verbose "Validation role definition ${role_definition_name} permissions..."
    for permission in "${CLOUD_INTEGRATION_PERMISSIONS[@]}"; do
      verbose "\t${permission}"
      if [ "$(verbose_run az role definition list --name "${role_definition_name}" --query "[?permissions[?contains(actions, '${permission}')]]" --output tsv --subscription "$subscription")" == "" ]; then
        verbose "\t\tPermission ${permission} is missing, attempting to add it..."

        # Get current actions as a comma-separated list
        current_actions=$(verbose_run az role definition list --name "${role_definition_name}" --query "[0].permissions[0].actions" --output tsv --subscription "$subscription" | tr '\n' ',' | sed 's/,$//')
        if [ -z "$current_actions" ]; then
          error "\t\tFailed to retrieve current actions for role definition ${role_definition_name}"
          exit 1
        fi

        # Check if permission already exists (double-check)
        if [[ "$current_actions" != *"$permission"* ]]; then
          # Build updated actions array - add the missing permission
          updated_actions="${current_actions},${permission}"

          # Build the updated role definition JSON
          role_id_val=$(verbose_run az role definition list --name "${role_definition_name}" --query "[0].id" --output tsv --subscription "$subscription")

          # Create all actions as a properly formatted JSON array
          actions_json="["
          IFS=',' read -ra ACTION_ARRAY <<< "$updated_actions"
          for i in "${!ACTION_ARRAY[@]}"; do
            if [ $i -gt 0 ]; then
              actions_json+=","
            fi
            actions_json+="\"${ACTION_ARRAY[$i]}\""
          done
          actions_json+="]"

          # Create the updated role definition
          updated_role_json="{
            \"roleName\":\"${role_definition_name}\",
            \"Id\":\"${role_id_val}\",
            \"IsCustom\":true,
            \"Description\":\"Read AKS clusters and manage their node pools\",
            \"Actions\":${actions_json},
            \"NotActions\":[],
            \"AssignableScopes\":[\"/subscriptions/${subscription}\"]
          }"

          # Try to update the role definition
          if ! echo "$updated_role_json" | verbose_run az role definition update --role-definition @- --subscription "$subscription"; then
            error "\t\tRole definition ${role_definition_name} does not have the required permissions ${permission}"
            error "\t\tFailed to add missing permission ${permission} to role definition"
            exit 1
          else
            verbose "\t\tSuccessfully added permission ${permission} to role definition ${role_definition_name}"
            sleep 5  # Wait for Azure to propagate the changes
          fi
        fi
      fi
      done
    fi
}



# Command to run
COMMAND=""

# Check if first argument is a command
if [ "$#" -gt 0 ]; then
  case $1 in
    cost|node-integration|all)
      COMMAND="$1"
      shift
      ;;
    *)
      # If first arg is not a command, default to help
      COMMAND="help"
      ;;
  esac
else
  COMMAND="help"
fi

# Cloud integration permissions (custom role)
CLOUD_INTEGRATION_PERMISSIONS=(
  "Microsoft.ContainerService/managedClusters/read"
  "Microsoft.ContainerService/managedClusters/write"
  "Microsoft.ContainerService/managedClusters/agentPools/read"
  "Microsoft.ContainerService/managedClusters/agentPools/write"
  "Microsoft.ContainerService/managedClusters/agentPools/delete"
  "Microsoft.ContainerService/managedClusters/agentPools/deleteMachines/action"
  "Microsoft.OperationsManagement/solutions/write"
  "Microsoft.OperationalInsights/workspaces/read"
  "Microsoft.OperationalInsights/workspaces/sharedkeys/read"
  "Microsoft.Network/virtualNetworks/subnets/join/action"
  "Microsoft.Quota/quotas/read"
  "Microsoft.Quota/usages/read"
  "Microsoft.ManagedIdentity/userAssignedIdentities/assign/action"
  "Microsoft.Network/ddosProtectionPlans/join/action"
  "Microsoft.Network/privateDnsZones/join/action"
  "Microsoft.Network/routeTables/join/action"
  "Microsoft.Network/publicIPPrefixes/join/action"
  "Microsoft.Compute/diskEncryptionSets/read"
  "Microsoft.Compute/capacityReservationGroups/deploy/action"
)

# Built-in Reader role name
READER_ROLE="Reader"

COST_EXPORT_PERMISSIONS=(
  "Microsoft.CostManagement/exports/read"
  "Microsoft.CostManagement/exports/action"
)

help=false
subscription=""
application=""
app_name="ScaleOps-Integration-$(date +%s)"
role_definition_name="ScaleOps Cloud Node Integration"
remove=false
create_application=false
skip_checks=false

# Cost billing integration defaults
location="westeurope"
cost_export_name="scaleops-cost-report"
export_type="AmortizedCost"
export_directory_name="scaleops-cost-exports"
resource_group_name="scaleops-cost-integration"
storage_account_name_prefix="scaleopscost"
storage_account_name=""
container_name="cost-exports"
cost_role_definition_name="ScaleOps Cost Export"
role_id="2a2b9908-6ea1-4ae2-8e65-a410df84e7d1"
start_date=$(date -u +%Y-%m-%dT00:00:00Z)
end_date="2099-01-01T00:00:00Z"
cost_app_name="ScaleOps-Cloud-Billing-Integration-$(date +%s)"
use_workload_identity=false
installation_namespace="scaleops-system"
cluster_name=""
minimal=false
additional_permissions=()

while [ "$#" -gt 0 ]; do
  case $1 in
  --subscription | -s | --subscription-id)
    subscription="$2"
    shift
    ;;
  --application | -a | --application-id)
    application="$2"
    shift
    ;;
  --app-name)
    app_name="$2"
    shift
    ;;
  --role-name)
    role_definition_name="$2"
    shift
    ;;
  --skip-checks)
    skip_checks=true
    ;;
  --location)
    location="$2"
    shift
    ;;
  --cost-export-name)
    cost_export_name="$2"
    shift
    ;;
  --export-type)
    export_type="$2"
    shift
    ;;
  --export-directory-name)
    export_directory_name="$2"
    shift
    ;;
  --resource-group-name)
    resource_group_name="$2"
    shift
    ;;
  --storage-account-name)
    storage_account_name="$2"
    shift
    ;;
  --container-name)
    container_name="$2"
    shift
    ;;
  --cost-role-name)
    cost_role_definition_name="$2"
    shift
    ;;
  --use-workload-identity)
    use_workload_identity=true
    ;;
  --installation-namespace)
    installation_namespace="$2"
    shift
    ;;
  --cluster-name)
    cluster_name="$2"
    shift
    ;;
  --minimal)
    minimal=true
    ;;
  --additional-permissions)
    IFS=',' read -ra additional_permissions <<< "$2"
    shift
    ;;
  --remove | -r)
    remove=true
    ;;
  --create-application)
    create_application=true
    ;;
  --dry-run)
    dry_run=true
    ;;
  --verbose | -v)
    verbose=true
    ;;
  --help | -h)
    help=true
    ;;
  -*)
    echo "error: unknown option $1"
    exit 1
    ;;
  esac
  shift
done

# Show help if requested or if command is help
if [ "$help" == "true" ] || [ "$COMMAND" == "help" ]; then
  echo "ScaleOps Azure Integration"
  echo
  echo "Usage:"
  echo "  $(basename "$0") <command> [flags]"
  echo
  echo "Commands:"
  echo "  cost               Set up Azure cost management integration"
  echo "  node-integration   Set up Azure node integration (AKS permissions)"
  echo "  all                Set up both cost and cloud node integration"
  echo
  echo "Flags:"
  echo "  -s <UUID>, --subscription <UUID>, --subscription-id <UUID>      Azure subscription ID (required for all commands)"
  echo "  -a <UUID>, --application <UUID>, --application-id <UUID>       Azure application ID (required for all commands unless --create-application is used)"
  echo "  --create-application                  Create a new Azure AD application instead of using existing"
  echo "  --app-name <name>                     Name for the Azure AD application (default: \"${app_name}\")"
  echo "  --role-name <name>                    Name for the role definition (default: \"${role_definition_name}\")"
  echo "  --minimal                             Use minimal permissions with custom role only (default: use built-in Reader role)"
  echo "  --additional-permissions <p1,p2,...>  Comma-separated list of extra Azure permissions to add to the custom role"
  echo "  --use-workload-identity               Enable workload identity setup (default: false)"
  echo "  --installation-namespace <NAMESPACE>  Installation namespace (default: scaleops-system)"
  echo "  --cluster-name <NAME>                 AKS cluster name (required for workload identity)"
  echo "  --skip-checks                         Skip authorization policies validations"
  echo "  -r, --remove                          Remove the role definition and role assignment (node-integration only)"
  echo "  --dry-run                             Show what would be created without making changes"
  echo "  -v, --verbose                         Print all az commands before executing them"
  echo "  -h, --help                            Show this help message"
  echo
  echo "Cost-specific flags:"
  echo "  --location <loc>                      Azure region (default: ${location})"
  echo "  --cost-export-name <name>             Cost export name (default: ${cost_export_name})"
  echo "  --export-type <type>               Cost export type - 'AmortizedCost' for EA, 'ActualCost' for WebDirect (default: ${export_type})"
  echo "  --export-directory-name <name>        Cost export directory name (default: ${export_directory_name})"
  echo "  --resource-group-name <name>          Resource group name (default: ${resource_group_name})"
  echo "  --storage-account-name <name>         Storage account name (default: ${storage_account_name_prefix}, default will be made unique)"
  echo "  --container-name <name>               Container name (default: ${container_name})"
  echo "  --cost-role-name <name>               Cost export role name (default: ${cost_role_definition_name})"
  echo
  echo "Examples:"
  echo "  $(basename "$0") cost --subscription abc123 --application def456"
  echo "  $(basename "$0") node-integration --subscription abc123 --create-application"
  echo "  $(basename "$0") all --subscription abc123 --application def456"
  exit 0
fi

# Validate flags for all commands that need Azure credentials
if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "node-integration" ] || [ "$COMMAND" == "all" ]; then
  if [ -z "${subscription}" ]; then
    error "Subscription ID is required for $COMMAND command"
    exit 1
  fi

  if [ "$create_application" == "false" ] && [ -z "${application}" ]; then
    error "Either --application <UUID> or --create-application must be provided for $COMMAND command"
    exit 1
  fi

  if [ "$create_application" == "true" ] && [ -n "${application}" ]; then
    error "Cannot specify both --application and --create-application"
    exit 1
  fi
fi

if [ "$export_type" == "Usage" ]; then
  error "Usage export type is not supported for cost integration"
  exit 1
fi

# Set role definition names with subscription suffix if unspecified
if [ -n "${subscription}" ]; then
  sub_prefix=$(echo "$subscription" | sed 's/-//g' | cut -c1-4)
  if [ "$role_definition_name" == "ScaleOps Cloud Node Integration" ]; then
    role_definition_name="${role_definition_name}-${sub_prefix}"
  fi
  if [ "$cost_role_definition_name" == "ScaleOps Cost Export" ]; then
    cost_role_definition_name="${cost_role_definition_name}-${sub_prefix}"
  fi
fi

# Create Azure AD application if requested
check_azure_cli

if [ -z "$storage_account_name" ] && [ -n "$subscription" ]; then
  # Take last 10 characters of subscription ID and remove hyphens to make name unique
  sub_suffix=$(echo "$subscription" | sed 's/-//g' | tail -c 11)
  storage_account_name="${storage_account_name_prefix}${sub_suffix}"
fi

if [ "${#additional_permissions[@]}" -gt 0 ]; then
  CLOUD_INTEGRATION_PERMISSIONS+=("${additional_permissions[@]}")
fi

########################################################
# Dry Run Mode
########################################################
if $dry_run; then
  if $remove; then
    if [ "$COMMAND" == "node-integration" ] || [ "$COMMAND" == "all" ]; then
      info "Would remove role definition: ${role_definition_name}"
      info "Would remove role assignment for the service principal: ${application} and role: ${role_definition_name}"
      if [ "$minimal" != "true" ]; then
        info "Would remove Reader role assignment for the service principal: ${application}"
      fi
    fi
    if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "all" ]; then
      info "Would remove cost export: ${cost_export_name}"
      info "Would remove role assignment for the service principal: ${application} and roles: ${cost_role_definition_name}"
      info "Would remove role definition: ${cost_role_definition_name}"
      info "Would remove container: ${container_name}"
      info "Would remove storage account: ${storage_account_name}"
      info "Would remove resource group: ${resource_group_name}"
    fi
    exit 0
  fi
  info "Dry run mode is enabled, no changes will be made"
  if [ "$skip_checks" != "true" ]; then
    info "Would validate authorization policies"
  fi
  if [ "$create_application" == "true" ]; then
      info "Would create Azure AD application: ${app_name}"
      info "Would create client secret for the application"
  fi
  if [ "$COMMAND" == "node-integration" ] || [ "$COMMAND" == "all" ]; then
    info "Would create role definition: ${role_definition_name} with permissions:"
    for permission in "${CLOUD_INTEGRATION_PERMISSIONS[@]}"; do
      info "\t${permission}"
    done
    info "Would create service principal for the application: ${application}"
    info "Would create role assignment for the service principal: ${application} and role: ${role_definition_name}"
    if [ "$minimal" != "true" ]; then
      info "Would assign built-in Reader role to service principal"
    fi
  fi
  if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "all" ]; then
    info "Would create cost export role definition: ${cost_role_definition_name} with permissions:"
    for permission in "${COST_EXPORT_PERMISSIONS[@]}"; do
      info "\t${permission}"
    done
    info "Would create role assignment for the service principal for application: ${application} and roles: ${cost_role_definition_name}, Blob Data Reader"
    info "Would create resource group: ${resource_group_name}"
    info "Would create storage account: ${storage_account_name}"
    info "Would create container: ${container_name}"
    info "Would create cost export: ${cost_export_name}"
    info "Would create role assignment for the service principal for application: ${application}"
  fi
  exit 0
fi

########################################################
# Remove Mode
########################################################
if $remove; then
  if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "all" ]; then
    info "Starting removal of resources for command: $COMMAND"
    export_scope="/subscriptions/$subscription"
    ctr_scope="/subscriptions/${subscription}/resourceGroups/${resource_group_name}/providers/Microsoft.Storage/storageAccounts/${storage_account_name}/blobServices/default/containers/${container_name}"
    
    sp_object_id=$(capture_run az ad sp show --id "$application" --query id -o tsv 2>/dev/null || echo "")
    
    # Remove cost export
    if check_cost_export_exists; then
      info "Removing cost export $cost_export_name"
      verbose_run az rest --method delete --url "${export_scope}/providers/Microsoft.CostManagement/exports/${cost_export_name}?api-version=2025-03-01"
    else
      info "Cost export $cost_export_name does not exist"
    fi
    
    # Remove role assignments (try both application and sp_object_id for backward compatibility)
    if [ -n "$sp_object_id" ]; then
      # Remove using service principal object ID (correct method)
      if silent_run az role assignment list --assignee "$sp_object_id" --role "$role_id" --scope "$ctr_scope" --query '[0]'; then
        info "Removing Storage Blob Data Reader role assignment for service principal $sp_object_id"
        verbose_run az role assignment delete --assignee "$sp_object_id" --role "$role_id" --scope "$ctr_scope"
      else
        info "Storage Blob Data Reader role assignment does not exist for service principal $sp_object_id"
      fi
      
      # Remove cost export role assignment
      if silent_run az role assignment list --assignee "$sp_object_id" --role "${cost_role_definition_name}" --scope "/subscriptions/${subscription}" --query '[0]'; then
        info "Removing cost export role assignment for service principal $sp_object_id"
        verbose_run az role assignment delete --assignee "$sp_object_id" --role "${cost_role_definition_name}" --scope "/subscriptions/${subscription}"
      else
        info "Cost export role assignment does not exist for service principal $sp_object_id"
      fi
    else
      # Fallback: try removing using application (for backward compatibility with old assignments)
      if silent_run az role assignment list --assignee "$application" --role "$role_id" --scope "$ctr_scope" --query '[0]'; then
        info "Removing Storage Blob Data Reader role assignment for app $application (fallback)"
        verbose_run az role assignment delete --assignee "$application" --role "$role_id" --scope "$ctr_scope"
      else
        info "Storage Blob Data Reader role assignment does not exist for app $application"
      fi
      
      if silent_run az role assignment list --assignee "$application" --role "${cost_role_definition_name}" --scope "/subscriptions/${subscription}" --query '[0]'; then
        info "Removing cost export role assignment for app $application (fallback)"
        verbose_run az role assignment delete --assignee "$application" --role "${cost_role_definition_name}" --scope "/subscriptions/${subscription}"
      else
        info "Cost export role assignment does not exist for app $application"
      fi
    fi
    
    # Remove cost export role definition
    if [ "$(capture_run az role definition list --name "${cost_role_definition_name}")" != "[]" ]; then
      info "Removing cost export role definition ${cost_role_definition_name}"
      verbose_run az role definition delete --name "${cost_role_definition_name}"
    else
      info "Cost export role definition ${cost_role_definition_name} does not exist"
    fi
    
    # Remove container
    verbose "Checking if container $container_name exists in storage account $storage_account_name"
    if timeout 10 bash -c "az storage container show -n '$container_name' --account-name '$storage_account_name' --subscription '$subscription' >/dev/null 2>&1"; then
      info "Removing container $container_name"
      verbose_run az storage container delete -n "$container_name" --account-name "$storage_account_name" --subscription "$subscription"
    else
      info "Container $container_name does not exist or storage account $storage_account_name does not exist"
    fi
      
    verbose "Checking if storage account $storage_account_name exists"
    if timeout 10 bash -c "az storage account show -n '$storage_account_name' -g '$resource_group_name' --subscription '$subscription' >/dev/null 2>&1"; then
      info "Removing storage account $storage_account_name"
      verbose_run az storage account delete -n "$storage_account_name" -g "$resource_group_name" --subscription "$subscription" --yes
    else
      info "Storage account $storage_account_name does not exist"
    fi
    
    # Remove resource group
    verbose "Checking if resource group $resource_group_name exists"
    if timeout 10 bash -c "az group show -n '$resource_group_name' --subscription '$subscription' >/dev/null 2>&1"; then
      info "Checking if resource group $resource_group_name is empty before removal"
      resource_count=$(az resource list -g "$resource_group_name" --subscription "$subscription" --query 'length(@)' -o tsv 2>/dev/null || echo "0")
      if [ "$resource_count" -eq 0 ]; then
        info "Removing empty resource group $resource_group_name"
        verbose_run az group delete -n "$resource_group_name" --subscription "$subscription" --yes
      else
        info "Resource group $resource_group_name is not empty ($resource_count resources), skipping removal"
      fi
    else
      info "Resource group $resource_group_name does not exist"
    fi
    
    # Remove app (if created by this script)
    if $create_application && [ -n "${application:-}" ]; then
      verbose "Checking if service principal $application exists"
      if timeout 10 bash -c "az ad sp show --id '$application' >/dev/null 2>&1"; then
        info "Removing service principal $application"
        verbose_run az ad sp delete --id "$application"
      else
        info "Service principal $application does not exist"
      fi
      
      verbose "Checking if application $application exists"
      if timeout 10 bash -c "az ad app show --id '$application' >/dev/null 2>&1"; then
        info "Removing application $application"
        verbose_run az ad app delete --id "$application"
      else
        info "Application $application does not exist"
      fi
    elif [ -n "${application:-}" ]; then
      info "Skipping service principal/application removal (not created by this script)"
    fi
    info "Removal complete – all resources deleted for command: $COMMAND"
  fi
  if [ "$COMMAND" == "node-integration" ] || [ "$COMMAND" == "all" ]; then
    info "Starting removal of resources for command: $COMMAND"

    # Remove Reader role assignment (if not minimal)
    if [ "$minimal" != "true" ]; then
      info "Removing Reader role assignment for the service principal: ${application}"
      verbose_run az role assignment delete --role "${READER_ROLE}" --scope "/subscriptions/${subscription}" --assignee "${application}" || true
    fi

    info "Removing role assignment for the service principal: ${application} and role: ${role_definition_name}"
    verbose_run az role assignment delete --role "${role_definition_name}" || true
    info "Removing role definition: ${role_definition_name}"
    verbose_run az role definition delete --name "${role_definition_name}" || true
    info "Removal complete – all resources deleted for command: $COMMAND"
  fi
  exit 0
fi


########################################################
# Validate Authorization Policies
########################################################
if [ "$skip_checks" != "true" ]; then
  info "Validating authorization policies"
  if [ "$create_application" == "true" ]; then
    can_create_applications=$(verbose_run az rest --method GET --url https://graph.microsoft.com/v1.0/policies/authorizationPolicy --query 'defaultUserRolePermissions.allowedToCreateApps' -o tsv)
    if [ "$can_create_applications" != "true" ]; then
      error "Current user does not have permission to create applications in Azure AD"
      exit 1
    fi
  fi
  info "Authorization policy checks completed successfully"
else
  info "Skipping authorization policy checks"
fi


########################################################
# Create Application
########################################################
if [ "$create_application" == "true" ]; then
  info "Creating Azure AD application: ${app_name}..."
  application=$(verbose_run az ad app create --display-name "${app_name}" --query "appId" --output tsv)
  if [ -z "$application" ]; then
    error "Failed to create Azure AD application"
    exit 1
  fi
  info "Azure AD application created with ID: ${application}"
  
  info "Creating client secret for application..."
  application_secret=$(verbose_run az ad app credential reset --id "${application}" --query "password" --output tsv --years 5)
  if [ -z "$application_secret" ]; then
    error "Failed to create client secret for application"
    exit 1
  fi
  info "Client secret created successfully"
fi

########################################################
# Cost Integration
########################################################
run_cost_integration() {
  info "Setting up Azure Cost Management integration..."

  # Create resource group
  if silent_run az group show -n "$resource_group_name" --subscription "$subscription"; then
    info "Resource group $resource_group_name already exists - skipping creation"
  else
    info "Creating resource group $resource_group_name"
    verbose_run az group create -n "$resource_group_name" -l "$location" --subscription "$subscription"
  fi

  # Register providers (only if not already registered)
  storage_state=$(capture_run az provider show --namespace Microsoft.Storage --subscription "$subscription" --query registrationState -o tsv)
  if [ "$storage_state" != "Registered" ]; then
    info "Registering Microsoft.Storage provider"
    verbose_run az provider register --namespace Microsoft.Storage --subscription "$subscription" --wait
  fi

  costmgmt_state=$(capture_run az provider show --namespace Microsoft.CostManagementExports --subscription "$subscription" --query registrationState -o tsv)
  if [ "$costmgmt_state" != "Registered" ]; then
    info "Registering Microsoft.CostManagementExports provider"
    verbose_run az provider register --namespace Microsoft.CostManagementExports --subscription "$subscription" --wait
  fi

  # Create storage account
  if silent_run az storage account show -n "$storage_account_name" -g "$resource_group_name" --subscription "$subscription"; then
    info "Storage account $storage_account_name already exists - skipping creation"
  else
    info "Creating storage account $storage_account_name"
    verbose_run az storage account create \
      -n "$storage_account_name" -g "$resource_group_name" -l "$location" --subscription "$subscription" \
      --sku Standard_LRS --kind StorageV2 --allow-blob-public-access false
  fi

  # Create container
  if silent_run az storage container show -n "$container_name" --account-name "$storage_account_name" --subscription "$subscription"; then
    info "Container $container_name already exists - skipping creation"
  else
    info "Creating container $container_name"
    verbose_run az storage container create \
      -n "$container_name" --account-name "$storage_account_name" --subscription "$subscription"
  fi

  # Check if cost export exists
  export_scope="/subscriptions/$subscription"
  expected_storage_account_id="/subscriptions/${subscription}/resourceGroups/${resource_group_name}/providers/Microsoft.Storage/storageAccounts/${storage_account_name}"
  ctr_scope="/subscriptions/${subscription}/resourceGroups/${resource_group_name}/providers/Microsoft.Storage/storageAccounts/${storage_account_name}/blobServices/default/containers/${container_name}"

  verbose "Checking if cost export $cost_export_name exists..."
  if check_cost_export_exists; then
    info "Cost export $cost_export_name exists - skipping creation"
  else
    info "Creating daily cost export with subscription scope"
    create_cost_export "$start_date"
  fi
  # create service principal if not exists
  service_principal=$(capture_run az ad sp show --id "$application" --query "id" --output tsv 2>/dev/null)

  if [ -z "$service_principal" ]; then
    info "Creating service principal for application: ${application}..."
    if ! verbose_run az ad sp create --id "$application"; then
      error "Failed to create service principal ${application}"
      exit 1
    fi
    sleep 10
    service_principal=$(capture_run az ad sp show --id "$application" --query "id" --output tsv 2>/dev/null)
  else
    info "Service principal ${application} already exists - skipping creation"
  fi

  sp_object_id=$(capture_run az ad sp show --id "$application" --query id -o tsv)
  if [ -z "$sp_object_id" ]; then
    error "Failed to get service principal object ID for application ${application}"
    exit 1
  fi

  # Check if role assignments exist
  storage_role_assignment_exists=false
  existing_storage_assignment=$(capture_run timeout 10s az role assignment list --assignee "$sp_object_id" --role "$role_id" --scope "$ctr_scope" --query '[0].principalId' -o tsv --subscription "$subscription" || echo "")
  if [ -n "$existing_storage_assignment" ] && [ "$existing_storage_assignment" != "null" ]; then
      storage_role_assignment_exists=true
  fi

  if $storage_role_assignment_exists; then
    info "Storage Blob Data Reader role assignment already exists for service principal ${sp_object_id} on container $container_name (app: ${application}) (service_principal: ${sp_object_id}) - skipping creation"
  else
    info "Granting Storage Blob Data Reader to service principal ${sp_object_id} (app: ${application}) (Container Scope)"
    verbose_run az role assignment create \
      --assignee "$sp_object_id" \
      --role "$role_id" \
      --scope "$ctr_scope" \
      --subscription "$subscription"
  fi

  # Create cost export role and assignment
  create_cost_export_role
  assign_cost_export_role

  # Set up workload identity if requested (only for individual commands, not "all")
  if [ "$COMMAND" != "all" ]; then
    setup_workload_identity
  fi

  info "Cost integration setup complete!"

  if [ "$COMMAND" != "all" ] && [ "$create_application" == "true" ]; then
    # Get tenant ID
    tenant_id=$(capture_run az account show --query tenantId -o tsv)

    echo
    info "Please add the following values to your helm chart:"
    echo
    helm "cloudBillingIntegration:"
    helm "  azure:"
    helm "    enabled: true"
    helm "    tenantId: \"${tenant_id}\""
    helm "    subscriptionId: \"${subscription}\""
    if [ "$use_workload_identity" != "true" ]; then
      helm "    applicationId: \"${application}\""
      helm "    applicationSecret: \"${application_secret}\""
    fi
    helm "    containerName: \"${container_name}\""
    helm "    storageAccountName: \"${storage_account_name}\""
    helm "    costExportName: \"${cost_export_name}\""
    helm "    exportDirectoryName: \"${export_directory_name}\""
    
    if [ "$use_workload_identity" == "true" ]; then
      helm "agent:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
      helm "dashboard:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
      helm "recommender:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
      helm "updater:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
    fi
    echo
  fi
}

########################################################
# Cloud Node Integration
########################################################
run_cloud_integration() {
  create_cloud_integration_role
  service_principal=$(capture_run az ad sp show --id "$application" --query "id" --output tsv 2>/dev/null)
  if [ -z "$service_principal" ]; then
    info "Creating service principal for application: ${application}..."
    if ! verbose_run az ad sp create --id "$application"; then
      error "Failed to create service principal ${application}"
      exit 1
    fi
    sleep 10
    info "Service principal for application: ${application} created"
    service_principal=$(capture_run az ad sp show --id "$application" --query "id" --output tsv 2>/dev/null)
  else
    info "Service principal ${application} already exists - skipping creation"
  fi

  if [ -z "$service_principal" ]; then
    error "Failed to get service principal, do you have the right permissions?"
    exit 1
  fi

  # Assign built-in Reader role (unless --minimal is specified)
  if [ "$minimal" != "true" ]; then
    # Check if Reader role assignment already exists
    if ! existing_reader_check=$(capture_run az role assignment list --subscription "$subscription" --role "${READER_ROLE}" --assignee "$service_principal" --scope "/subscriptions/${subscription}"); then
      error "Failed to check existing Reader role assignments"
      exit 1
    fi

    if [ "$existing_reader_check" == "[]" ]; then
      info "Assigning built-in Reader role to service principal ${service_principal}..."
      for _ in $(seq 1 10); do
        if verbose_run az role assignment create --assignee-object-id "$service_principal" --assignee-principal-type "ServicePrincipal" --role "${READER_ROLE}" --scope "/subscriptions/${subscription}"; then
            break
        fi
        sleep 5
      done
    else
      info "Reader role assignment already exists for service principal ${service_principal} - skipping"
    fi
  fi

  # Check if we need to create custom role assignment
  should_create_assignment=true

  # Only check existing assignments if we have a valid service principal
  if ! existing_assignment_check=$(capture_run az role assignment list --subscription "$subscription" --role "${role_definition_name}" --assignee "$service_principal" --scope "/subscriptions/${subscription}"); then
    error "Failed to check existing role assignments"
    exit 1
  fi

  if [ "$existing_assignment_check" != "[]" ]; then
    should_create_assignment=false
    info "Role assignment already exists for service principal ${service_principal} - skipping creation"
  fi

  if [ "$should_create_assignment" == "true" ]; then
    info "Assigning role definition ${role_definition_name} to application ${application}..."
    for _ in $(seq 1 10); do
      if verbose_run az role assignment create --assignee-object-id "$service_principal" --assignee-principal-type "ServicePrincipal" --role "$role_definition_name" --scope "/subscriptions/${subscription}"; then
          break
      fi
      sleep 5
    done
  fi

  # Set up workload identity if requested (only for individual commands, not "all")
  if [ "$COMMAND" != "all" ]; then
    setup_workload_identity
  fi

  info "Cloud node integration setup complete!"

  # Print helm values if we created an application
  if [ "$COMMAND" != "all" ] && [ "$create_application" == "true" ]; then
    # Get tenant ID
    tenant_id=$(capture_run az account show --query tenantId -o tsv)

    echo
    info "Please add the following values to your helm chart:"
    echo
    helm "cloudNodeIntegration:"
    helm "  azure:"
    helm "    enabled: true"
    helm "    tenantId: \"${tenant_id}\""
    helm "    subscriptionId: \"${subscription}\""
    if [ "$use_workload_identity" != "true" ]; then
      helm "    applicationId: \"${application}\""
      helm "    applicationSecret: \"${application_secret}\""
    fi
    
    if [ "$use_workload_identity" == "true" ]; then
      helm "agent:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
      helm "dashboard:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
      helm "recommender:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
      helm "updater:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      azure.workload.identity/client-id: \"${application}\""
    fi
    echo
  fi
}

show_combined_helm_values() {
  # Get tenant ID
  tenant_id=$(capture_run az account show --query tenantId -o tsv)
  
  echo
  info "Please add the following combined values to your helm chart:"
  echo
  
  # Cloud billing integration
  helm "cloudBillingIntegration:"
  helm "  azure:"
  helm "    enabled: true"
  helm "    tenantId: \"${tenant_id}\""
  helm "    subscriptionId: \"${subscription}\""
  if [ "$use_workload_identity" != "true" ]; then
    helm "    applicationId: \"${application}\""
    helm "    applicationSecret: \"${application_secret}\""
  fi
  helm "    containerName: \"${container_name}\""
  helm "    storageAccountName: \"${storage_account_name}\""
  helm "    costExportName: \"${cost_export_name}\""
  helm "    exportDirectoryName: \"${export_directory_name}\""
  
  # Cloud node integration
  helm "cloudNodeIntegration:"
  helm "  azure:"
  helm "    enabled: true"
  helm "    tenantId: \"${tenant_id}\""
  helm "    subscriptionId: \"${subscription}\""
  if [ "$use_workload_identity" != "true" ]; then
    helm "    applicationId: \"${application}\""
    helm "    applicationSecret: \"${application_secret}\""
  fi
  
  # Workload identity annotations (if enabled)
  if [ "$use_workload_identity" == "true" ]; then
    helm "agent:"
    helm "  serviceAccount:"
    helm "    annotations:"
    helm "      azure.workload.identity/client-id: \"${application}\""
    helm "dashboard:"
    helm "  serviceAccount:"
    helm "    annotations:"
    helm "      azure.workload.identity/client-id: \"${application}\""
    helm "recommender:"
    helm "  serviceAccount:"
    helm "    annotations:"
    helm "      azure.workload.identity/client-id: \"${application}\""
    helm "updater:"
    helm "  serviceAccount:"
    helm "    annotations:"
    helm "      azure.workload.identity/client-id: \"${application}\""
  fi
  
  echo
}

# Execute the appropriate command
case "$COMMAND" in
  cost)
    run_cost_integration
    ;;
  node-integration)
    run_cloud_integration
    ;;
  all)
    info "Running both cost and cloud node integration..."
    run_cloud_integration
    echo
    run_cost_integration
    setup_workload_identity
    
    # Show combined helm values
    if [ "$create_application" == "true" ]; then
      show_combined_helm_values
    fi
    ;;
  *)
    error "Unknown command: $COMMAND"
    exit 1
    ;;
esac