#!/bin/bash

# Nomad 集群领导者发现与访问脚本
# 此脚本自动发现当前 Nomad 集群领导者并执行相应命令

# 默认服务器列表（可根据实际情况修改）
SERVERS=(
    "100.116.158.95"  # bj-semaphore.global
    "100.81.26.3"     # ash1d.global
    "100.103.147.94"  # ash2e.global
    "100.90.159.68"   # ch2.global
    "100.86.141.112"  # ch3.global
    "100.98.209.50"   # bj-onecloud1.global
    "100.120.225.29"  # de.global
)

# 超时设置（秒）
TIMEOUT=5

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 打印帮助信息
function show_help() {
    echo "Nomad 集群领导者发现与访问脚本"
    echo ""
    echo "用法: $0 [选项] [nomad命令]"
    echo ""
    echo "选项:"
    echo "  -h, --help          显示此帮助信息"
    echo "  -s, --server IP     指定初始服务器IP"
    echo "  -t, --timeout SECS  设置超时时间（默认: $TIMEOUT 秒）"
    echo "  -l, --list-servers  列出所有配置的服务器"
    echo "  -c, --check-leader  仅检查领导者，不执行命令"
    echo ""
    echo "示例:"
    echo "  $0 node status                    # 使用自动发现的领导者查看节点状态"
    echo "  $0 -s 100.116.158.95 job status   # 指定初始服务器查看作业状态"
    echo "  $0 -c                             # 仅检查当前领导者"
    echo ""
}

# 列出所有配置的服务器
function list_servers() {
    echo -e "${YELLOW}配置的服务器列表:${NC}"
    for server in "${SERVERS[@]}"; do
        echo "  - $server"
    done
}

# 发现领导者
function discover_leader() {
    local initial_server=$1
    
    # 如果指定了初始服务器，先尝试使用它
    if [ -n "$initial_server" ]; then
        echo -e "${YELLOW}尝试从服务器 $initial_server 发现领导者...${NC}" >&2
        leader=$(curl -s --max-time $TIMEOUT "http://${initial_server}:4646/v1/status/leader" 2>/dev/null | sed 's/"//g')
        if [ -n "$leader" ] && [ "$leader" != "" ]; then
            # 将RPC端口(4647)替换为HTTP端口(4646)
            leader=$(echo "$leader" | sed 's/:4647$/:4646/')
            echo -e "${GREEN}发现领导者: $leader${NC}" >&2
            echo "$leader"
            return 0
        fi
        echo -e "${RED}无法从 $initial_server 获取领导者信息${NC}" >&2
    fi
    
    # 遍历所有服务器尝试发现领导者
    echo -e "${YELLOW}遍历所有服务器寻找领导者...${NC}" >&2
    for server in "${SERVERS[@]}"; do
        echo -n "  检查 $server ... " >&2
        leader=$(curl -s --max-time $TIMEOUT "http://${server}:4646/v1/status/leader" 2>/dev/null | sed 's/"//g')
        if [ -n "$leader" ] && [ "$leader" != "" ]; then
            # 将RPC端口(4647)替换为HTTP端口(4646)
            leader=$(echo "$leader" | sed 's/:4647$/:4646/')
            echo -e "${GREEN}成功${NC}" >&2
            echo -e "${GREEN}发现领导者: $leader${NC}" >&2
            echo "$leader"
            return 0
        else
            echo -e "${RED}失败${NC}" >&2
        fi
    done
    
    echo -e "${RED}无法发现领导者，请检查集群状态${NC}" >&2
    return 1
}

# 解析命令行参数
INITIAL_SERVER=""
CHECK_LEADER_ONLY=false
NOMAD_COMMAND=()

while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help)
            show_help
            exit 0
            ;;
        -s|--server)
            INITIAL_SERVER="$2"
            shift 2
            ;;
        -t|--timeout)
            TIMEOUT="$2"
            shift 2
            ;;
        -l|--list-servers)
            list_servers
            exit 0
            ;;
        -c|--check-leader)
            CHECK_LEADER_ONLY=true
            shift
            ;;
        *)
            NOMAD_COMMAND+=("$1")
            shift
            ;;
    esac
done

# 主逻辑
echo -e "${YELLOW}Nomad 集群领导者发现与访问脚本${NC}" >&2
echo "==================================" >&2

# 发现领导者
LEADER=$(discover_leader "$INITIAL_SERVER")
if [ $? -ne 0 ]; then
    exit 1
fi

# 提取领导者IP和端口
LEADER_IP=$(echo "$LEADER" | cut -d':' -f1)
LEADER_PORT=$(echo "$LEADER" | cut -d':' -f2)

# 如果仅检查领导者，则退出
if [ "$CHECK_LEADER_ONLY" = true ]; then
    echo -e "${GREEN}当前领导者: $LEADER${NC}" >&2
    exit 0
fi

# 如果没有指定命令，显示交互式菜单
if [ ${#NOMAD_COMMAND[@]} -eq 0 ]; then
    echo -e "${YELLOW}未指定命令，请选择要执行的操作:${NC}" >&2
    echo "1) 查看节点状态" >&2
    echo "2) 查看作业状态" >&2
    echo "3) 查看服务器成员" >&2
    echo "4) 查看集群状态" >&2
    echo "5) 自定义命令" >&2
    echo "0) 退出" >&2
    
    read -p "请输入选项 (0-5): " choice
    
    case $choice in
        1) NOMAD_COMMAND=("node" "status") ;;
        2) NOMAD_COMMAND=("job" "status") ;;
        3) NOMAD_COMMAND=("server" "members") ;;
        4) NOMAD_COMMAND=("operator" "raft" "list-peers") ;;
        5) 
            read -p "请输入完整的 Nomad 命令: " -a NOMAD_COMMAND
            ;;
        0) exit 0 ;;
        *) 
            echo -e "${RED}无效选项${NC}" >&2
            exit 1
            ;;
    esac
fi

# 执行命令
echo -e "${YELLOW}执行命令: nomad ${NOMAD_COMMAND[*]} -address=http://${LEADER}${NC}" >&2
nomad "${NOMAD_COMMAND[@]}" -address="http://${LEADER}"

# 检查命令执行结果
if [ $? -eq 0 ]; then
    echo -e "${GREEN}命令执行成功${NC}" >&2
else
    echo -e "${RED}命令执行失败，可能需要重新发现领导者${NC}" >&2
    echo -e "${YELLOW}尝试重新发现领导者...${NC}" >&2
    NEW_LEADER=$(discover_leader)
    if [ $? -eq 0 ] && [ "$NEW_LEADER" != "$LEADER" ]; then
        echo -e "${YELLOW}领导者已更改，重新执行命令...${NC}" >&2
        nomad "${NOMAD_COMMAND[@]}" -address="http://${NEW_LEADER}"
    else
        echo -e "${RED}无法恢复，请检查集群状态${NC}" >&2
        exit 1
    fi
fi