#!/usr/bin/env bash

set -euo pipefail

NOMAD_SECRETS_DIR="${NOMAD_SECRETS_DIR:-}"
NOMAD_TOKEN="${NOMAD_TOKEN:-}"
NOMAD_ADDR="${NOMAD_ADDR:-localhost}"

TARGET_NAMESPACE="${TARGET_NAMESPACE:-default}"
TARGET_JOB="${TARGET_JOB:-}"
TARGET_GROUP="${TARGET_GROUP:-}"

TARGET_COUNT="${TARGET_COUNT:-0}"

CURL_ARGS=(
  --silent
  --show-error
  --fail-with-body
)

if [[ -n "${NOMAD_TOKEN}" ]]; then
  CURL_ARGS+=(-H "Authorization: Bearer ${NOMAD_TOKEN}")
fi
if [[ -n "${NOMAD_SECRETS_DIR}" ]]; then
  CURL_ARGS+=(--unix-socket "${NOMAD_SECRETS_DIR}/api.sock")
fi

if ! [[ "$TARGET_COUNT" =~ ^[+-]?[0-9]+$ ]]; then
  echo "$TARGET_COUNT is not a number"
  exit 1
fi

curl "${CURL_ARGS[@]}" \
  -X POST \
  "${NOMAD_ADDR}/v1/job/${TARGET_JOB}/scale?namespace=${TARGET_NAMESPACE}" \
  -H "Content-Type: application/json" \
  -d "$(jq -n \
    --arg group "$TARGET_GROUP" \
    --arg count "$TARGET_COUNT" \
    '{
      Count: ($count|tonumber),
      Target: {
        "Group": ($group)
      },
      Message: "Scaling up to perform backup"
    }'
  )"

while true; do
  allocations=$(
    curl "${CURL_ARGS[@]}" \
      "${NOMAD_ADDR}/v1/job/${TARGET_JOB}/allocations?namespace=${TARGET_NAMESPACE}"
  )

  count=$(
    jq --arg group "$TARGET_GROUP" '
      [
        .[]
        | select(.TaskGroup == $group)
        | select(
            .ClientStatus != "complete"
            and .ClientStatus != "dead"
          )
      ] | length
    ' <<< "$allocations"
  )

  echo "Remaining allocations: ${count}"

  if [[ "$count" -eq $TARGET_COUNT ]]; then
    break
  fi

  sleep 5
done