#!/usr/bin/env bash
#
# setup-ssh.sh - 配置本机到远程主机的 SSH 免密连接并设置别名
#
# 用法:
#   ./setup-ssh.sh <用户名> <IP> <密码> [别名] [端口]
#
# 示例:
#   ./setup-ssh.sh book 192.168.1.124 mypass
#   生成别名: book-124  (用户名 + IP 最后一段数字)
#   ./setup-ssh.sh book 192.168.1.124 mypass myserver 2222
#   使用自定义别名: myserver, 端口: 2222
#
# 之后即可使用: ssh book-124  (或自定义别名)
#
set -euo pipefail

# ---------- 参数校验 ----------
if [[ $# -lt 3 ]]; then
    echo "用法: $0 <用户名> <IP> <密码> [别名] [端口]" >&2
    echo "示例: $0 book 192.168.1.124 mypass myserver 2222" >&2
    exit 1
fi

USER_NAME="$1"
IP="$2"
PASSWORD="$3"
ALIAS_ARG="${4:-}"
PORT="${5:-22}"

# 校验 IP 格式
if ! [[ "$IP" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
    echo "错误: IP 格式不正确: $IP" >&2
    exit 1
fi

# 别名: 若传入了别名参数则使用传入值, 否则用 用户名 + "-" + IP 最后一段
if [[ -n "$ALIAS_ARG" ]]; then
    ALIAS="$ALIAS_ARG"
else
    LAST_OCTET="${IP##*.}"
    ALIAS="${USER_NAME}-${LAST_OCTET}"
fi

SSH_DIR="$HOME/.ssh"
KEY_FILE="$SSH_DIR/id_ed25519"
CONFIG_FILE="$SSH_DIR/config"

# ---------- 依赖检查 ----------
if ! command -v sshpass >/dev/null 2>&1; then
    echo "缺少 sshpass, 正在尝试安装..." >&2
    if command -v apt-get >/dev/null 2>&1; then
        sudo apt-get update -qq && sudo apt-get install -y sshpass
    elif command -v yum >/dev/null 2>&1; then
        sudo yum install -y sshpass
    elif command -v apk >/dev/null 2>&1; then
        sudo apk add --no-cache sshpass
    else
        echo "错误: 请先手动安装 sshpass" >&2
        exit 1
    fi
fi

# ---------- 准备 .ssh 目录 ----------
mkdir -p "$SSH_DIR"
chmod 700 "$SSH_DIR"

# ---------- 生成密钥(如不存在) ----------
if [[ ! -f "$KEY_FILE" ]]; then
    echo "生成 SSH 密钥: $KEY_FILE"
    ssh-keygen -t ed25519 -f "$KEY_FILE" -N "" -C "$(whoami)@$(hostname)"
fi

# ---------- 复制公钥到远程 ----------
echo "复制公钥到 ${USER_NAME}@${IP}:${PORT} ..."
sshpass -p "$PASSWORD" ssh-copy-id \
    -i "${KEY_FILE}.pub" \
    -o StrictHostKeyChecking=no \
    -o UserKnownHostsFile="$SSH_DIR/known_hosts" \
    -p "$PORT" \
    "${USER_NAME}@${IP}"

# ---------- 写入/更新 ssh config 别名 ----------
touch "$CONFIG_FILE"
chmod 600 "$CONFIG_FILE"

# 若已存在同名 Host 块, 先删除
if grep -qE "^[[:space:]]*Host[[:space:]]+${ALIAS}([[:space:]]|\$)" "$CONFIG_FILE"; then
    echo "更新已存在的别名: $ALIAS"
    tmp_file="$(mktemp)"
    awk -v alias="$ALIAS" '
        BEGIN { skip = 0 }
        /^[[:space:]]*Host[[:space:]]+/ {
            if ($2 == alias) { skip = 1; next } else { skip = 0 }
        }
        skip == 1 && /^[[:space:]]*Host[[:space:]]+/ { skip = 0 }
        skip == 0 { print }
    ' "$CONFIG_FILE" > "$tmp_file"
    mv "$tmp_file" "$CONFIG_FILE"
    chmod 600 "$CONFIG_FILE"
fi

cat >> "$CONFIG_FILE" <<EOF

Host ${ALIAS}
    HostName ${IP}
    User ${USER_NAME}
    Port ${PORT}
    IdentityFile ${KEY_FILE}
EOF

echo
echo "✅ 配置完成!"
echo "   别名: ${ALIAS}"
echo "   连接: ssh ${ALIAS}"