#!/usr/bin/env bash
# run-prometheus-test.sh - run the Prometheus binary for each cluster and collect metrics to files
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/test-common.sh"
BIN_DIR="${GENERATED_BIN_DIR}"
CONFIG_DIR="${GENERATED_CONFIG_DIR}"

# Wait for the Prometheus HTTP server to become ready (up to 30s)
wait_ready() {
  local port="$1"
  local i
  for i in $(seq 1 30); do
    if curl -sf "http://127.0.0.1:${port}/healthz" > /dev/null 2>&1; then
      return 0
    fi
    sleep 1
  done
  echo "  ❌ service did not become ready within 30s (port=${port})"
  return 1
}

check_representative_metrics() {
  local service_plan="$1"
  local metrics_file="$2"

  case "${service_plan}" in
    Dedicated)
      if grep -q '^tidbcloud_node_cpu_capacity_cores' "${metrics_file}" \
        || grep -q '^tidbcloud_component_uptime' "${metrics_file}" \
        || grep -q '^tidbcloud_db_connections_total' "${metrics_file}"; then
        echo "  ✅ found Dedicated representative metrics"
      else
        echo "  ⚠️  Dedicated payload does not contain representative dedicated-only metrics"
      fi
      ;;
    Premium|Essential|Starter)
      local missing=0
      if ! grep -q '^tidbcloud_db_total_connection' "${metrics_file}"; then
        echo "  ⚠️  missing tidbcloud_db_total_connection"
        missing=1
      fi
      if ! grep -q '^tidbcloud_db_row_storage_used_bytes' "${metrics_file}"; then
        echo "  ⚠️  missing tidbcloud_db_row_storage_used_bytes"
        missing=1
      fi
      if [[ "${missing}" == "0" ]]; then
        echo "  ✅ found NextGen representative DB/storage metrics"
      fi
      ;;
  esac
}

# run_prometheus_test <binary> <cluster_id> <service_plan> <branch> <listen_port>
run_prometheus_test() {
  local binary="$1"
  local cluster_id="$2"
  local service_plan="$3"
  local branch="$4"
  local port="$5"

  local output_dir="${RESULTS_DIR}/${cluster_id}-${service_plan}/${branch}/prometheus"
  mkdir -p "${output_dir}"

  # Generate a temporary config and replace all listening ports to avoid collisions
  local tmp_config
  tmp_config=$(mktemp --suffix=.yaml)
  sed -e "s/0\.0\.0\.0:8080/0.0.0.0:${port}/" \
      -e "s/grpcPort: 10001/grpcPort: $(( 10001 + port - 18080 ))/" \
      -e "s/metricsPort: 9091/metricsPort: $(( 9091 + port - 18080 ))/" \
      "${CONFIG_DIR}/prometheus-config-${service_plan}-${cluster_id}.yaml" > "${tmp_config}"

  local file_prefix="${service_plan}-${cluster_id}"

  echo "  [${branch}][${service_plan}] starting ${binary} (port=${port})..."
  "${BIN_DIR}/${binary}" --config-path "${tmp_config}" \
    > "${output_dir}/server-${file_prefix}.log" 2>&1 &
  local pid=$!

  # Ensure cleanup on script exit
  local cleanup_done=0
  cleanup() {
    if [[ "${cleanup_done}" == "0" ]]; then
      cleanup_done=1
      kill "${pid}" 2>/dev/null || true
      rm -f "${tmp_config}"
    fi
  }
  trap cleanup EXIT INT TERM

  if ! wait_ready "${port}"; then
    echo "  --- last 20 log lines ---"
    tail -20 "${output_dir}/server-${file_prefix}.log" || true
    cleanup
    trap - EXIT INT TERM
    return 1
  fi

  # Wait for the first Auth0 token fetch and Thanos collection cycle (about 20-40s)
  echo "  waiting 40s for the first collection cycle to finish..."
  sleep 40

  local metrics_file="${output_dir}/${file_prefix}.txt"
  echo "  collecting metrics..."
  if curl -sf \
       -H "Authorization: Bearer ${PROMETHEUS_TOKEN}" \
       "http://127.0.0.1:${port}/public_api/v1/clusters/${cluster_id}/metrics" \
       -o "${metrics_file}"; then
    local line_count
    line_count=$(wc -l < "${metrics_file}")
    echo "  ✅ saved to ${metrics_file#"${TOOL_ROOT}/"} (${line_count} lines)"
    check_representative_metrics "${service_plan}" "${metrics_file}"
  else
    echo "  ❌ collection failed (HTTP error or connection refused), see ${output_dir#"${TOOL_ROOT}/"}/server-${file_prefix}.log"
  fi

  cleanup
  trap - EXIT INT TERM
}

run_prometheus_test_if_ready() {
  local binary="$1"
  local cluster_id="$2"
  local service_plan="$3"
  local branch="$4"
  local port="$5"

  if [[ -z "${cluster_id}" ]]; then
    echo "  [${branch}][${service_plan}] skipping: cluster_id is empty, you can retry later"
    return 0
  fi

  run_prometheus_test "${binary}" "${cluster_id}" "${service_plan}" "${branch}" "${port}"
}

echo "================================================================"
echo " Prometheus Pipeline Test"
echo "================================================================"
echo ""
echo "Results round directory: ${RESULTS_DIR}"
echo ""

mkdir -p "${RESULTS_DIR}"

# run_prometheus_pair - start feature and baseline together for the same cluster to keep timing aligned
# Usage: run_prometheus_pair <cluster_id> <service_plan> <feature_port> <baseline_port>
run_prometheus_pair() {
  local cluster_id="$1"
  local service_plan="$2"
  local feature_port="$3"
  local baseline_port="$4"

  if [[ -z "${cluster_id}" ]]; then
    echo "  [${service_plan}] skipping: cluster_id is empty"
    return 0
  fi

  echo "--- ${service_plan} (starting feature:${feature_port} + ${BASELINE_LABEL}:${baseline_port} together) ---"
  run_prometheus_test "prometheus-feature" "${cluster_id}" "${service_plan}" "feature" "${feature_port}" &
  local pid_feature=$!
  run_prometheus_test "prometheus-${BASELINE_LABEL}" "${cluster_id}" "${service_plan}" "${BASELINE_LABEL}" "${baseline_port}" &
  local pid_baseline=$!

  wait "${pid_feature}" || true
  wait "${pid_baseline}" || true
  echo ""
}

# Collect from all configured clusters in parallel with unique ports
declare -a PIDS=()

run_prometheus_pair "${CLUSTER_ID_DEDICATED}" "Dedicated" 18080 18090 &
PIDS+=($!)

run_prometheus_test_if_ready "prometheus-feature" "${CLUSTER_ID_PREMIUM}"   "Premium"   "feature" 18081 &
PIDS+=($!)

run_prometheus_test_if_ready "prometheus-feature" "${CLUSTER_ID_ESSENTIAL}" "Essential" "feature" 18082 &
PIDS+=($!)

run_prometheus_test_if_ready "prometheus-feature" "${CLUSTER_ID_STARTER}"   "Starter"   "feature" 18083 &
PIDS+=($!)

ALL_OK=true
for pid in "${PIDS[@]}"; do
  wait "${pid}" || ALL_OK=false
done

echo ""
${ALL_OK} && echo "✅ Prometheus test completed" || { echo "⚠️  Some tests failed"; exit 1; }