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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
INVENTORY_DEFAULT="${ROOT_DIR}/ansible/inventory.ini"
K3S_NODES_GROUP_SERVER="k3s_server"
K3S_NODES_GROUP_WORKER="k3s_worker"
TIMEOUT_SEC="${TIMEOUT_SEC:-5}"

usage() {
  cat <<'EOF'
用法:
  test-ssh.sh [选项]

选项:
  --inventory <path>          ansible inventory 路径（默认 ansible/inventory.ini）
  --timeout <sec>             默认 5
  -h, --help                  显示帮助
EOF
}

parse_k3s_hosts() {
  local inventory="$1"
  local group="$2"

  if [[ ! -f "$inventory" ]]; then
    echo "[ERR] 找不到 inventory: $inventory" >&2
    return 1
  fi

  # 兼容 CRLF：先去掉 \r 再解析
  tr -d '\r' < "$inventory" | awk -v g="[$group]" '
    $0 == g { in_group=1; next }
    in_group && ($0 ~ /^\[/) { in_group=0 }
    in_group && $0 !~ /^($|#)/ {
      host=""; ip=""
      host=$1
      for (i=1; i<=NF; i++) {
        if ($i ~ /^ansible_host=/) {
          split($i, a, "=")
          ip=a[2]
        }
      }
      if (ip != "") { print ip }
      else if (host != "") { print host }
    }
  '
}

INVENTORY_PATH="$INVENTORY_DEFAULT"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --inventory) INVENTORY_PATH="${2:-}"; shift 2 ;;
    --timeout) TIMEOUT_SEC="${2:-5}"; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "[ERR] 未知参数: $1"; usage; exit 1 ;;
  esac
done

check_key() {
  local path="$1"
  if [[ ! -f "$path" ]]; then
    echo "[ERR] 私钥不存在: $path"
    echo "      请先运行 scripts/ssh/setup-k3s-workers-ssh.sh 或手动生成密钥，并在远端写入对应公钥。"
    return 1
  fi
   # 自动矫正权限，避免 "UNPROTECTED PRIVATE KEY FILE" 报错
  chmod 600 "$path" 2>/dev/null || true
  return 0
}

run_test() {
  local name="$1"
  local host="$2"
  local key="$3"
  local failed=0

  echo
  echo "=== 测试: ${name} (${host}) ==="

  if ! check_key "$key"; then
    return 1
  fi

  echo "[INFO] 使用私钥: $key"
  local ssh_base=(ssh -o BatchMode=yes -o StrictHostKeyChecking=accept-new -o ConnectTimeout="${TIMEOUT_SEC}" -i "$key" -o IdentitiesOnly=yes "$host")

  if "${ssh_base[@]}" "echo ok" >/dev/null 2>&1; then
    echo "[OK] SSH key 登录成功"
  else
    echo "[FAIL] SSH key 登录失败"
    failed=1
  fi

  if "${ssh_base[@]}" "hostname; id -u; whoami" >/dev/null 2>&1; then
    echo "[OK] 远端基础命令执行成功"
  else
    echo "[FAIL] 远端基础命令执行失败"
    failed=1
  fi

  if "${ssh_base[@]}" "sudo -n true" >/dev/null 2>&1; then
    echo "[OK] 远端 sudo 免密可用"
  else
    echo "[WARN] 远端 sudo 需要密码（不影响 SSH，但会影响自动排障）"
  fi

  if "${ssh_base[@]}" "sudo -n ip -br a >/dev/null 2>&1 || ip -br a >/dev/null 2>&1"; then
    echo "[OK] 网络命令可执行"
  else
    echo "[FAIL] 网络命令不可执行"
    failed=1
  fi

  return $failed
}

echo "开始 SSH 连通性测试..."

if [[ ! -f "$INVENTORY_PATH" ]]; then
  echo "[ERR] 找不到 inventory: $INVENTORY_PATH" >&2
  exit 1
fi

K3S_HOSTS=()
for grp in "$K3S_NODES_GROUP_SERVER" "$K3S_NODES_GROUP_WORKER"; do
  while IFS= read -r h; do
    [[ -n "$h" ]] && K3S_HOSTS+=("$h")
  done < <(parse_k3s_hosts "$INVENTORY_PATH" "$grp" 2>/dev/null || true)
done

if [[ "${#K3S_HOSTS[@]}" -eq 0 ]]; then
  echo "[ERR] 未在 ${INVENTORY_PATH} 中找到任何 k3s_server / k3s_worker 主机" >&2
  exit 1
fi

echo "[INFO] 将测试以下节点的 root + 证书登录（按 inventory 中的 k3s_server / k3s_worker 分组）："
for h in "${K3S_HOSTS[@]}"; do
  echo "  - root@${h}"
done

rc=0
for h in "${K3S_HOSTS[@]}"; do
  key_path="$HOME/.ssh/id_ed25519_k3s_${h}"
  run_test "$h" "root@${h}" "$key_path" || rc=1
done

echo
if [[ $rc -eq 0 ]]; then
  echo "[PASS] SSH 测试通过，可用于 Ansible/运维自动化。"
else
  echo "[FAIL] 存在失败项，请先修复 SSH/key/sudo 配置。"
fi
exit $rc