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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
INVENTORY_DEFAULT="${ROOT_DIR}/ansible/inventory.ini"
SSH_USER_DEFAULT="jack"
# 配置所有 k3s 节点（控制节点 + 工作节点），便于 Ansible 以 root SSH 连接
K3S_NODES_GROUP="k3s_nodes"

print_title() {
  echo
  echo "=== $1 ==="
}

ask_default() {
  local prompt="$1"
  local def="$2"
  local v
  printf "%s [%s]: " "$prompt" "$def" >&2
  read -r v
  echo "${v:-$def}"
}

ensure_cmd() {
  local c="$1"
  if ! command -v "$c" >/dev/null 2>&1; then
    echo "[ERR] 缺少命令: $c"
    exit 1
  fi
}

gen_key_if_missing() {
  local key_path="$1"
  mkdir -p "$(dirname "$key_path")"
  if [[ -f "$key_path" ]]; then
    echo "[SKIP] 已存在密钥: $key_path"
    chmod 600 "$key_path" 2>/dev/null || true
    return 0
  fi
  echo "[RUN ] 生成密钥: $key_path"
  ssh-keygen -t ed25519 -f "$key_path" -C "k3s-cluster" -N ""
  chmod 600 "$key_path" 2>/dev/null || true

  # 如选择生成 PuTTY 私钥，则调用 puttygen；若命令不存在则给出安装提示并退出
  if [[ "${GENERATE_PUTTY_PPK:-n}" == "y" ]]; then
    if ! command -v puttygen >/dev/null 2>&1; then
      echo "[ERR] 已选择生成 PuTTY 私钥，但当前系统未安装 puttygen。" >&2
      echo "      请先安装 puttygen 后重新运行本脚本，或在提示时选择不生成 PuTTY 私钥。" >&2
      echo "" >&2
      echo "      常见系统安装示例：" >&2
      echo "        Fedora / CentOS / RHEL : sudo dnf install putty 或 sudo dnf install putty-tools" >&2
      echo "        Debian / Ubuntu        : sudo apt install putty-tools" >&2
      echo "        openSUSE               : sudo zypper install putty" >&2
      exit 1
    fi
    local ppk_path="${key_path}.ppk"
    echo "[RUN ] 生成 PuTTY 私钥: $ppk_path"
    puttygen "$key_path" -o "$ppk_path"
  fi
}

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

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

  # 兼容 CRLF：先去掉 \r 再解析。注意 [group] 在正则中为字符类，用字符串相等匹配
  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 }
    }
  '
}

copy_key_to_host() {
  local key="$1"
  local pub="$2"
  local user="$3"
  local host="$4"
  local pass="${5:-}"

  echo "[RUN ] 配置 ${user}@${host}"

  # 情况 1：提供了密码，使用 sshpass 让所有 ssh/scp 命令自动带密码（不再弹出交互）
  if [[ -n "$pass" ]]; then
    sshpass -p "$pass" ssh -o StrictHostKeyChecking=accept-new "${user}@${host}" "mkdir -p ~/.ssh; :" 2>/dev/null || true
    sshpass -p "$pass" scp "$pub" "${user}@${host}:/tmp/k3s_pubkey.pub"
    sshpass -p "$pass" ssh "${user}@${host}" "grep -v 'k3s-cluster' ~/.ssh/authorized_keys > /tmp/ak.tmp 2>/dev/null || true; tr -d '\r' < /tmp/k3s_pubkey.pub >> /tmp/ak.tmp; mv /tmp/ak.tmp ~/.ssh/authorized_keys; chmod 600 ~/.ssh/authorized_keys; rm -f /tmp/k3s_pubkey.pub"

    echo "[RUN ] 将同一公钥写入 root"
    sshpass -p "$pass" scp "$pub" "${user}@${host}:/tmp/k3s_pubkey.pub"
    printf '%s\n' "$pass" | sshpass -p "$pass" ssh -t "${user}@${host}" "sudo -S bash -c 'mkdir -p /root/.ssh; touch /root/.ssh/authorized_keys; grep -v \"k3s-cluster\" /root/.ssh/authorized_keys > /tmp/ak.tmp 2>/dev/null || true; tr -d \"\\r\" < /tmp/k3s_pubkey.pub >> /tmp/ak.tmp; mv /tmp/ak.tmp /root/.ssh/authorized_keys; chmod 600 /root/.ssh/authorized_keys'; rm -f /tmp/k3s_pubkey.pub"
    return 0
  fi

  # 情况 2：未提供密码，退回到简单的逐步交互模式（每一步按需提示密码）
  ssh "${user}@${host}" "mkdir -p ~/.ssh; :" 2>/dev/null || true
  scp "$pub" "${user}@${host}:/tmp/k3s_pubkey.pub"
  ssh "${user}@${host}" "grep -v 'k3s-cluster' ~/.ssh/authorized_keys > /tmp/ak.tmp 2>/dev/null || true; tr -d '\r' < /tmp/k3s_pubkey.pub >> /tmp/ak.tmp; mv /tmp/ak.tmp ~/.ssh/authorized_keys; chmod 600 ~/.ssh/authorized_keys; rm -f /tmp/k3s_pubkey.pub"

  echo "[RUN ] 将同一公钥写入 root"
  scp "$pub" "${user}@${host}:/tmp/k3s_pubkey.pub"
  ssh -t "${user}@${host}" "sudo bash -c 'mkdir -p /root/.ssh; touch /root/.ssh/authorized_keys; grep -v \"k3s-cluster\" /root/.ssh/authorized_keys > /tmp/ak.tmp 2>/dev/null || true; tr -d \"\\r\" < /tmp/k3s_pubkey.pub >> /tmp/ak.tmp; mv /tmp/ak.tmp /root/.ssh/authorized_keys; chmod 600 /root/.ssh/authorized_keys'; rm -f /tmp/k3s_pubkey.pub"
}

print_title "K3s 节点 SSH 密钥批量配置（控制节点 + 工作节点，每节点一把密钥）"

ensure_cmd ssh-keygen
ensure_cmd ssh-copy-id
ensure_cmd ssh

# 默认显示相对路径（相对于仓库根）
INVENTORY_REL="ansible/inventory.ini"
INVENTORY_PATH="$(ask_default "Ansible inventory 路径（相对仓库根 ${ROOT_DIR}）" "$INVENTORY_REL")"
[[ "$INVENTORY_PATH" != /* ]] && INVENTORY_PATH="${ROOT_DIR}/${INVENTORY_PATH}"
[[ ! -f "$INVENTORY_PATH" ]] && { echo "[ERR] 找不到 inventory: $INVENTORY_PATH" >&2; exit 1; }

# 交互输入：用户名（有默认值）、密码（可选，用于后续 SSH/sudo）、是否生成 PuTTY 私钥
echo ""
echo "--- 交互输入用户名与密码 ---"
SSH_USER="$(ask_default "SSH 登录用户名（直接回车使用默认 jack）" "$SSH_USER_DEFAULT")"

print_title "1) 解析节点列表（含控制节点 + 工作节点）"
# k3s_nodes 为 children 组，需分别解析 k3s_server 与 k3s_worker 后合并
K3S_HOSTS=()
for grp in k3s_server k3s_worker; do
  while IFS= read -r h; do
    [[ -n "$h" ]] && K3S_HOSTS+=("$h")
  done < <(parse_worker_hosts "$INVENTORY_PATH" "$grp" 2>/dev/null || true)
done

if [[ "${#K3S_HOSTS[@]}" -eq 0 ]]; then
  echo "[WARN] 未在 ${INVENTORY_PATH} 中找到 k3s_server 或 k3s_worker 的任何主机" >&2
  echo "请检查：1) 是否存在 [k3s_server]、[k3s_worker] 节；2) 节下是否有主机行（含 ansible_host=）；3) 文件是否为 CRLF 换行（可用 sed -i 's/\r$//' 修复）" >&2
  exit 1
fi

echo "[INFO] 将配置以下节点（含控制节点）："
for h in "${K3S_HOSTS[@]}"; do
  echo "  - ${h}"
done

# 交互输入密码：用于 SSH 登录与 sudo，仅存于脚本变量、脚本结束即丢弃
JACK_PASS=""
printf "请输入 %s 的密码（用于 SSH 与 sudo，仅本次使用、不落盘；直接回车则每步手动输入）: " "$SSH_USER" >&2
read -rs JACK_PASS
echo "" >&2
# 如果用户输入了密码，但本机未安装 sshpass，则给出提示并退出
if [[ -n "${JACK_PASS:-}" ]]; then
  if ! command -v sshpass >/dev/null 2>&1; then
    echo "[ERR] 脚本已进入“一次输入密码并复用”的模式，但当前系统未安装 sshpass。" >&2
    echo "      请先安装 sshpass 后重新运行本脚本，或在密码处直接回车改为逐步交互输入模式。" >&2
    echo "" >&2
    echo "      常见系统安装示例：" >&2
    echo "        Fedora / CentOS / RHEL : sudo dnf install sshpass" >&2
    echo "        Debian / Ubuntu        : sudo apt install sshpass" >&2
    echo "        openSUSE               : sudo zypper install sshpass" >&2
    exit 1
  fi
fi

# 是否同时为每把新生成的 OpenSSH 密钥生成一份 PuTTY 私钥（.ppk）
GENERATE_PUTTY_PPK="$(ask_default "是否为新密钥同时生成 PuTTY 私钥（.ppk）？(y/N)" "N")"
GENERATE_PUTTY_PPK="${GENERATE_PUTTY_PPK,,}"  # 转小写
if [[ "$GENERATE_PUTTY_PPK" == "y" ]]; then
  GENERATE_PUTTY_PPK="y"
else
  GENERATE_PUTTY_PPK="n"
fi

print_title "2) 为每个节点生成 SSH 密钥（如不存在）并分发到 jack + root"
declare -A KEY_PATHS=()
for h in "${K3S_HOSTS[@]}"; do
  # 针对每个节点，单独一个密钥文件，默认名包含 IP/主机名
  DEFAULT_KEY_PATH="${HOME}/.ssh/id_ed25519_k3s_${h}"
  KEY_PATH="$(ask_default "  私钥路径（用于 ${SSH_USER}@${h}）" "$DEFAULT_KEY_PATH")"
  gen_key_if_missing "$KEY_PATH"
  copy_key_to_host "$KEY_PATH" "${KEY_PATH}.pub" "$SSH_USER" "$h" "${JACK_PASS:-}"
  KEY_PATHS["$h"]="$KEY_PATH"
done

# 丢弃密码，避免残留在 shell 环境
unset -v JACK_PASS

print_title "完成"
echo "每个节点对应的私钥路径（供 Ansible inventory 中 ansible_ssh_private_key_file 使用）："
for h in "${K3S_HOSTS[@]}"; do
  echo "  root@${h} -> ${KEY_PATHS[$h]}"
done
echo
echo "登录示例："
FIRST_HOST="${K3S_HOSTS[0]}"
echo "  ssh -i \"${KEY_PATHS[$FIRST_HOST]}\" root@${FIRST_HOST}"