#!/bin/bash

# Minimal permissions - used with --minimal flag (custom role)
MINIMAL_PERMISSIONS=(
  "container.clusters.get"
  "container.clusters.list"
  "container.clusters.update"
  "container.operations.get"
  "compute.regions.get"
  "compute.instances.delete"
  "compute.instanceGroupManagers.update"
  "compute.instanceGroupManagers.list"
  "compute.instanceTemplates.delete"
  "compute.instanceTemplates.get"
  "compute.instanceTemplates.list"
)

# Full viewer roles - used by default (predefined roles)
VIEWER_ROLES=(
  "roles/bigquery.dataViewer"
  "roles/compute.viewer"
  "roles/compute.networkViewer"
  "roles/container.viewer"
  "roles/aiplatform.viewer"
  "roles/monitoring.viewer"
  "roles/cloudsql.viewer"
  "roles/alloydb.viewer"
  "roles/datastore.viewer"
  "roles/bigtable.viewer"
  "roles/spanner.viewer"
  "roles/run.viewer"
  "roles/redis.viewer"
  "roles/dataflow.viewer"
  "roles/cloudfunctions.viewer"
  "roles/storage.objectViewer"
  "roles/cloudasset.viewer"
  "roles/pubsub.viewer"
  "roles/file.viewer"
  "roles/netapp.viewer"
)

# Write permissions - always needed for node operations (custom role)
WRITE_PERMISSIONS=(
  "compute.instances.delete"
  "compute.instanceGroupManagers.update"
  "compute.instanceTemplates.delete"
  "container.clusters.update"
)


COMMAND=""

if [ "$#" -gt 0 ]; then
  case $1 in
    cost|node-integration|all)
      COMMAND="$1"
      shift
      ;;
    *)
      COMMAND="help"
      ;;
  esac
else
  COMMAND="help"
fi

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
}

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

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
}

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
}

check_gcloud_cli() {
  if ! command -v gcloud >/dev/null 2>&1; then
    error 'gcloud not found'
    info "Please install gcloud from https://cloud.google.com/sdk/docs/install"
    exit 2
  fi
  if ! capture_run gcloud auth list --filter=status:ACTIVE --format="value(account)"; then
    error "You are not logged in to Google Cloud CLI."
    error "Please run 'gcloud auth login' first and then retry."
    exit 2
  fi
}

check_bq_cli() {
  if ! command -v bq >/dev/null 2>&1; then
    info "BigQuery CLI (bq) not found, installing..."
      if ! silent_run gcloud components install bq; then
    error "Failed to install BigQuery CLI component"
    exit 2
  fi
  fi
}

check_jq_cli() {
  if ! command -v jq >/dev/null 2>&1; then
    error "jq command not found. Please install jq: https://stedolan.github.io/jq/download/"
    exit 2
  fi
}

check_organization_policies() {
  local project="$1"
  
  local sa_creation_enforced=false
  local domain_restriction_enforced=false
  
  # Check for service account creation/key creation policies
  for constraint in iam.managed.disableServiceAccountCreation iam.disableServiceAccountCreation iam.managed.disableServiceAccountKeyCreation iam.disableServiceAccountKeyCreation; do
    verbose "Checking organization policy: constraints/$constraint"
    local sa_key_policy=$(capture_run gcloud resource-manager org-policies list \
        --filter="constraint:constraints/$constraint" \
        --project="${project}" \
        --format="value(constraint)" 2>/dev/null || echo "")

    if [ -n "$sa_key_policy" ]; then
      local policy_enforced=$(capture_run gcloud resource-manager org-policies describe \
        "constraints/$constraint" \
        --project="${project}" \
        --effective \
        --format="value(booleanPolicy.enforced)" 2>/dev/null || echo "false")

      if [ "$policy_enforced" == "True" ]; then
        sa_creation_enforced=true
        break
      fi
    fi
  done

  # Check for domain restriction policies
  for constraint in iam.managed.allowedPolicyMemberDomains iam.allowedPolicyMemberDomains; do
    verbose "Checking organization policy: constraints/$constraint"
    local external_sa_policy=$(capture_run gcloud resource-manager org-policies list \
      --project="${project}" \
      --filter="constraint:constraints/$constraint" \
      --format="value(constraint)" 2>/dev/null || echo "")

    if [ -n "$external_sa_policy" ]; then
      local policy_json=$(capture_run gcloud resource-manager org-policies describe \
        "constraints/$constraint" \
        --project="${project}" \
        --effective \
        --format="json" 2>/dev/null || echo "{}")
      local all_values_rule=$(echo "$policy_json" | jq -r '.listPolicy.allValues // empty' 2>/dev/null || echo "")

      if [ "$all_values_rule" != "ALLOW" ]; then
        domain_restriction_enforced=true
        break
      fi
    fi
  done

  # Apply logic based on which policies are enforced
  if [ "$sa_creation_enforced" == "true" ]; then
    # If SA creation is blocked, always use workload identity (regardless of domain restrictions)
    error "Organization policies prevent service account creation/key generation."
    error ""
    error "Please rerun the setup using Workload Identity instead by adding the --use-workload-identity flag."
    error ""
    if [ "$COMMAND" == "cost" ]; then
      error "For more information, see: https://docs.scaleops.com/cloud-billing-integration/google/workload-identity"
    else
      error "For more information, see: https://docs.scaleops.com/integrations/cloud-node/google/workload-identity"
    fi
    exit 1
  elif [ "$domain_restriction_enforced" == "true" ]; then
    # If only domain restrictions (but SA creation allowed), use created service account
    error "Organization policies restrict external service account access."
    error ""
    error "Please rerun the setup using a created service account instead by adding the --create-service-account flag."
    error ""
    if [ "$COMMAND" == "cost" ]; then
      error "For more information, see: https://docs.scaleops.com/cloud-billing-integration/google/advanced-setup#google-cloud-cost-integration-setup"
    else
      error "For more information, see: https://docs.scaleops.com/integrations/cloud-node/google/advanced-setup#google-integration-setup"
    fi
    exit 1
  fi
}

# Legacy functions for backward compatibility - now just call the unified function
check_organization_create_policies() {
  check_organization_policies "$1"
}

check_organization_domain_policies() {
  check_organization_policies "$1"
}

setup_workload_identity() {
  if [ "$use_workload_identity" == "true" ]; then
    info "Setting up Workload Identity"
    
    # Determine which projects we need to set up workload identity for
    projects_to_setup=()
    
    if [ -n "${project_id}" ]; then
      projects_to_setup+=("${project_id}")
    fi
    
    if [ -n "${cost_project_id}" ] && [ "${cost_project_id}" != "${project_id}" ]; then
      projects_to_setup+=("${cost_project_id}")
    fi
    
    # If no projects specified, this shouldn't happen but handle gracefully
    if [ ${#projects_to_setup[@]} -eq 0 ]; then
      error "No project ID specified for workload identity setup"
      exit 1
    fi
    
    SERVICE_ACCOUNTS=("scaleops-dashboards" "scaleops-agent" "scaleops-recommender" "scaleops-updater")
    
    for workload_project_id in "${projects_to_setup[@]}"; do
      info "Setting up workload identity for project: ${workload_project_id}"
      
      for SA in "${SERVICE_ACCOUNTS[@]}"; do
        # Extract the service account project from the service account email
        service_account_project=$(echo "${service_account}" | cut -d'@' -f2 | cut -d'.' -f1)
        
        # Check if workload identity binding already exists
        workload_identity_binding_exists=$(capture_run gcloud iam service-accounts get-iam-policy "${service_account}" \
          --flatten="bindings[].members" \
          --format="table(bindings.role)" \
          --filter="bindings.role=roles/iam.workloadIdentityUser AND bindings.members=serviceAccount:${workload_project_id}.svc.id.goog[${installation_namespace}/${SA}]" \
          --project="${service_account_project}")
        
        if [ "$workload_identity_binding_exists" == "" ]; then
          info "Setting up workload identity for ${SA} in namespace ${installation_namespace} (project: ${workload_project_id})"
          if ! silent_run gcloud iam service-accounts add-iam-policy-binding "${service_account}" \
            --role roles/iam.workloadIdentityUser \
            --member "serviceAccount:${workload_project_id}.svc.id.goog[${installation_namespace}/${SA}]" \
            --project="${service_account_project}" --condition=None; then
            error "Failed to set up workload identity for ${SA} in project ${workload_project_id}"
            exit 1
          fi
        else
          info "Workload identity binding already exists for ${SA} in project ${workload_project_id} - skipping"
        fi
      done
    done
  fi
}

impersonate_cluster_service_account() {
  project_id="$1"
  service_account="$2"
  cluster="$3"
  location="$4"

  info "Granting Service Account User permissions to all node pools for: ${cluster} (location: ${location})"
  for sa in $(capture_run gcloud container node-pools list --cluster "${cluster}" --project="${project_id}" --location="${location}" --format="value(config.serviceAccount)" | sort | uniq); do
    if [ "$(echo "${sa}" | xargs)" == "default" ]; then
      sa="$(capture_run gcloud projects describe "${project_id}" --format="value(projectNumber)")-compute@developer.gserviceaccount.com"
    fi
    impersonate_service_account "${project_id}" "${service_account}" "${sa}"
  done
}

impersonate_service_account() {
  project_id="$1"
  service_account="$2"
  sa="$3"

  sa_trimmed=$(echo "${sa}" | xargs)  # Remove whitespace

  # Check if binding already exists
  existing_binding=$(capture_run gcloud iam service-accounts get-iam-policy "${sa_trimmed}" \
    --flatten="bindings[].members" \
    --format="table(bindings.role)" \
    --filter="bindings.role=roles/iam.serviceAccountUser AND bindings.members=serviceAccount:${service_account}")

  if [ "$existing_binding" == "" ]; then
    info "Granting Service Account User permissions to ScaleOps SA for: ${sa_trimmed}"
    if ! silent_run gcloud iam service-accounts add-iam-policy-binding "${sa_trimmed}" \
      --member="serviceAccount:${service_account}" \
      --role="roles/iam.serviceAccountUser" --project="${project_id}" --condition=None; then
      error "Failed to grant Service Account User permissions for: ${sa_trimmed}"
      exit 1
    fi
  else
    info "Service Account User permissions already exist for: ${sa_trimmed} - skipping"
  fi
}


get_role_path() {
  echo "projects/${project_id}/roles/${role_id}"
}

check_dataset_exists() {
  local dataset="$1"
  capture_run bq show --dataset "${project_id}:${dataset}" >/dev/null 2>&1
}

# Check if role permissions match required permissions
permissions_match() {
  # Get current permissions
  local current_permissions
  current_permissions=$(verbose_run gcloud iam roles describe "${role_id}" --format="value(includedPermissions[].join(','))" --project="${project_id}" 2>/dev/null | tr ',' '\n' | sort)

  if [ "$current_permissions" = "$(printf '%s\n' "${PERMISSIONS[@]}" | sort)" ]; then
    return 0
  fi
  return 1
}

help=false
if [ "$#" -eq 0 ]; then
  help=true
fi

project_id=""
cost_project_id=""
service_account=""
role_id="ScaleOpsIntegration"
role_display_name="ScaleOps Cloud Node Integration"
service_account_name="scaleops-$(date +%s)"
allow_service_accounts=""
cluster=""
all_clusters=false
remove=false
create_service_account=false
bq_dataset_name=""
use_workload_identity=false
installation_namespace="scaleops-system"
skip_checks=false
minimal=false


while [ "$#" -gt 0 ]; do
  case $1 in
  --project-id | -p)
    project_id="$2"
    shift
    ;;
  --cost-project-id)
    cost_project_id="$2"
    shift
    ;;
  --service-account | -s)
    service_account="$2"
    shift
    ;;
  --role-id | -i)
    role_id="$2"
    shift
    ;;
  --role-display-name)
    role_display_name="$2"
    shift
    ;;
  --service-account-name)
    service_account_name="$2"
    shift
    ;;
  --allow-service-accounts)
    allow_service_accounts="$2"
    shift
    ;;
  --cluster)
    cluster="$2"
    shift
    ;;
  --all-clusters)
    all_clusters=true
    ;;
  --remove | -r)
    remove=true
    ;;
  --create-service-account)
    create_service_account=true
    ;;
  --dataset-name | -d)
    bq_dataset_name="$2"
    shift
    ;;
  --use-workload-identity)
    use_workload_identity=true
    ;;
  --skip-checks)
    skip_checks=true
    ;;
  --minimal)
    minimal=true
    ;;
  --installation-namespace)
    installation_namespace="$2"
    shift
    ;;
  --dry-run)
    dry_run=true
    ;;
  --verbose | -v)
    verbose=true
    ;;
  --help | -h)
    help=true
    ;;
  -*)
    echo "error: unknown option $1"
    exit 1
    ;;
  esac
  shift
done


if [ "$help" == "true" ] || [ "$COMMAND" == "help" ]; then
  echo "ScaleOps GCP Integration"
  echo
  echo "Usage:"
  echo "  $(basename "$0") <command> [flags]"
  echo
  echo "Commands:"
  echo "  cost               Set up GCP cost integration with BigQuery"
  echo "  node-integration   Set up GCP node integration permissions"
  echo "  all                Set up both cost and node integration"
  echo
  echo "Flags:"
  echo "  -p <PROJECT_ID>, --project-id <PROJECT_ID>    GCP project ID (required)"
  echo "  -s <EMAIL>, --service-account <EMAIL>         Service account email (required unless --create-service-account is used)"
  echo "  -i <NAME>, --role-id <NAME>                   ID for the role (default: \"${role_id}\")"
  echo "  --service-account-name <NAME>                 Name for the service account (default: \"${service_account_name}\")"
  echo "  --role-display-name <NAME>                    Display name for the role (default: \"${role_display_name}\")"
  echo "  --create-service-account                      Create a new service account instead of using existing"
  echo "  --skip-checks                                 Skip organization policies validations"
  echo
  echo "Cost integration specific flags:"
  echo "  --cost-project-id <PROJECT_ID>                Project ID for BigQuery cost data (required for cost integration)"
  echo "  -d <DATASET>, --dataset-name <DATASET>        BigQuery dataset name (required for cost integration)"
  echo
  echo "Node integration specific flags:"
  echo "  --minimal                                     Use minimal permissions with custom role only (default: use predefined viewer roles)"
  echo "  --allow-service-accounts <SA1,SA2>            Comma-separated list of service accounts to grant impersonation"
  echo "                                                 permissions for (required for node operations with custom SAs)"
  echo "  --cluster <CLUSTER_NAME>                      Grant impersonation permissions for all node groups in this cluster"
  echo "  --all-clusters                                Grant impersonation permissions for all node groups in all clusters in the project"
  echo "  --use-workload-identity                       Enable workload identity setup (default: false)"
  echo "  --installation-namespace <NAMESPACE>          Installation namespace (default: scaleops-system)"
  echo
  echo "  -r, --remove                                  Remove the custom role and IAM policy binding"
  echo "  --dry-run                                     Show what would be created without making changes"
  echo "  -v, --verbose                                 Print all gcloud commands before executing them"
  echo "  -h, --help                                    Show this help message"
  echo
  echo "Examples:"
  echo "  $(basename "$0") cost --project-id my-project-123 --cost-project-id my-billing-project --create-service-account --dataset-name my-billing-dataset"
  echo "  $(basename "$0") cost --project-id my-project-123 --cost-project-id my-billing-project --service-account sa@project.iam.gserviceaccount.com --dataset-name my-billing-dataset"
  echo "  $(basename "$0") node-integration --project-id my-project-123 --create-service-account"
  echo "  $(basename "$0") node-integration --project-id my-project-123 --service-account sa@project.iam.gserviceaccount.com \\"
  echo "    --allow-service-accounts node-sa@project.iam.gserviceaccount.com,gpu-sa@project.iam.gserviceaccount.com"
  echo "  $(basename "$0") node-integration --project-id my-project-123 --service-account sa@project.iam.gserviceaccount.com --use-workload-identity"
  echo "  $(basename "$0") all --project-id my-project-123 --cost-project-id my-billing-project --create-service-account --dataset-name my-billing-dataset"
  exit 0
fi
if [ "$COMMAND" == "all" ] || [ "$COMMAND" == "node-integration" ]; then
  if [ -z "${project_id}"  ]; then
    error "--project-id is required"
    exit 1
  fi
fi

if [ "$create_service_account" == "false" ] && [ -z "${service_account}" ]; then
  error "Either --service-account or --create-service-account is required"
  exit 1
fi

if [ "$create_service_account" == "true" ] && [ -n "${service_account}" ]; then
  error "Cannot specify both --service-account and --create-service-account"
  exit 1
fi

SCALEOPS_MANAGED_PROJECT="${SCALEOPS_MANAGED_PROJECT:-scaleops-platform}"

if [ "$use_workload_identity" == "true" ] && [ -n "${service_account}" ]; then
  sa_project=$(echo "${service_account}" | cut -d'@' -f2 | cut -d'.' -f1)
  if [ "$sa_project" == "$SCALEOPS_MANAGED_PROJECT" ]; then
    error "Cannot use --use-workload-identity with the ScaleOps managed service account."
    error "The service account '${service_account}' is managed by ScaleOps and cannot be bound to your Workload Identity pool."
    error ""
    error "Options:"
    error "  1. Remove --use-workload-identity to use the managed service account with key-based credentials."
    error "  2. Replace --service-account with --create-service-account to create a new SA in your project:"
    error "     --create-service-account --use-workload-identity"
    exit 1
  fi
fi

if [ "$create_service_account" == "true" ] && [ "$COMMAND" == "cost" ] && [ -z "${project_id}" ] && [ -z "${cost_project_id}" ]; then
  error "Either --project-id or --cost-project-id is required when using --create-service-account with cost command"
  exit 1
fi

if [ -n "${allow_service_accounts}" ] && [ "$COMMAND" != "node-integration" ] && [ "$COMMAND" != "all" ]; then
  error "--allow-service-accounts can only be used with node-integration or all command"
  exit 1
fi
if [ -n "${cluster}" ] && [ "$COMMAND" != "node-integration" ] && [ "$COMMAND" != "all" ]; then
  error "--cluster can only be used with node-integration or all command"
  exit 1
fi
if [ "$all_clusters" == "true" ] && [ "$COMMAND" != "node-integration" ] && [ "$COMMAND" != "all" ]; then
  error "--all-clusters can only be used with node-integration or all command"
  exit 1
fi

if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "all" ]; then
  if [ -z "${bq_dataset_name}" ]; then
    error "--dataset-name is required for cost integration"
    exit 1
  fi
  if [ -z "${cost_project_id}" ]; then
    error "--cost-project-id is required for cost integration"
    exit 1
  fi
fi


check_gcloud_cli


########################################################
# Dry Run Mode
########################################################
if $dry_run; then
  if $remove; then
    if [ "$COMMAND" == "node-integration" ] || [ "$COMMAND" == "all" ]; then
      info "Would remove custom role: ${role_id}"
      info "Would remove IAM policy binding for service account: ${service_account} and role: ${role_id}"

      if [ "$minimal" != "true" ]; then
        info "Would remove predefined viewer role bindings:"
        for role in "${VIEWER_ROLES[@]}"; do
          info "\t${role}"
        done
      fi

      if [ -n "${allow_service_accounts}" ]; then
        IFS=',' read -ra SA_ARRAY <<< "${allow_service_accounts}"
        for sa in "${SA_ARRAY[@]}"; do
          info "Would remove Service Account User permissions from ScaleOps SA for: ${sa}"
        done
      fi
    fi
    if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "all" ]; then
      info "Would remove BigQuery Data Viewer permissions from service account on dataset: ${bq_dataset_name} in project: ${cost_project_id}"
      info "Would remove BigQuery Job User role from service account"
    fi
    exit 0
  fi
  info "Dry run mode is enabled, no changes will be made"
  if [ "$skip_checks" != "true" ]; then
    info "Would validate organization policies"
  fi
  
  if [ "$COMMAND" == "node-integration" ] || [ "$COMMAND" == "all" ]; then
    if [ "$minimal" == "true" ]; then
      info "Mode: Minimal (custom role only)"
      info "Would create custom role ${role_id} with permissions:"
      for permission in "${MINIMAL_PERMISSIONS[@]}"; do
        info "\t${permission}"
      done
      info "Would create IAM policy binding for service account: ${service_account} and custom role: ${role_id}"
    else
      info "Mode: Full (predefined viewer roles + custom write role)"
      info "Would grant the following predefined viewer roles:"
      for role in "${VIEWER_ROLES[@]}"; do
        info "\t${role}"
      done
      info "Would create custom role ${role_id} for write permissions:"
      for permission in "${WRITE_PERMISSIONS[@]}"; do
        info "\t${permission}"
      done
      info "Would create IAM policy binding for service account: ${service_account}"
    fi
    if [ -n "${allow_service_accounts}" ]; then
      IFS=',' read -ra SA_ARRAY <<< "${allow_service_accounts}"
      for sa in "${SA_ARRAY[@]}"; do
        info "Would grant Service Account User permissions to ScaleOps SA for: ${sa}"
      done
    fi
    if [ -n "${cluster}" ]; then
      info "Would grant Service Account User permissions to ScaleOps SA for cluster ${cluster}'s node pools"
    fi
    if [ "$all_clusters" == "true" ]; then
      info "Would grant Service Account User permissions to ScaleOps SA for all node pools in all clusters"
    fi
  fi
  if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "all" ]; then
    info "Would grant BigQuery Data Viewer permissions to service account on dataset: ${bq_dataset_name} in project: ${cost_project_id}"
    info "Would grant BigQuery Job User role to service account for query execution"
  fi
  exit 0
fi

########################################################
# Remove Mode
########################################################
if $remove; then
  if [ "$COMMAND" == "node-integration" ] || [ "$COMMAND" == "all" ]; then
    if ! verbose_run gcloud config set project "${project_id}" --quiet; then
      error "Failed to set gcloud project"
      exit 1
    fi
    info "Starting removal of resources for command: $COMMAND"
    role_path=$(get_role_path)

    # Remove custom role binding
    info "Removing IAM policy binding for service account: ${service_account} and role: ${role_id}"
    capture_run gcloud projects remove-iam-policy-binding "${project_id}" \
      --member="serviceAccount:${service_account}" \
      --role="${role_path}" || true

    # Remove predefined viewer role bindings (if they exist from non-minimal mode)
    if [ "$minimal" != "true" ]; then
      info "Removing predefined viewer role bindings..."
      for viewer_role in "${VIEWER_ROLES[@]}"; do
        info "Removing role: ${viewer_role}"
        capture_run gcloud projects remove-iam-policy-binding "${project_id}" \
          --member="serviceAccount:${service_account}" \
          --role="${viewer_role}" || true
      done
    fi

    if [ -n "${allow_service_accounts}" ]; then
      IFS=',' read -ra SA_ARRAY <<< "${allow_service_accounts}"
      for sa in "${SA_ARRAY[@]}"; do
        sa_trimmed=$(echo "${sa}" | xargs)  # Remove whitespace
        info "Removing Service Account User permissions from ScaleOps SA for: ${sa_trimmed}"
        capture_run gcloud iam service-accounts remove-iam-policy-binding "${sa_trimmed}" \
          --member="serviceAccount:${service_account}" \
          --role="roles/iam.serviceAccountUser" --project="${project_id}" || true
      done
    fi

    info "Removing custom role: ${role_id}"
    capture_run gcloud iam roles delete "${role_id}" --project="${project_id}" || true

  fi
  if [ "$COMMAND" == "cost" ] || [ "$COMMAND" == "all" ]; then
    if ! verbose_run gcloud config set project "${cost_project_id}" --quiet; then
      error "Failed to set gcloud project"
      exit 1
    fi
    
    info "Removing BigQuery permissions for service account: ${service_account}"
    
    check_bq_cli
    check_jq_cli
    
    dataset_config_file="/tmp/bq_dataset_config_remove_$(date +%s).json"
    
    dataset_output=$(bq show --format=prettyjson "${cost_project_id}:${bq_dataset_name}" 2>/dev/null)
    if [ $? -eq 0 ] && [ -n "$dataset_output" ]; then
      echo "$dataset_output" > "${dataset_config_file}"
      # Remove service account from access array
      jq --arg sa "${service_account}" '.access = [.access[] | select(.userByEmail != $sa)]' "${dataset_config_file}" > "${dataset_config_file}.tmp" && mv "${dataset_config_file}.tmp" "${dataset_config_file}"
      
      # Update dataset
      capture_run bq update --source="${dataset_config_file}" "${cost_project_id}:${bq_dataset_name}" || true
    fi
    
    rm -f "${dataset_config_file}"

    # Remove BigQuery Job User role
    info "Removing BigQuery Job User role from service account"
    capture_run gcloud projects remove-iam-policy-binding "${cost_project_id}" \
      --member="serviceAccount:${service_account}" \
      --role="roles/bigquery.jobUser" || true
  fi
  info "Removal complete – all resources deleted for command: $COMMAND"
  exit 0
fi

########################################################
# Validate Organization Policies
########################################################
if [ "$skip_checks" != "true" ]; then
  # if project_id is not set, use the cost_project_id (when running the script for only cost integration)
  if [ -z "${project_id}" ]; then
    project_id="${cost_project_id}"
  fi
  info "Validating organization policies"
  if [ "$create_service_account" == "true" ]; then
    check_organization_create_policies "${project_id}"
  else
    check_organization_domain_policies "${project_id}"
  fi
  verbose "Organization policy checks completed successfully"
else
  info "Skipping organization policy checks"
fi

########################################################
# Create Service Account
########################################################
# Create service account if requested
if [ "$create_service_account" == "true" ]; then
  # if project_id is not set, use the cost_project_id (when running the script for only cost integration)
  if [ -z "${project_id}" ]; then
    project_id="${cost_project_id}"
  fi

  service_account="${service_account_name}@${project_id}.iam.gserviceaccount.com"

  info "Creating service account: ${service_account}"
  if ! silent_run gcloud iam service-accounts create "${service_account_name}" \
    --display-name="ScaleOps Integration Service Account" \
    --description="Service account for ScaleOps cloud integration" --project="${project_id}"; then
    error "Failed to create service account"
    exit 1
  fi
  sleep 5
    
  # Only create keys if not using workload identity
  if [ "$use_workload_identity" != "true" ]; then
    key_file="/tmp/scaleops-sa-key-$(date +%s).json"
    if ! silent_run gcloud iam service-accounts keys create "${key_file}" \
      --iam-account="${service_account}" \
      --project="${project_id}"; then
      error "Failed to create service account keys"
      exit 1
    fi
      
    service_account_credentials=$(cat "${key_file}")
    rm -f "${key_file}"
  fi
  info "Service account ${service_account} created"
fi


########################################################
# Cost Integration
########################################################
run_cost_integration() {

  info "Setting up GCP Cost Integration"
  if ! verbose_run gcloud config set project "${cost_project_id}" --quiet; then
    error "Failed to set gcloud project"
    exit 1
  fi
  # Grant BigQuery permissions to service account
  check_bq_cli
  check_jq_cli
  
  dataset_config_file="/tmp/bq_dataset_config_$(date +%s).json"
  
  # Get current dataset configuration
  info "Checking existing BigQuery dataset permissions"
  dataset_output=$(bq show --format=prettyjson "${cost_project_id}:${bq_dataset_name}" 2>/dev/null)
  if [ $? -ne 0 ] || [ -z "$dataset_output" ]; then
    error "Failed to get dataset configuration"
    exit 1
  fi
  echo "$dataset_output" > "${dataset_config_file}"
  
  # Check if service account already has READER access
  has_reader=$(jq -r --arg sa "${service_account}" '.access[] | select(.userByEmail == $sa and .role == "READER") | .role' "${dataset_config_file}" | wc -l | tr -d ' ')
  
  permissions_updated=false
  
  if [ "$has_reader" -eq 0 ]; then
    info "Adding BigQuery Data Viewer permissions for dataset: ${bq_dataset_name}"
    # Add READER access for the service account
    jq --arg sa "${service_account}" '.access += [{"role": "READER", "userByEmail": $sa}]' "${dataset_config_file}" > "${dataset_config_file}.tmp" && mv "${dataset_config_file}.tmp" "${dataset_config_file}"
    permissions_updated=true
  else
    info "BigQuery Data Viewer permissions already exist for dataset - skipping"
  fi
  
  # Update dataset if permissions were added
  if [ "$permissions_updated" = true ]; then
    info "Updating BigQuery dataset permissions"
    if ! bq update --source="${dataset_config_file}" "${cost_project_id}:${bq_dataset_name}"; then
      error "Failed to update BigQuery dataset permissions"
      rm -f "${dataset_config_file}"
      exit 1
    fi
  fi
  
  # Cleanup
  rm -f "${dataset_config_file}"

  # Grant BigQuery Job User role for query execution
  job_user_binding_exists=$(capture_run gcloud projects get-iam-policy "${cost_project_id}" \
    --flatten="bindings[].members" \
    --format="table(bindings.role)" \
    --filter="bindings.role:roles/bigquery.jobUser AND bindings.members:serviceAccount:${service_account}")

  if [ "$job_user_binding_exists" == "" ]; then
    info "Granting BigQuery Job User role to service account for query execution"
    if ! silent_run gcloud projects add-iam-policy-binding "${cost_project_id}" \
      --member="serviceAccount:${service_account}" \
      --role="roles/bigquery.jobUser" --condition=None; then
      error "Failed to grant BigQuery Job User role"
      exit 1
    fi
  else
    info "BigQuery Job User role already granted - skipping"
  fi

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

  if [ "$COMMAND" != "all" ] && ([ "$create_service_account" == "true" ] || [ "$use_workload_identity" == "true" ]); then
    echo
    info "Please add the following values to your helm chart:"
    echo
    helm "cloudBillingIntegration:"
    helm "  google:"
    helm "    enabled: true"
    helm "    projectId: ${cost_project_id}"
    helm "    datasetName: ${bq_dataset_name}"
    if [ "$use_workload_identity" == "true" ]; then
      helm ""
      helm "agent:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "dashboard:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "recommender:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "updater:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
    elif [ "$create_service_account" == "true" ]; then
      helm "    serviceAccountCredentials: |"
      echo -ne "\033[1;34m"
      echo "${service_account_credentials}" | sed 's/^/      /'
      echo -ne "\033[0m"
    fi
    echo
  fi

  info "GCP Cost Integration setup complete!"
} 

########################################################
# Cloud Node Integration
########################################################
run_cloud_node_integration() {
  info "Setting up GCP Cloud Node Integration"

  if ! verbose_run gcloud config set project "${project_id}" --quiet; then
    error "Failed to set gcloud project for node integration"
    exit 1
  fi

  if [ "$minimal" == "true" ]; then
    # Minimal mode: use custom role with minimal permissions only
    info "Using minimal permissions mode (custom role only)"

    role_exists=$(capture_run gcloud iam roles describe "${role_id}" --project="${project_id}" --format="value(name)" 2>/dev/null || echo "")

    if [ -z "$role_exists" ]; then
        info "Creating custom role ${role_id}"
        permissions_list=$(printf '%s,' "${MINIMAL_PERMISSIONS[@]}")
        if ! silent_run gcloud iam roles create "${role_id}" --title="${role_display_name}" --project="${project_id}" --permissions="${permissions_list%,}" --description="Custom role for ScaleOps Cloud Integration (minimal)"; then
          error "Failed to create custom role ${role_id}"
          exit 1
        fi
        sleep 5
    else
      info "Custom role ${role_id} already exists - checking permissions"
      # Get current permissions and compare with MINIMAL_PERMISSIONS
      current_permissions=$(verbose_run gcloud iam roles describe "${role_id}" --format="value(includedPermissions[].join(','))" --project="${project_id}" 2>/dev/null | tr ',' '\n' | sort)
      expected_permissions=$(printf '%s\n' "${MINIMAL_PERMISSIONS[@]}" | sort)

      if [ "$current_permissions" != "$expected_permissions" ]; then
          info "Role permissions differ from required - updating permissions"
          permissions_list=$(printf '%s,' "${MINIMAL_PERMISSIONS[@]}")

          if ! silent_run gcloud iam roles update "${role_id}" --project="${project_id}" --permissions="${permissions_list%,}"; then
            error "Failed to update custom role ${role_id}"
            exit 1
          fi
      fi
    fi

    if [ -n "$role_exists" ]; then
      role_deleted=$(capture_run gcloud iam roles describe "${role_id}" --project="${project_id}" --format="value(deleted)" 2>/dev/null || echo "")

      if [ "$role_deleted" == "True" ]; then
          capture_run gcloud iam roles undelete "${role_id}" --project="${project_id}" || true
      fi
    fi

    role_path=$(get_role_path)

    binding_exists=$(capture_run gcloud projects get-iam-policy "${project_id}" \
      --flatten="bindings[].members" \
      --format="table(bindings.role)" \
      --filter="bindings.role=${role_path} AND bindings.members=serviceAccount:${service_account}")

    if [ "$binding_exists" == "" ]; then
        info "Creating IAM policy binding for role: ${role_path} and service account: ${service_account}"
        if ! silent_run gcloud projects add-iam-policy-binding "${project_id}" \
          --member="serviceAccount:${service_account}" \
          --role="${role_path}" --project="${project_id}" --condition=None; then
          error "Failed to create IAM policy binding"
          exit 1
        fi
    else
      info "IAM policy binding already exists for service account: ${service_account} - skipping creation"
    fi
  else
    # Default mode: use predefined viewer roles + custom role for write permissions
    info "Using full permissions mode (predefined viewer roles + custom write role)"

    # Grant predefined viewer roles
    for viewer_role in "${VIEWER_ROLES[@]}"; do
      binding_exists=$(capture_run gcloud projects get-iam-policy "${project_id}" \
        --flatten="bindings[].members" \
        --format="table(bindings.role)" \
        --filter="bindings.role=${viewer_role} AND bindings.members=serviceAccount:${service_account}")

      if [ "$binding_exists" == "" ]; then
          if ! silent_run gcloud projects add-iam-policy-binding "${project_id}" \
            --member="serviceAccount:${service_account}" \
            --role="${viewer_role}" --project="${project_id}" --condition=None; then
            error "Failed to grant role ${viewer_role}"
            exit 1
          fi
      else
        info "Role ${viewer_role} already granted - skipping"
      fi
    done

    # Create custom role for write permissions
    role_exists=$(capture_run gcloud iam roles describe "${role_id}" --project="${project_id}" --format="value(name)" 2>/dev/null || echo "")

    if [ -z "$role_exists" ]; then
        info "Creating custom role ${role_id} for write permissions"
        permissions_list=$(printf '%s,' "${WRITE_PERMISSIONS[@]}")
        if ! silent_run gcloud iam roles create "${role_id}" --title="${role_display_name}" --project="${project_id}" --permissions="${permissions_list%,}" --description="Custom role for ScaleOps Cloud Integration (write permissions)"; then
          error "Failed to create custom role ${role_id}"
          exit 1
        fi
        sleep 5
    else
      info "Custom role ${role_id} already exists - checking permissions"
      # Get current permissions and compare with WRITE_PERMISSIONS
      current_permissions=$(verbose_run gcloud iam roles describe "${role_id}" --format="value(includedPermissions[].join(','))" --project="${project_id}" 2>/dev/null | tr ',' '\n' | sort)
      expected_permissions=$(printf '%s\n' "${WRITE_PERMISSIONS[@]}" | sort)

      if [ "$current_permissions" != "$expected_permissions" ]; then
          info "Role permissions differ from required - updating permissions"
          permissions_list=$(printf '%s,' "${WRITE_PERMISSIONS[@]}")

          if ! silent_run gcloud iam roles update "${role_id}" --project="${project_id}" --permissions="${permissions_list%,}"; then
            error "Failed to update custom role ${role_id}"
            exit 1
          fi
      fi
    fi

    if [ -n "$role_exists" ]; then
      role_deleted=$(capture_run gcloud iam roles describe "${role_id}" --project="${project_id}" --format="value(deleted)" 2>/dev/null || echo "")

      if [ "$role_deleted" == "True" ]; then
          capture_run gcloud iam roles undelete "${role_id}" --project="${project_id}" || true
      fi
    fi

    role_path=$(get_role_path)

    binding_exists=$(capture_run gcloud projects get-iam-policy "${project_id}" \
      --flatten="bindings[].members" \
      --format="table(bindings.role)" \
      --filter="bindings.role=${role_path} AND bindings.members=serviceAccount:${service_account}")

    if [ "$binding_exists" == "" ]; then
        info "Creating IAM policy binding for custom write role: ${role_path}"
        if ! silent_run gcloud projects add-iam-policy-binding "${project_id}" \
          --member="serviceAccount:${service_account}" \
          --role="${role_path}" --project="${project_id}" --condition=None; then
          error "Failed to create IAM policy binding"
          exit 1
        fi
    else
      info "IAM policy binding already exists for custom write role - skipping"
    fi
  fi

  # Grant impersonation permissions to additional service accounts
  if [ -n "${allow_service_accounts}" ]; then
    IFS=',' read -ra SA_ARRAY <<< "${allow_service_accounts}"
    for sa in "${SA_ARRAY[@]}"; do
      if [ "$(echo "${sa}" | xargs)" == "default" ]; then
        sa="$(capture_run gcloud projects describe "${project_id}" --format="value(projectNumber)")-compute@developer.gserviceaccount.com"
      fi
      impersonate_service_account "${project_id}" "${service_account}" "${sa}"
    done
  fi

  if [ -n "${cluster}" ]; then
    clusters=$(capture_run gcloud container clusters list --format="csv[separator=',',no-heading](name,location)" --project="${project_id}" --filter="name=${cluster}")
    if [ -z "$clusters" ]; then
      error "No clusters found in project ${project_id}"
      exit 1
    fi
    for cl in $clusters; do
      while IFS=$',' read -r name location; do
        impersonate_cluster_service_account "${project_id}" "${service_account}" "${name}" "${location}"
      done <<< "$cl"
    done
  fi

  if [ "$all_clusters" == "true" ]; then
    clusters=$(capture_run gcloud container clusters list --format="csv[separator=',',no-heading](name,location)" --project="${project_id}")
    if [ -z "$clusters" ]; then
      error "No clusters found in project ${project_id}"
      exit 1
    fi
    for cl in $clusters; do
      while IFS=$',' read -r name location; do
        impersonate_cluster_service_account "${project_id}" "${service_account}" "${name}" "${location}"
      done <<< "$cl"
    done
  fi

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

  if [ "$COMMAND" != "all" ] && ([ "$create_service_account" == "true" ] || [ "$use_workload_identity" == "true" ]); then
    echo
    info "Please add the following values to your helm chart:"
    echo
    helm "cloudNodeIntegration:"
    helm "  google:"
    helm "    enabled: true"
    if [ "$use_workload_identity" == "true" ]; then
      helm ""
      helm "agent:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "dashboard:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "recommender:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "updater:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
    elif [ "$create_service_account" == "true" ]; then
      helm "    serviceAccountCredentials: |"
      echo -ne "\033[1;34m"
      echo "${service_account_credentials}" | sed 's/^/      /'
      echo -ne "\033[0m"
    fi
    echo
  fi

  info "GCP Cloud Node Integration setup complete!"
}


show_combined_helm_values() {
  if [ "$COMMAND" == "all" ] && ([ "$create_service_account" == "true" ] || [ "$use_workload_identity" == "true" ]); then
    echo
    info "Combined helm values for both cost and node integration:"
    echo
    helm "cloudBillingIntegration:"
    helm "  google:"
    helm "    enabled: true"
    helm "    projectId: ${cost_project_id}"
    helm "    datasetName: ${bq_dataset_name}"
    if [ "$use_workload_identity" == "false" ] && [ "$create_service_account" == "true" ]; then
      helm "    serviceAccountCredentials: |"
      echo -ne "\033[1;34m"
      echo "${service_account_credentials}" | sed 's/^/      /'
      echo -ne "\033[0m"
    fi
    echo
    helm "cloudNodeIntegration:"
    helm "  google:"
    helm "    enabled: true"
    if [ "$use_workload_identity" == "false" ] && [ "$create_service_account" == "true" ]; then
      helm "    serviceAccountCredentials: |"
      echo -ne "\033[1;34m"
      echo "${service_account_credentials}" | sed 's/^/      /'
      echo -ne "\033[0m"
    fi
    echo
    if [ "$use_workload_identity" == "true" ]; then
      helm "agent:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "dashboard:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "recommender:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      helm ""
      helm "updater:"
      helm "  serviceAccount:"
      helm "    annotations:"
      helm "      iam.gke.io/gcp-service-account: ${service_account}"
      echo
    fi
  fi
}

case "$COMMAND" in
  cost)
    run_cost_integration
    ;;
  node-integration)
    run_cloud_node_integration
    ;;
  all)
    run_cost_integration
    run_cloud_node_integration
    setup_workload_identity
    show_combined_helm_values
    ;;
  *)
    error "Unknown command: $COMMAND"
    exit 1
    ;;
esac