#!/usr/bin/env bash

set -euo pipefail

preflight_only=false
case ${1:-} in
  --preflight-only)
    preflight_only=true
    ;;
  '')
    ;;
  *)
    echo "Usage: $0 [--preflight-only]" >&2
    exit 2
    ;;
esac

script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
deployment_dir=$(cd "$script_dir/.." && pwd)
cd "$deployment_dir"

inventory_path=${LABS_INVENTORY_PATH:-inventory/}
playbook_path=${LABS_PLAYBOOK_PATH:-playbooks/labs_service.yml}
expected_hosts=${LABS_EXPECTED_HOSTS:-2}
public_health_url=${LABS_PUBLIC_HEALTH_URL:-https://labs.tidb.io/health}
target_group_names=(
  labs-pingcap-com-manual-lb-tg
  labs-tidb-io-asg-lb-tg
)

curl_health() {
  curl \
    --fail \
    --silent \
    --show-error \
    --location \
    --retry 5 \
    --retry-delay 3 \
    --retry-all-errors \
    --max-time 20 \
    "$public_health_url"
  echo
}

echo "Running deployment preflight checks"
ansible-playbook -i "$inventory_path" "$playbook_path" --syntax-check
ansible -i "$inventory_path" lab_servers -m ansible.builtin.ping
curl_health

hosts=()
while IFS= read -r host; do
  hosts+=("$host")
done < <(
  ansible-inventory -i "$inventory_path" --list |
    jq -r '.lab_servers.hosts[]' |
    sort
)

if [[ ${#hosts[@]} -ne $expected_hosts ]]; then
  echo "Expected $expected_hosts Labs hosts, found ${#hosts[@]}" >&2
  exit 1
fi

target_group_arns=()
for target_group_name in "${target_group_names[@]}"; do
  target_group_arn=$(aws elbv2 describe-target-groups \
    --names "$target_group_name" \
    --query 'TargetGroups[0].TargetGroupArn' \
    --output text)

  if [[ -z $target_group_arn || $target_group_arn == None ]]; then
    echo "Target group not found: $target_group_name" >&2
    exit 1
  fi

  target_group_arns+=("$target_group_arn")
done

for host in "${hosts[@]}"; do
  instance_id=$(ansible-inventory \
    -i "$inventory_path" \
    --host "$host" |
    jq -r '.ec2_instance_id')

  if [[ ! $instance_id =~ ^i-[0-9a-f]+$ ]]; then
    echo "Invalid EC2 instance ID for $host: $instance_id" >&2
    exit 1
  fi

  echo "Checking $instance_id before deployment"
  for target_group_arn in "${target_group_arns[@]}"; do
    aws elbv2 wait target-in-service \
      --target-group-arn "$target_group_arn" \
      --targets "Id=$instance_id,Port=80"
  done

  if [[ $preflight_only == true ]]; then
    continue
  fi

  echo "Deploying $host ($instance_id)"
  ansible-playbook \
    -i "$inventory_path" \
    "$playbook_path" \
    --limit "$host"

  echo "Waiting for $instance_id to become healthy"
  for target_group_arn in "${target_group_arns[@]}"; do
    aws elbv2 wait target-in-service \
      --target-group-arn "$target_group_arn" \
      --targets "Id=$instance_id,Port=80"
  done

  curl_health
done

if [[ $preflight_only == true ]]; then
  echo "Deployment preflight completed successfully"
  exit 0
fi

echo "All Labs servers deployed successfully"
curl_health