#!/usr/bin/env bash
# Usage: ./sf-org-lic -to NUSHUB-PROD
# 
# Salesforce org license utilization report with proper error handling:
# - Validates org exists before running queries  
# - Reports clear error messages instead of silent failures
# - Uses strict error checking to catch more issues
# - Fixed SOQL queries to only use existing fields
# - Focuses on User Licenses and Permission Set Licenses (the main license types)
set -euo pipefail

ORG=""
while [[ $# -gt 0 ]]; do
  case "$1" in
    -to|--target-org) ORG="$2"; shift 2;;
    -hp|--help) 
      echo "Usage: $0 -to <ORG_ALIAS>"
      echo ""
      echo "Generate Salesforce license utilization report for an org"
      echo ""
      echo "Options:"
      echo "  -to  Target org alias or username (required)"
      echo "  -hp  Show this help message"
      echo ""
      echo "Examples:"
      echo "  $0 -to PROD-ORG"
      echo "  $0 -to dev@company.com"
      echo ""
      echo "Reports:"
      echo "  • User Licenses - Core Salesforce user licenses"
      echo "  • Permission Set Licenses - Add-on feature licenses"
      exit 0;;
    *) echo "Unknown arg: $1" >&2; exit 1;;
  esac
done
[[ -z "$ORG" ]] && { echo "Usage: $0 -to <ORG_ALIAS>" >&2; exit 1; }

command -v sf >/dev/null 2>&1 || { echo "'sf' CLI is required." >&2; exit 1; }
command -v jq >/dev/null 2>&1 || { echo "'jq' is required." >&2; exit 1; }
command -v column >/dev/null 2>&1 || { echo "'column' is required." >&2; exit 1; }

BOLD="$(printf '\033[1m')"; DIM="$(printf '\033[2m')"
CYAN="$(printf '\033[36m')"; GREEN="$(printf '\033[32m')"
RESET="$(printf '\033[0m')"

underline() { printf '%*s\n' "${#1}" '' | tr ' ' '-'; }

# Helper function for error messages
error_exit() { echo "Error: $1" >&2; exit "${2:-1}"; }

# Validate that the org exists and is authorized
validate_org() {
  local org="$1"
  echo "Validating org: $org..." >&2
  
  if ! sf org display --target-org "$org" --json >/dev/null 2>&1; then
    # Try to get available orgs to suggest alternatives
    local available_orgs
    available_orgs=$(sf org list --json 2>/dev/null | jq -r '.result.other[]?.alias // empty' 2>/dev/null | tr '\n' ', ' | sed 's/,$//' || echo "none")
    if [[ -z "$available_orgs" ]]; then
      available_orgs="none"
    fi
    error_exit "org alias '$org' not found or not authorized. Available orgs: $available_orgs" 2
  fi
}

# Normalize records location across sf CLI variants
JQ_RECS_DEF='def recs: (.result.records? // .records? // .Result.records? // []);'

run_json_with_error_handling() {
  local soql="$1"
  local json
  local exit_code
  
  # Temporarily disable strict error checking for this function
  set +e
  json=$(sf data query --target-org "$ORG" --json --query "$soql" 2>&1)
  exit_code=$?
  set -e
  
  # If the sf command itself failed, return the special error code
  if [[ $exit_code -ne 0 ]]; then
    return 42
  fi
  
  # Check if the result contains error information (sf CLI error responses)
  if echo "$json" | jq -e '.name // .error // (.message and (.commandName // .status))' >/dev/null 2>&1; then
    local error_msg
    error_msg=$(echo "$json" | jq -r '.message // .error // "Unknown error"' 2>/dev/null || echo "JSON parse error")
    
    # Return special error code for object not found/no access (non-fatal)
    if [[ "$error_msg" == *"not supported"* ]] || [[ "$error_msg" == *"does not exist"* ]] || [[ "$error_msg" == *"INVALID_TYPE"* ]] || [[ "$error_msg" == *"No such column"* ]]; then
      return 42  # Special return code for missing objects
    fi
    
    error_exit "SOQL query failed: $error_msg" 4
  fi
  
  echo "$json"
  return 0
}

print_table_or_empty() {
  # $1=json, $2=jq expr (builds a TSV row array from each record), $3=headers CSV
  local json="$1" jexpr="$2" headers="$3"
  local count
  count="$(echo "$json" | jq -r "$JQ_RECS_DEF recs | length" 2>/dev/null || echo 0)"
  echo "$headers" | awk '{gsub(/,/, "\t"); print}' | column -t -s $'\t'
  if [[ "$count" -eq 0 ]]; then
    echo "(no rows)"
    echo
    return
  fi
  echo "$json" \
    | jq -r "$JQ_RECS_DEF recs[] | $jexpr | @tsv" \
    | column -t -s $'\t'
  echo
}

print_totals() {
  local json="$1"
  local total used remaining
  total=$(echo "$json" | jq -r "$JQ_RECS_DEF [recs[] | (try .TotalLicenses|tonumber // 0)] | add // 0")
  used=$( echo "$json" | jq -r "$JQ_RECS_DEF [recs[] | (try .UsedLicenses|tonumber // 0)] | add // 0")
  remaining=$(( total - used ))
  echo -e "${CYAN}Totals:${RESET} Total=${BOLD}${total}${RESET}  Used=${BOLD}${used}${RESET}  Remaining=${GREEN}${BOLD}${remaining}${RESET}"
  echo
}

header() {
  echo
  echo -e "${BOLD}Salesforce License Utilization — ${ORG}${RESET}"
  echo -e "${DIM}Generated: $(TZ=Asia/Manila date '+%Y-%m-%d %H:%M:%S %Z')${RESET}"
  echo
}
section() { echo -e "${BOLD}$1${RESET}"; underline "$1"; }

# SOQL queries for the two main license types that exist in most orgs
SOQL_USER='SELECT Id, Name, TotalLicenses, UsedLicenses FROM UserLicense ORDER BY Name LIMIT 200'
SOQL_PSL='SELECT Id, MasterLabel, DeveloperName, TotalLicenses, UsedLicenses FROM PermissionSetLicense ORDER BY MasterLabel LIMIT 200'

# Validate org before proceeding
validate_org "$ORG"

header

# 1) User Licenses - Core Salesforce licenses
section "User Licenses"
if JUSER="$(run_json_with_error_handling "$SOQL_USER")"; then
  if ! echo "$JUSER" | jq -e . >/dev/null 2>&1; then
    echo "Failed to query UserLicense (no valid JSON)."
    echo
  else
    print_table_or_empty "$JUSER" \
      '[ .Name,
         (try .TotalLicenses|tonumber // 0),
         (try .UsedLicenses|tonumber // 0),
         ((try .TotalLicenses|tonumber // 0) - (try .UsedLicenses|tonumber // 0))
       ]' \
      "Name,Total,Used,Remaining"
    print_totals "$JUSER"
  fi
elif [[ $? -eq 42 ]]; then
  echo "UserLicense object not available in this org."
  echo
fi

# 2) Permission Set Licenses - Add-on feature licenses  
section "Permission Set Licenses"
if JPSL="$(run_json_with_error_handling "$SOQL_PSL")"; then
  if ! echo "$JPSL" | jq -e . >/dev/null 2>&1; then
    echo "Failed to query PermissionSetLicense (no valid JSON)."
    echo
  else
    print_table_or_empty "$JPSL" \
      '[ .MasterLabel,
         .DeveloperName,
         (try .TotalLicenses|tonumber // 0),
         (try .UsedLicenses|tonumber // 0),
         ((try .TotalLicenses|tonumber // 0) - (try .UsedLicenses|tonumber // 0))
       ]' \
      "MasterLabel,DeveloperName,Total,Used,Remaining"
    print_totals "$JPSL"
  fi
elif [[ $? -eq 42 ]]; then
  echo "PermissionSetLicense object not available in this org."
  echo
fi

echo -e "${DIM}Note: This report covers the main Salesforce license types available in most orgs.${RESET}"
echo -e "${DIM}FeatureLicense objects are not commonly available and have been excluded.${RESET}"