#!/usr/bin/env bash
# 等待 TCP host:port 就绪的工具（wait-for-it 自包含实现）。
# 思路来自 https://github.com/vishnubob/wait-for-it（MIT），精简为项目内置脚本，
# 避免依赖已废弃的 compose/wait-for-it 镜像。
#
# 用法：
#   wait-for-it.sh host:port [-t timeout] [-s] [-q] [-- command...]
#   -t / --timeout   超时秒数（默认 15）
#   -s / --strict    探测失败时不执行后续命令、返回非零
#   -q / --quiet     不输出状态
#   -- command...    探测成功后执行的命令
set -o pipefail

cmdname=$(basename "$0")
echoerr() { [ "${QUIET:-0}" != 1 ] && echo "$@" 1>&2; }

usage() {
    cat <<EOF >&2
Usage: $cmdname host:port [-t timeout] [-s] [-q] [-- command...]
EOF
    exit 1
}

HOST=""
PORT=""
TIMEOUT=15
STRICT=0
QUIET=0
CLI=()

while [ $# -gt 0 ]; do
    case "$1" in
        *:*) HOST=${1%%:*}; PORT=${1##*:}; shift ;;
        -h) HOST="$2"; shift 2 ;;
        --host=*) HOST="${1#*=}"; shift ;;
        -p) PORT="$2"; shift 2 ;;
        --port=*) PORT="${1#*=}"; shift ;;
        -t) TIMEOUT="$2"; shift 2 ;;
        --timeout=*) TIMEOUT="${1#*=}"; shift ;;
        -s|--strict) STRICT=1; shift ;;
        -q|--quiet) QUIET=1; shift ;;
        --) shift; CLI=("$@"); break ;;
        *) usage ;;
    esac
done

if [ -z "$HOST" ] || [ -z "$PORT" ]; then
    usage
fi

# 探测方式：优先 nc，否则用 bash 内建 /dev/tcp（python:3.12-slim 自带 bash）。
USE_NC=0
if command -v nc >/dev/null 2>&1; then
    USE_NC=1
fi

probe() {
    if [ "$USE_NC" -eq 1 ]; then
        nc -z "$HOST" "$PORT"
    else
        (echo > "/dev/tcp/$HOST/$PORT")
    fi
}

echoerr "$cmdname: 等待 $HOST:$PORT（超时 ${TIMEOUT}s）"
start=$(date +%s)
while :; do
    if probe >/dev/null 2>&1; then
        echoerr "$cmdname: $HOST:$PORT 就绪（耗时 $(($(date +%s) - start))s）"
        if [ "${#CLI[@]}" -gt 0 ]; then
            exec "${CLI[@]}"
        fi
        exit 0
    fi
    if [ "$TIMEOUT" -ge 0 ] && [ "$(date +%s)" -gt $((start + TIMEOUT)) ]; then
        break
    fi
    sleep 1
done

echoerr "$cmdname: $HOST:$PORT 在 ${TIMEOUT}s 内未就绪"
if [ "${#CLI[@]}" -gt 0 ]; then
    if [ "$STRICT" -eq 1 ]; then
        exit 1
    fi
    exec "${CLI[@]}"
fi
exit 1