#!/usr/bin/env bash
# Assign required Entra directory roles to ComputerGuru MSP service principals
# in a newly-consented customer tenant, and programmatically consent all other
# ComputerGuru apps so only Tenant Admin requires a manual customer consent click.
#
# Usage: onboard-tenant.sh <domain-or-tenant-id> [--dry-run]
#
# What this script does:
#   1. Resolves the tenant ID
#   2. Acquires a Tenant Admin token (fails gracefully if not consented)
#   3. Creates SPs for Security Investigator, Exchange Operator, User Manager,
#      and Defender Add-on (equivalent to admin consent for each)
#   4. Grants all required Graph/EXO/Defender appRoleAssignments to each SP
#   5. Assigns required directory roles to each SP
#   6. Prints a final status table
#
# Exit codes:
#   0  all roles present or successfully assigned
#   1  resolve failure
#   2  Tenant Admin not consented (consent URL printed)
#   3  vault error
#   10 partial failure (some roles could not be assigned)
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
__ROOT="${CLAUDETOOLS_ROOT:-$(cd "$SCRIPT_DIR/../../../.." && pwd)}"

TARGET="${1:?Usage: onboard-tenant.sh <domain-or-tenant-id> [--dry-run]}"
DRY_RUN=false
[[ "${2:-}" == "--dry-run" ]] && DRY_RUN=true

# ── App IDs ───────────────────────────────────────────────────────────────────
APP_SEC_INV="bfbc12a4-f0dd-4e12-b06d-997e7271e10c"
APP_EXCH_OP="b43e7342-5b4b-492f-890f-bb5a4f7f40e9"
APP_USER_MGR="64fac46b-8b44-41ad-93ee-7da03927576c"
APP_TENANT_ADMIN="709e6eed-0711-4875-9c44-2d3518c47063"
APP_DEFENDER="dbf8ad1a-54f4-4bb8-8a9e-ea5b9634635b"

# ── Resource app IDs (well-known Microsoft multi-tenant apps) ─────────────────
GRAPH_APP_ID="00000003-0000-0000-c000-000000000000"
EXO_APP_ID="00000002-0000-0ff1-ce00-000000000000"
DEFENDER_APP_ID="fc780465-2017-40d4-a0c5-307022471b92"

# ── Directory role GUIDs ──────────────────────────────────────────────────────
ROLE_EXCHANGE_ADMIN="29232cdf-9323-42fd-ade2-1d097af3e4de"
ROLE_USER_ADMIN="fe930be7-5e62-47db-91af-98c3a49a38b1"
ROLE_AUTH_ADMIN="c4e39bd9-1100-46d3-8c65-fb160da0071f"
ROLE_CA_ADMIN="b1be1c3e-b65d-4f19-8427-f6fa0d97feb9"

# ── Graph appRole GUIDs per app (from requiredResourceAccess in home tenant) ──
# Security Investigator — Graph
SEC_INV_GRAPH_ROLES=(
  "df021288-bdef-4463-88db-98f22de89214"
  "b0afded3-3588-46d8-8b3d-9842eff778da"
  "7ab1d382-f21e-4acd-a863-ba3e13f7da61"
  "40f97065-369a-49f4-947c-6a255697ae91"
  "810c84a8-4a9e-49e6-bf7d-12d183f40d01"
  "9a5d68dd-52b0-4cc2-bd40-abcf44ac3a30"
  "38d9df27-64da-44fd-b7c5-a6fbac20248f"
  "dc5007c0-2d7d-4c42-879c-2dab87571379"
  "246dd0d5-5bd0-4def-940b-0421030a5b68"
  "498476ce-e0fe-48b0-b801-37ba7e2685c6"
  "57f1cf28-c0c4-4ec3-9a30-19a2eaaf2f6e"  # BitLockerKey.Read.All
)
# Security Investigator — Exchange Online
SEC_INV_EXO_ROLES=(
  "dc890d15-9560-4a4c-9b7f-a736ec74ec40"
)

# Exchange Operator — Graph
EXCH_OP_GRAPH_ROLES=(
  "df021288-bdef-4463-88db-98f22de89214"
  "6931bccd-447a-43d1-b442-00a195474933"
  "e2a3a72e-5f79-4c64-b1b1-878b674786c9"
  "77f3a031-c388-4f99-b373-dc68676a979e"
  "498476ce-e0fe-48b0-b801-37ba7e2685c6"
)
# Exchange Operator — Exchange Online
EXCH_OP_EXO_ROLES=(
  "dc890d15-9560-4a4c-9b7f-a736ec74ec40"
  "dc50a0fb-09a3-484d-be87-e023b12c6440"
)

# User Manager — Graph only
USER_MGR_GRAPH_ROLES=(
  "741f803b-c850-494e-b5df-cde7c675a1ca"
  "19dbc75e-c2e2-444c-a770-ec69d8559fc7"
  "62a82d76-70ea-41e2-9197-370581804d09"
  "50483e42-d915-4231-9639-7fdb7fd190e5"
  "77f3a031-c388-4f99-b373-dc68676a979e"
  "498476ce-e0fe-48b0-b801-37ba7e2685c6"
)

# Defender Add-on — Graph
DEFENDER_GRAPH_ROLES=(
  "bf394140-e372-4bf9-a898-299cfc7564e5"
)
# Defender Add-on — Defender ATP
DEFENDER_ATP_ROLES=(
  "71fe6b80-7034-4028-9ed8-0f316df9c3ff"
  "ea8291d3-4b9a-44b5-bc3a-6cea3026dc79"
  "93489bf5-0fbc-4f2d-b901-33f2fe08ff05"
  "41269fc5-d04d-4bfd-bce7-43a51cea049a"
  "528ca142-c849-4a5b-935e-10b8b9c38a84"
)

CONSENT_BASE="https://login.microsoftonline.com"
CONSENT_REDIRECT="https://azcomputerguru.com"

# ── Helper: print consent URLs for all apps ───────────────────────────────────
print_consent_urls() {
  local tenant_id="$1"
  echo ""
  echo "[INFO] Consent URLs for tenant $tenant_id (provide to customer Global Admin):"
  echo "  [1] Tenant Admin (consent FIRST — needed for programmatic onboarding of all other apps):"
  echo "  ${CONSENT_BASE}/${tenant_id}/adminconsent?client_id=${APP_TENANT_ADMIN}&redirect_uri=${CONSENT_REDIRECT}&prompt=consent"
  echo ""
  echo "  After the admin accepts Tenant Admin consent, run:"
  echo "  bash $0 $TARGET"
  echo ""
  echo "  The script will then automatically consent all other apps in the suite."
  echo ""
  echo "  (Optional — only needed if Tenant Admin consent failed for individual apps):"
  echo "  [2] Security Investigator: ${CONSENT_BASE}/${tenant_id}/adminconsent?client_id=${APP_SEC_INV}&redirect_uri=${CONSENT_REDIRECT}&prompt=consent"
  echo "  [3] Exchange Operator:     ${CONSENT_BASE}/${tenant_id}/adminconsent?client_id=${APP_EXCH_OP}&redirect_uri=${CONSENT_REDIRECT}&prompt=consent"
  echo "  [4] User Manager:          ${CONSENT_BASE}/${tenant_id}/adminconsent?client_id=${APP_USER_MGR}&redirect_uri=${CONSENT_REDIRECT}&prompt=consent"
  echo "  [5] Defender Add-on (MDE-licensed tenants only): ${CONSENT_BASE}/${tenant_id}/adminconsent?client_id=${APP_DEFENDER}&redirect_uri=${CONSENT_REDIRECT}&prompt=consent"
}

# ── Helper: get SP OID in tenant ──────────────────────────────────────────────
get_sp_oid() {
  local token="$1"
  local app_id="$2"
  local resp
  resp=$(curl -s --max-time 15 \
    -H "Authorization: Bearer $token" \
    -G \
    --data-urlencode "\$filter=appId eq '${app_id}'" \
    --data-urlencode "\$select=id,displayName" \
    "https://graph.microsoft.com/v1.0/servicePrincipals")
  echo "$resp" | jq -r '.value[0].id // empty'
}

# ── Helper: create SP for our app if not present ──────────────────────────────
create_sp_if_missing() {
  local token="$1"
  local app_id="$2"
  local app_name="$3"

  local oid
  oid=$(get_sp_oid "$token" "$app_id")

  if [[ -n "$oid" ]]; then
    echo "  [OK] $app_name SP already present: $oid" >&2
    echo "$oid"
    return 0
  fi

  if [[ "$DRY_RUN" == "true" ]]; then
    echo "  [DRY-RUN] Would create SP for $app_name ($app_id)" >&2
    echo ""
    return 0
  fi

  local resp
  resp=$(curl -s --max-time 15 \
    -H "Authorization: Bearer $token" \
    -H "Content-Type: application/json" \
    -X POST \
    "https://graph.microsoft.com/v1.0/servicePrincipals" \
    -d "{\"appId\": \"$app_id\"}")

  local new_oid
  new_oid=$(echo "$resp" | jq -r '.id // empty')

  if [[ -z "$new_oid" ]]; then
    local err_code
    err_code=$(echo "$resp" | jq -r '.error.code // empty')
    if [[ "$err_code" == "Request_MultipleObjectsWithSameKeyValue" ]] || echo "$resp" | grep -qi "conflicting"; then
      oid=$(get_sp_oid "$token" "$app_id")
      echo "  [OK] $app_name SP already exists: $oid" >&2
      echo "$oid"
      return 0
    fi
    echo "  [ERROR] Failed to create SP for $app_name: $(echo "$resp" | jq -r '.error.message // empty')" >&2
    bash "$__ROOT/.claude/scripts/log-skill-error.sh" "remediation-tool" "onboard-tenant: failed to create service principal" --context "app=$app_name appId=$app_id msg=$(echo "$resp" | jq -r '.error.message // empty' | head -c 80)" >/dev/null 2>&1 || true
    return 1
  fi

  echo "  [CREATED] $app_name SP: $new_oid" >&2
  # Brief pause for Graph replication before granting appRoleAssignments
  sleep 5
  echo "$new_oid"
}

# ── Helper: grant single appRoleAssignment (idempotent) ───────────────────────
grant_app_role() {
  local token="$1"
  local principal_oid="$2"
  local resource_oid="$3"
  local role_id="$4"

  # Check if already granted
  local already
  already=$(curl -s --max-time 15 \
    -H "Authorization: Bearer $token" \
    "https://graph.microsoft.com/v1.0/servicePrincipals/$principal_oid/appRoleAssignments" \
    | jq --arg rid "$role_id" '([.value[]? | select(.appRoleId == $rid)] | length) > 0')

  if [[ "$already" == "true" ]]; then
    return 0
  fi

  if [[ "$DRY_RUN" == "true" ]]; then
    echo "    [DRY-RUN] Would grant role $role_id"
    return 0
  fi

  local body
  body=$(jq -n \
    --arg principal "$principal_oid" \
    --arg resource "$resource_oid" \
    --arg role "$role_id" \
    '{"principalId": $principal, "resourceId": $resource, "appRoleId": $role}')

  local resp
  resp=$(curl -s --max-time 15 \
    -H "Authorization: Bearer $token" \
    -H "Content-Type: application/json" \
    -X POST \
    "https://graph.microsoft.com/v1.0/servicePrincipals/$principal_oid/appRoleAssignments" \
    -d "$body")

  local granted_id
  granted_id=$(echo "$resp" | jq -r '.id // empty')

  if [[ -z "$granted_id" ]]; then
    local err_code
    err_code=$(echo "$resp" | jq -r '.error.code // empty')
    if [[ "$err_code" == "Request_MultipleObjectsWithSameKeyValue" ]] || echo "$resp" | grep -qi "conflicting"; then
      return 0
    fi
    echo "    [ERROR] grant_app_role failed for $role_id: $(echo "$resp" | jq -r '.error.message // "unknown"')" >&2
    bash "$__ROOT/.claude/scripts/log-skill-error.sh" "remediation-tool" "onboard-tenant: grant_app_role appRoleAssignment failed" --context "role=$role_id msg=$(echo "$resp" | jq -r '.error.message // "unknown"' | head -c 80)" >/dev/null 2>&1 || true
    return 1
  fi
}

# ── Helper: consent app + grant all permissions ───────────────────────────────
# Usage: consent_app <token> <app_id> <app_name> \
#          <graph_sp_oid> <exo_sp_oid_or_empty> <defender_sp_oid_or_empty> \
#          <graph_roles_varname> [<exo_roles_varname>] [<atp_roles_varname>]
consent_app() {
  local token="$1"
  local app_id="$2"
  local app_name="$3"
  local graph_sp_oid="$4"
  local exo_sp_oid="${5:-}"
  local defender_sp_oid="${6:-}"
  local graph_roles_varname="$7"
  local exo_roles_varname="${8:-}"
  local atp_roles_varname="${9:-}"

  echo ""
  echo "[CONSENT] $app_name ($app_id)"

  # Create SP (or confirm existing)
  local sp_oid
  sp_oid=$(create_sp_if_missing "$token" "$app_id" "$app_name")
  if [[ -z "$sp_oid" ]]; then
    echo "  [ERROR] Cannot proceed — SP creation failed" >&2
    return 1
  fi

  local grant_errors=0

  # Grant Graph permissions
  eval "local graph_roles=(\"\${${graph_roles_varname}[@]}\")"
  local granted=0 skipped=0 errors=0
  for role_id in "${graph_roles[@]}"; do
    if grant_app_role "$token" "$sp_oid" "$graph_sp_oid" "$role_id"; then
      ((granted++)) || true
    else
      ((errors++)) || true
      ((grant_errors++)) || true
    fi
  done
  echo "  Graph permissions: ${#graph_roles[@]} total — $errors error(s)"

  # Grant Exchange Online permissions
  if [[ -n "$exo_roles_varname" ]] && [[ -n "$exo_sp_oid" ]]; then
    eval "local exo_roles=(\"\${${exo_roles_varname}[@]}\")"
    local exo_errors=0
    for role_id in "${exo_roles[@]}"; do
      if ! grant_app_role "$token" "$sp_oid" "$exo_sp_oid" "$role_id"; then
        ((exo_errors++)) || true
        ((grant_errors++)) || true
      fi
    done
    echo "  Exchange Online permissions: ${#exo_roles[@]} total — $exo_errors error(s)"
  elif [[ -n "$exo_roles_varname" ]] && [[ -z "$exo_sp_oid" ]]; then
    echo "  [WARNING] Exchange Online SP not found — EXO permissions skipped"
  fi

  # Grant Defender ATP permissions
  if [[ -n "$atp_roles_varname" ]] && [[ -n "$defender_sp_oid" ]]; then
    eval "local atp_roles=(\"\${${atp_roles_varname}[@]}\")"
    local atp_errors=0
    for role_id in "${atp_roles[@]}"; do
      if ! grant_app_role "$token" "$sp_oid" "$defender_sp_oid" "$role_id"; then
        ((atp_errors++)) || true
        ((grant_errors++)) || true
      fi
    done
    echo "  Defender ATP permissions: ${#atp_roles[@]} total — $atp_errors error(s)"
  elif [[ -n "$atp_roles_varname" ]] && [[ -z "$defender_sp_oid" ]]; then
    echo "  [INFO] Defender ATP SP not found — tenant likely not MDE-licensed, skipping"
  fi

  if [[ $grant_errors -eq 0 ]]; then
    echo "  [OK] $app_name fully consented and permissions granted"
    return 0
  else
    echo "  [WARNING] $app_name consent completed with $grant_errors permission error(s)"
    return 1
  fi
}

# ── Helper: check if directory role already assigned ─────────────────────────
# NOTE: The "MISSING -> ASSIGNING" noise was NOT PIM, as previously suspected — the
# root cause was an unencoded space in the $filter (now %20-encoded), which made Graph
# return empty/error and this function always return false. The ACG tenant has no Entra
# ID P2, so PIM is not a factor here. The dual-query idea (also checking
# /roleManagement/directory/roleAssignmentSchedules) remains valid ONLY for P2 tenants
# where roles can be PIM-managed; return true if either query finds the role.
role_assigned() {
  local token="$1"
  local sp_oid="$2"
  local role_id="$3"
  local resp
  resp=$(curl -s --max-time 15 \
    -H "Authorization: Bearer $token" \
    "https://graph.microsoft.com/v1.0/roleManagement/directory/roleAssignments?\$filter=principalId%20eq%20'${sp_oid}'")
  echo "$resp" | jq --arg rid "$role_id" \
    '[.value[] | select(.roleDefinitionId == $rid)] | length > 0'
}

# ── Helper: assign directory role ─────────────────────────────────────────────
assign_role() {
  local token="$1"
  local sp_oid="$2"
  local role_id="$3"
  local role_name="$4"

  if [[ "$DRY_RUN" == "true" ]]; then
    echo "    [DRY-RUN] Would assign $role_name to SP $sp_oid"
    return 0
  fi

  local body
  body=$(jq -n \
    --arg role "$role_id" \
    --arg principal "$sp_oid" \
    '{"roleDefinitionId": $role, "principalId": $principal, "directoryScopeId": "/"}')

  local resp
  resp=$(curl -s --max-time 15 \
    -H "Authorization: Bearer $token" \
    -H "Content-Type: application/json" \
    -X POST \
    "https://graph.microsoft.com/v1.0/roleManagement/directory/roleAssignments" \
    -d "$body")

  local assigned_id
  assigned_id=$(echo "$resp" | jq -r '.id // empty')
  if [[ -z "$assigned_id" ]]; then
    local err_code
    err_code=$(echo "$resp" | jq -r '.error.code // empty')
    if [[ "$err_code" == "Conflict" ]] || [[ "$err_code" == "Request_MultipleObjectsWithSameKeyValue" ]] || \
       echo "$resp" | grep -qi "conflicting object"; then
      echo "    [OK] $role_name already assigned (conflict returned — idempotent)"
      return 0
    fi
    echo "    [ERROR] Failed to assign $role_name" >&2
    echo "    Response: $resp" >&2
    bash "$__ROOT/.claude/scripts/log-skill-error.sh" "remediation-tool" "onboard-tenant: failed to assign directory role" --context "role=$role_name sp=$sp_oid msg=$(echo "$resp" | jq -r '.error.message // empty' | head -c 80)" >/dev/null 2>&1 || true
    return 1
  fi
  echo "    [OK] $role_name assigned (assignment id=$assigned_id)"
}

# ── Step 1: Resolve tenant ────────────────────────────────────────────────────
echo "[INFO] Resolving tenant: $TARGET"
TENANT_ID=$("$SCRIPT_DIR/resolve-tenant.sh" "$TARGET")
if [[ -z "$TENANT_ID" ]]; then
  echo "[ERROR] Could not resolve tenant ID for: $TARGET" >&2
  bash "$__ROOT/.claude/scripts/log-skill-error.sh" "remediation-tool" "onboard-tenant: could not resolve tenant ID" --context "target=$TARGET" >/dev/null 2>&1 || true
  exit 1
fi

DISPLAY_NAME="$TARGET"
echo "[OK] Tenant: $DISPLAY_NAME ($TENANT_ID)"

# ── Step 2: Acquire Tenant Admin token ───────────────────────────────────────
echo "[INFO] Acquiring Tenant Admin token for $TENANT_ID..."
set +e
TENANT_ADMIN_TOKEN_OUT=$("$SCRIPT_DIR/get-token.sh" "$TENANT_ID" "tenant-admin" 2>/tmp/onboard-token-err.txt)
GET_TOKEN_EXIT=$?
TOKEN_ERR=$(cat /tmp/onboard-token-err.txt 2>/dev/null || true)
set -e

if [[ $GET_TOKEN_EXIT -ne 0 ]]; then
  if echo "$TOKEN_ERR" | grep -qi "7000229\|AADSTS7000229\|service principal\|not been authorized\|not found"; then
    echo "[WARNING] Tenant Admin app not yet consented in tenant $TENANT_ID"
    print_consent_urls "$TENANT_ID"
    exit 2
  fi
  echo "[ERROR] Failed to acquire Tenant Admin token (exit $GET_TOKEN_EXIT)" >&2
  echo "$TOKEN_ERR" >&2
  bash "$__ROOT/.claude/scripts/log-skill-error.sh" "remediation-tool" "onboard-tenant: failed to acquire Tenant Admin token" --context "tenant=$TENANT_ID exit=$GET_TOKEN_EXIT" >/dev/null 2>&1 || true
  exit 5
fi
TENANT_ADMIN_TOKEN="$TENANT_ADMIN_TOKEN_OUT"

TA_SP_OID=$(get_sp_oid "$TENANT_ADMIN_TOKEN" "$APP_TENANT_ADMIN")
if [[ -z "$TA_SP_OID" ]]; then
  echo "[WARNING] Tenant Admin SP not found in tenant — app not consented yet"
  print_consent_urls "$TENANT_ID"
  exit 2
fi
echo "[OK] Tenant Admin consented — SP: $TA_SP_OID"

[[ "$DRY_RUN" == "true" ]] && echo "[INFO] --dry-run mode: no changes will be made"

# ── Step 3: Locate resource SPs in customer tenant ───────────────────────────
echo ""
echo "[INFO] Locating resource service principals in tenant..."

GRAPH_SP_OID=$(get_sp_oid "$TENANT_ADMIN_TOKEN" "$GRAPH_APP_ID")
EXO_SP_OID=$(get_sp_oid "$TENANT_ADMIN_TOKEN" "$EXO_APP_ID")
DEFENDER_SP_OID=$(get_sp_oid "$TENANT_ADMIN_TOKEN" "$DEFENDER_APP_ID")

[[ -n "$GRAPH_SP_OID" ]]    && echo "  [OK] Microsoft Graph SP:      $GRAPH_SP_OID" || echo "  [ERROR] Microsoft Graph SP not found — cannot proceed" >&2
[[ -n "$EXO_SP_OID" ]]      && echo "  [OK] Exchange Online SP:      $EXO_SP_OID"   || echo "  [WARNING] Exchange Online SP not found (no Exchange license?)"
[[ -n "$DEFENDER_SP_OID" ]] && echo "  [OK] Defender ATP SP:         $DEFENDER_SP_OID" || echo "  [INFO] Defender ATP SP not found (tenant likely not MDE-licensed)"

if [[ -z "$GRAPH_SP_OID" ]]; then
  echo "[ERROR] Microsoft Graph SP missing — cannot grant app permissions" >&2
  bash "$__ROOT/.claude/scripts/log-skill-error.sh" "remediation-tool" "onboard-tenant: Microsoft Graph SP missing in tenant — cannot grant app permissions" --context "tenant=$TENANT_ID" >/dev/null 2>&1 || true
  exit 1
fi

# ── Step 3.5: Backfill Tenant Admin SP — Policy.Read.All ─────────────────────
# Microsoft tightened LIST /identity/conditionalAccess/* to require Policy.Read.All
# (Policy.ReadWrite.ConditionalAccess no longer accepted for reads). Added to manifest
# 2026-04-29; tenants consented before that need backfill.
POLICY_READ_ALL_ID="246dd0d5-5bd0-4def-940b-0421030a5b68"
HAS_POLICY_READ=$(curl -s --max-time 15 \
  -H "Authorization: Bearer $TENANT_ADMIN_TOKEN" \
  "https://graph.microsoft.com/v1.0/servicePrincipals/$TA_SP_OID/appRoleAssignments" \
  | jq --arg rid "$POLICY_READ_ALL_ID" '[.value[]? | select(.appRoleId == $rid)] | length > 0')

if [[ "$HAS_POLICY_READ" != "true" ]]; then
  echo ""
  echo "[INFO] Tenant Admin SP missing Policy.Read.All — backfilling..."
  if [[ "$DRY_RUN" == "true" ]]; then
    echo "  [DRY-RUN] Would grant Policy.Read.All ($POLICY_READ_ALL_ID) to Tenant Admin SP"
  elif grant_app_role "$TENANT_ADMIN_TOKEN" "$TA_SP_OID" "$GRAPH_SP_OID" "$POLICY_READ_ALL_ID"; then
    echo "  [OK] Policy.Read.All granted to Tenant Admin SP"
  else
    echo "  [WARNING] Could not grant Policy.Read.All programmatically — customer re-consent may be needed"
    echo "    Re-consent URL: ${CONSENT_BASE}/${TENANT_ID}/adminconsent?client_id=${APP_TENANT_ADMIN}&redirect_uri=${CONSENT_REDIRECT}&prompt=consent"
  fi
fi

# ── Step 4: Programmatic consent — create SPs + grant appRoleAssignments ──────
echo ""
echo "[INFO] Consenting app suite in tenant (programmatic — no customer click needed)..."

CONSENT_PARTIAL=false

consent_app "$TENANT_ADMIN_TOKEN" "$APP_SEC_INV"  "Security Investigator" \
  "$GRAPH_SP_OID" "$EXO_SP_OID" "" \
  "SEC_INV_GRAPH_ROLES" "SEC_INV_EXO_ROLES" \
  || CONSENT_PARTIAL=true

consent_app "$TENANT_ADMIN_TOKEN" "$APP_EXCH_OP"  "Exchange Operator" \
  "$GRAPH_SP_OID" "$EXO_SP_OID" "" \
  "EXCH_OP_GRAPH_ROLES" "EXCH_OP_EXO_ROLES" \
  || CONSENT_PARTIAL=true

consent_app "$TENANT_ADMIN_TOKEN" "$APP_USER_MGR"  "User Manager" \
  "$GRAPH_SP_OID" "" "" \
  "USER_MGR_GRAPH_ROLES" \
  || CONSENT_PARTIAL=true

if [[ -n "$DEFENDER_SP_OID" ]]; then
  consent_app "$TENANT_ADMIN_TOKEN" "$APP_DEFENDER" "Defender Add-on" \
    "$GRAPH_SP_OID" "" "$DEFENDER_SP_OID" \
    "DEFENDER_GRAPH_ROLES" "" "DEFENDER_ATP_ROLES" \
    || CONSENT_PARTIAL=true
else
  echo ""
  echo "[INFO] Skipping Defender Add-on consent (no MDE license detected)"
fi

# ── Step 5: Check/assign directory roles per SP ───────────────────────────────
declare -A STATUS_MAP

echo ""
echo "[INFO] Checking and assigning directory roles..."

SEC_INV_OID=$(get_sp_oid "$TENANT_ADMIN_TOKEN" "$APP_SEC_INV")
EXCH_OP_OID=$(get_sp_oid "$TENANT_ADMIN_TOKEN" "$APP_EXCH_OP")
USER_MGR_OID=$(get_sp_oid "$TENANT_ADMIN_TOKEN" "$APP_USER_MGR")

PARTIAL_FAILURE=false

# Security Investigator -> Exchange Administrator
if [[ -z "$SEC_INV_OID" ]]; then
  echo "[WARNING] Security Investigator SP still not found after consent attempt"
  STATUS_MAP["Security Investigator:Exchange Administrator"]="MISSING SP"
else
  echo ""
  echo "[CHECK] Security Investigator SP: $SEC_INV_OID"
  IS_PRESENT=$(role_assigned "$TENANT_ADMIN_TOKEN" "$SEC_INV_OID" "$ROLE_EXCHANGE_ADMIN")
  if [[ "$IS_PRESENT" == "true" ]]; then
    echo "  Exchange Administrator: PRESENT"
    STATUS_MAP["Security Investigator:Exchange Administrator"]="OK"
  else
    echo "  Exchange Administrator: MISSING -> ASSIGNING..."
    if assign_role "$TENANT_ADMIN_TOKEN" "$SEC_INV_OID" "$ROLE_EXCHANGE_ADMIN" "Exchange Administrator"; then
      STATUS_MAP["Security Investigator:Exchange Administrator"]=$( [[ "$DRY_RUN" == "true" ]] && echo "DRY-RUN" || echo "ASSIGNED" )
    else
      STATUS_MAP["Security Investigator:Exchange Administrator"]="ERROR"
      PARTIAL_FAILURE=true
    fi
  fi
fi

# Exchange Operator -> Exchange Administrator
if [[ -z "$EXCH_OP_OID" ]]; then
  echo "[WARNING] Exchange Operator SP still not found after consent attempt"
  STATUS_MAP["Exchange Operator:Exchange Administrator"]="MISSING SP"
else
  echo ""
  echo "[CHECK] Exchange Operator SP: $EXCH_OP_OID"
  IS_PRESENT=$(role_assigned "$TENANT_ADMIN_TOKEN" "$EXCH_OP_OID" "$ROLE_EXCHANGE_ADMIN")
  if [[ "$IS_PRESENT" == "true" ]]; then
    echo "  Exchange Administrator: PRESENT"
    STATUS_MAP["Exchange Operator:Exchange Administrator"]="OK"
  else
    echo "  Exchange Administrator: MISSING -> ASSIGNING..."
    if assign_role "$TENANT_ADMIN_TOKEN" "$EXCH_OP_OID" "$ROLE_EXCHANGE_ADMIN" "Exchange Administrator"; then
      STATUS_MAP["Exchange Operator:Exchange Administrator"]=$( [[ "$DRY_RUN" == "true" ]] && echo "DRY-RUN" || echo "ASSIGNED" )
    else
      STATUS_MAP["Exchange Operator:Exchange Administrator"]="ERROR"
      PARTIAL_FAILURE=true
    fi
  fi
fi

# Tenant Admin -> Conditional Access Administrator
# Required because Microsoft Graph evaluates SP membership in CA-relevant directory
# roles in addition to OAuth scopes when calling /identity/conditionalAccess/* endpoints.
# Without this role, the Tenant Admin SP gets 403 even with Policy.ReadWrite.ConditionalAccess
# in its token. Discovered Cascades 2026-04-28; backfilled here so future tenants don't trip on it.
echo ""
echo "[CHECK] Tenant Admin SP: $TA_SP_OID"
# Activate CA Admin role in tenant first (idempotent — returns 400 if already activated)
ACTIVATE_RESP=$(curl -s --max-time 15 \
  -H "Authorization: Bearer $TENANT_ADMIN_TOKEN" \
  -H "Content-Type: application/json" \
  -X POST \
  "https://graph.microsoft.com/v1.0/directoryRoles" \
  -d "{\"roleTemplateId\":\"$ROLE_CA_ADMIN\"}" 2>/dev/null || true)
# Don't fail on activation — directly-assigned role POST handles the case where already active

IS_CA=$(role_assigned "$TENANT_ADMIN_TOKEN" "$TA_SP_OID" "$ROLE_CA_ADMIN")
if [[ "$IS_CA" == "true" ]]; then
  echo "  Conditional Access Administrator: PRESENT"
  STATUS_MAP["Tenant Admin:Conditional Access Administrator"]="OK"
else
  echo "  Conditional Access Administrator: MISSING -> ASSIGNING..."
  if assign_role "$TENANT_ADMIN_TOKEN" "$TA_SP_OID" "$ROLE_CA_ADMIN" "Conditional Access Administrator"; then
    STATUS_MAP["Tenant Admin:Conditional Access Administrator"]=$( [[ "$DRY_RUN" == "true" ]] && echo "DRY-RUN" || echo "ASSIGNED" )
  else
    STATUS_MAP["Tenant Admin:Conditional Access Administrator"]="ERROR"
    PARTIAL_FAILURE=true
  fi
fi

# User Manager -> User Administrator + Authentication Administrator
if [[ -z "$USER_MGR_OID" ]]; then
  echo "[WARNING] User Manager SP still not found after consent attempt"
  STATUS_MAP["User Manager:User Administrator"]="MISSING SP"
  STATUS_MAP["User Manager:Authentication Administrator"]="MISSING SP"
else
  echo ""
  echo "[CHECK] User Manager SP: $USER_MGR_OID"

  IS_UA=$(role_assigned "$TENANT_ADMIN_TOKEN" "$USER_MGR_OID" "$ROLE_USER_ADMIN")
  if [[ "$IS_UA" == "true" ]]; then
    echo "  User Administrator: PRESENT"
    STATUS_MAP["User Manager:User Administrator"]="OK"
  else
    echo "  User Administrator: MISSING -> ASSIGNING..."
    if assign_role "$TENANT_ADMIN_TOKEN" "$USER_MGR_OID" "$ROLE_USER_ADMIN" "User Administrator"; then
      STATUS_MAP["User Manager:User Administrator"]=$( [[ "$DRY_RUN" == "true" ]] && echo "DRY-RUN" || echo "ASSIGNED" )
    else
      STATUS_MAP["User Manager:User Administrator"]="ERROR"
      PARTIAL_FAILURE=true
    fi
  fi

  IS_AA=$(role_assigned "$TENANT_ADMIN_TOKEN" "$USER_MGR_OID" "$ROLE_AUTH_ADMIN")
  if [[ "$IS_AA" == "true" ]]; then
    echo "  Authentication Administrator: PRESENT"
    STATUS_MAP["User Manager:Authentication Administrator"]="OK"
  else
    echo "  Authentication Administrator: MISSING -> ASSIGNING..."
    if assign_role "$TENANT_ADMIN_TOKEN" "$USER_MGR_OID" "$ROLE_AUTH_ADMIN" "Authentication Administrator"; then
      STATUS_MAP["User Manager:Authentication Administrator"]=$( [[ "$DRY_RUN" == "true" ]] && echo "DRY-RUN" || echo "ASSIGNED" )
    else
      STATUS_MAP["User Manager:Authentication Administrator"]="ERROR"
      PARTIAL_FAILURE=true
    fi
  fi
fi

# ── Step 6: Final status table ────────────────────────────────────────────────
echo ""
if [[ "$PARTIAL_FAILURE" == "true" ]] || [[ "$CONSENT_PARTIAL" == "true" ]]; then
  echo "[WARNING] Onboarding completed with errors for $DISPLAY_NAME"
else
  if [[ "$DRY_RUN" == "true" ]]; then
    echo "[INFO] Dry-run complete for $DISPLAY_NAME ($TENANT_ID) — no changes made"
  else
    echo "[SUCCESS] Onboarding complete for $DISPLAY_NAME"
  fi
fi

echo "SP roles status:"
TA_CA="${STATUS_MAP["Tenant Admin:Conditional Access Administrator"]:-SKIPPED}"
echo "  Tenant Admin:"
printf "    Conditional Access Administrator: %s\n" "[$TA_CA]"

SEC_EXCH="${STATUS_MAP["Security Investigator:Exchange Administrator"]:-SKIPPED}"
echo "  Security Investigator:"
printf "    Exchange Administrator:       %s\n" "[$SEC_EXCH]"

EO_EXCH="${STATUS_MAP["Exchange Operator:Exchange Administrator"]:-SKIPPED}"
echo "  Exchange Operator:"
printf "    Exchange Administrator:       %s\n" "[$EO_EXCH]"

UA="${STATUS_MAP["User Manager:User Administrator"]:-SKIPPED}"
AA="${STATUS_MAP["User Manager:Authentication Administrator"]:-SKIPPED}"
echo "  User Manager:"
printf "    User Administrator:           %s\n" "[$UA]"
printf "    Authentication Administrator: %s\n" "[$AA]"

if [[ "$PARTIAL_FAILURE" == "true" ]]; then
  exit 10
fi
exit 0
