#!/usr/bin/env bash
# setup-endpoint-service.sh
#
# Creates a VPC Endpoint Service from a Kong NLB and registers a private DNS
# name with Route53 domain ownership verification.
#
# Usage:
#   ./setup-endpoint-service.sh \
#     --nlb-name      o11y-dev-kong-privatelink \
#     --private-dns   www.ds.observability-dev.pingcap.cloud \
#     --hosted-zone   Z0833482527F3J3BE8AO \
#     --region        us-west-2 \
#     --profile       clinic-dev
#
# All flags are required. The script is idempotent: re-running it when the
# endpoint service already exists only updates the private DNS if it changed,
# and skips verification if the domain is already verified.

set -euo pipefail

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------
NLB_NAME=""
PRIVATE_DNS=""
HOSTED_ZONE=""
REGION=""
AWS_PROFILE=""

usage() {
  grep '^#' "$0" | grep -v '#!/' | sed 's/^# \{0,1\}//'
  exit 1
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --nlb-name)    NLB_NAME="$2";    shift 2 ;;
    --private-dns) PRIVATE_DNS="$2"; shift 2 ;;
    --hosted-zone) HOSTED_ZONE="$2"; shift 2 ;;
    --region)      REGION="$2";      shift 2 ;;
    --profile)     AWS_PROFILE="$2"; shift 2 ;;
    -h|--help)     usage ;;
    *) echo "Unknown argument: $1"; usage ;;
  esac
done

for flag in NLB_NAME PRIVATE_DNS HOSTED_ZONE REGION AWS_PROFILE; do
  [[ -n "${!flag}" ]] || { echo "ERROR: --$(echo "$flag" | tr '[:upper:]_' '[:lower:]-') is required"; exit 1; }
done

AWS="aws --region $REGION --profile $AWS_PROFILE"

log()  { echo "[$(date -u '+%H:%M:%S')] $*"; }
ok()   { echo "[$(date -u '+%H:%M:%S')] ✓ $*"; }
err()  { echo "[$(date -u '+%H:%M:%S')] ✗ $*" >&2; exit 1; }

# ---------------------------------------------------------------------------
# 1. Resolve NLB ARN
# ---------------------------------------------------------------------------
log "Looking up NLB: $NLB_NAME"
NLB_ARN=$($AWS elbv2 describe-load-balancers \
  --names "$NLB_NAME" \
  --query "LoadBalancers[0].LoadBalancerArn" \
  --output text 2>/dev/null) || err "NLB '$NLB_NAME' not found in $REGION"

[[ "$NLB_ARN" == "None" ]] && err "NLB '$NLB_NAME' not found in $REGION"
ok "NLB ARN: $NLB_ARN"

# ---------------------------------------------------------------------------
# 2. Create or locate the endpoint service
# ---------------------------------------------------------------------------
log "Checking for existing endpoint service backed by this NLB..."
# No server-side filter for NLB ARN exists; match client-side via JMESPath.
SERVICE_ID=$($AWS ec2 describe-vpc-endpoint-service-configurations \
  --query "ServiceConfigurations[?contains(NetworkLoadBalancerArns, '$NLB_ARN')].ServiceId | [0]" \
  --output text 2>/dev/null)

if [[ -z "$SERVICE_ID" || "$SERVICE_ID" == "None" ]]; then
  log "Creating endpoint service from NLB..."
  SERVICE_ID=$($AWS ec2 create-vpc-endpoint-service-configuration \
    --network-load-balancer-arns "$NLB_ARN" \
    --acceptance-required \
    --query "ServiceConfiguration.ServiceId" \
    --output text)
  ok "Created endpoint service: $SERVICE_ID"
else
  ok "Endpoint service already exists: $SERVICE_ID"
fi

# ---------------------------------------------------------------------------
# 3. Set the private DNS name (idempotent)
# ---------------------------------------------------------------------------
CURRENT_DNS=$($AWS ec2 describe-vpc-endpoint-service-configurations \
  --service-ids "$SERVICE_ID" \
  --query "ServiceConfigurations[0].PrivateDnsName" \
  --output text)

if [[ "$CURRENT_DNS" != "$PRIVATE_DNS" ]]; then
  log "Setting private DNS name: $PRIVATE_DNS"
  $AWS ec2 modify-vpc-endpoint-service-configuration \
    --service-id "$SERVICE_ID" \
    --private-dns-name "$PRIVATE_DNS" > /dev/null
  ok "Private DNS name set"
else
  ok "Private DNS name already set: $PRIVATE_DNS"
fi

# ---------------------------------------------------------------------------
# 4. Check verification state
# ---------------------------------------------------------------------------
get_dns_config() {
  $AWS ec2 describe-vpc-endpoint-service-configurations \
    --service-ids "$SERVICE_ID" \
    --query "ServiceConfigurations[0].PrivateDnsNameConfiguration" \
    --output json
}

DNS_CONFIG=$(get_dns_config)
VERIFY_STATE=$(echo "$DNS_CONFIG" | python3 -c "import sys,json; print(json.load(sys.stdin).get('State',''))")

if [[ "$VERIFY_STATE" == "verified" ]]; then
  ok "Domain ownership already verified"
else
  # -------------------------------------------------------------------------
  # 5. Add Route53 TXT record for domain ownership verification
  # -------------------------------------------------------------------------
  TXT_NAME=$(echo "$DNS_CONFIG"  | python3 -c "import sys,json; print(json.load(sys.stdin)['Name'])")
  TXT_VALUE=$(echo "$DNS_CONFIG" | python3 -c "import sys,json; print(json.load(sys.stdin)['Value'])")
  FQDN="${TXT_NAME}.$(echo "$PRIVATE_DNS" | cut -d. -f2-)"

  log "Adding Route53 TXT record: $FQDN = \"$TXT_VALUE\""

  # Check if TXT record already exists with the right value
  EXISTING=$($AWS route53 list-resource-record-sets \
    --hosted-zone-id "$HOSTED_ZONE" \
    --query "ResourceRecordSets[?Name=='${FQDN}.'].ResourceRecords[0].Value" \
    --output text 2>/dev/null || true)

  EXPECTED="\"$TXT_VALUE\""
  if [[ "$EXISTING" == "$EXPECTED" ]]; then
    ok "TXT record already present"
  else
    $AWS route53 change-resource-record-sets \
      --hosted-zone-id "$HOSTED_ZONE" \
      --change-batch "{
        \"Changes\": [{
          \"Action\": \"UPSERT\",
          \"ResourceRecordSet\": {
            \"Name\": \"$FQDN\",
            \"Type\": \"TXT\",
            \"TTL\": 300,
            \"ResourceRecords\": [{\"Value\": \"\\\"$TXT_VALUE\\\"\"}]
          }
        }]
      }" > /dev/null
    ok "TXT record created: $FQDN"
  fi

  # -------------------------------------------------------------------------
  # 6. Trigger verification
  # -------------------------------------------------------------------------
  log "Triggering domain ownership verification..."
  $AWS ec2 start-vpc-endpoint-service-private-dns-verification \
    --service-id "$SERVICE_ID" > /dev/null
  ok "Verification triggered"

  # -------------------------------------------------------------------------
  # 7. Poll until verified (up to 5 minutes)
  # -------------------------------------------------------------------------
  log "Waiting for verification (this may take up to 2 minutes)..."
  for i in $(seq 1 30); do
    sleep 10
    VERIFY_STATE=$(get_dns_config | python3 -c "import sys,json; print(json.load(sys.stdin).get('State',''))")
    log "  State: $VERIFY_STATE (${i}/30)"
    [[ "$VERIFY_STATE" == "verified" ]] && break
    [[ "$VERIFY_STATE" == "failed" ]] && err "Verification failed. Check that the TXT record is correct and DNS has propagated."
  done

  [[ "$VERIFY_STATE" == "verified" ]] || err "Timed out waiting for verification (state: $VERIFY_STATE)"
  ok "Domain ownership verified"
fi

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------
echo ""
echo "================================================"
echo "  Endpoint Service : $SERVICE_ID"
echo "  Private DNS      : $PRIVATE_DNS"
echo "  Verification     : verified"
echo "================================================"
echo ""
echo "Add this to the Kong proxy service annotations in values.yaml:"
echo "  # VPC Endpoint Service: $SERVICE_ID"
echo "  # Private DNS: $PRIVATE_DNS"