#!/usr/bin/env bash
set -euo pipefail

: "${AWS_PROFILE:?AWS_PROFILE is required}"
: "${AWS_REGION:?AWS_REGION is required}"
: "${ENDPOINT_SERVICE_ID:?ENDPOINT_SERVICE_ID is required}"
: "${DESIRED_AVAILABILITY_ZONE_NAMES:?DESIRED_AVAILABILITY_ZONE_NAMES is required}"

read -r -a desired_zones <<<"${DESIRED_AVAILABILITY_ZONE_NAMES}"
read -r -a current_zones <<<"$(
  aws ec2 describe-vpc-endpoint-service-configurations \
    --profile "${AWS_PROFILE}" \
    --region "${AWS_REGION}" \
    --service-ids "${ENDPOINT_SERVICE_ID}" \
    --query 'ServiceConfigurations[0].AvailabilityZones' \
    --output text
)"

declare -A current=()
for zone in "${current_zones[@]}"; do
  current["${zone}"]=1
done

missing=()
for zone in "${desired_zones[@]}"; do
  if [[ -z "${current[${zone}]:-}" ]]; then
    missing+=("${zone}")
  fi
done

if ((${#missing[@]} > 0)); then
  # AWS derives supported zones from the associated load balancers. Calling
  # ModifyVpcEndpointServiceConfiguration with only the service ID refreshes
  # zones that were enabled on an NLB after the service was created.
  aws ec2 modify-vpc-endpoint-service-configuration \
    --profile "${AWS_PROFILE}" \
    --region "${AWS_REGION}" \
    --service-id "${ENDPOINT_SERVICE_ID}" \
    >/dev/null
fi

for _ in {1..12}; do
  read -r -a current_zones <<<"$(
    aws ec2 describe-vpc-endpoint-service-configurations \
      --profile "${AWS_PROFILE}" \
      --region "${AWS_REGION}" \
      --service-ids "${ENDPOINT_SERVICE_ID}" \
      --query 'ServiceConfigurations[0].AvailabilityZones' \
      --output text
  )"
  current=()
  for zone in "${current_zones[@]}"; do
    current["${zone}"]=1
  done
  complete=true
  for zone in "${desired_zones[@]}"; do
    if [[ -z "${current[${zone}]:-}" ]]; then
      complete=false
      break
    fi
  done
  if [[ "${complete}" == true ]]; then
    printf 'Endpoint Service %s supports all requested AZs: %s\n' \
      "${ENDPOINT_SERVICE_ID}" "${DESIRED_AVAILABILITY_ZONE_NAMES}"
    exit 0
  fi
  sleep 5
done

printf 'Endpoint Service %s did not expose all requested AZs within 60 seconds.\n' \
  "${ENDPOINT_SERVICE_ID}" >&2
exit 1