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

# ES endpoint and wait strategy
ES="${ES:-http://localhost:9200}"
es_wait_attempts="${ES_WAIT_ATTEMPTS:-60}"      # total attempts to wait for ES
es_wait_interval="${ES_WAIT_INTERVAL:-2}"        # seconds between attempts

echo "[i] 查询 ES 端点：$ES"

wait_for_es() {
  local attempt=1
  while (( attempt <= es_wait_attempts )); do
    # 等待集群达到至少 yellow 状态；请求失败则重试
    if curl -fsS "$ES/_cluster/health?wait_for_status=yellow&timeout=1s" >/dev/null 2>&1; then
      echo "[ok] Elasticsearch 已就绪 (attempt=${attempt}/${es_wait_attempts})"
      return 0
    fi
    echo "[..] 等待 Elasticsearch 可用中 (${attempt}/${es_wait_attempts})"
    sleep "${es_wait_interval}"
    (( attempt++ ))
  done
  echo "[err] Elasticsearch 在 ${es_wait_attempts} 次尝试后仍不可用"
  return 1
}

safe_count() {
  # 对缺失索引返回 0，避免 404 触发失败
  local pattern="$1"
  local json
  json=$(curl -fsS "$ES/${pattern}/_count?ignore_unavailable=true&allow_no_indices=true" 2>/dev/null || echo '{}')
  echo "$json" | sed -E 's/.*"count":([0-9]+).*/\1/' | awk 'NF{print $0;exit} END{if(NR==0)print 0}'
}

wait_for_es

# 列出相关索引（可能为空，允许）
curl -fsS "$ES/_cat/indices?v" | egrep 'train-|infer-|logstash' || true

# 打印计数，缺失索引按 0 处理
printf "train-* 计数："; safe_count "train-*"; echo
printf "infer-* 计数："; safe_count "infer-*"; echo
