#!/bin/bash
set -euo pipefail

#  =========================== 全局配置与常量 ==========================
# 颜色定义（兼容终端）
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # 无颜色

# 打印工具函数（统一格式）
info() { echo -e "${YELLOW}[INFO]${NC} $1"; }
success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; }
step() { echo -e "\n${BLUE}[STEP $1/$2]${NC} $3"; }

# 配置参数（默认值）
ACTION=""
BACKENDS=""
HAPROXY_PORT=6443
HAPROXY_CONF="/etc/haproxy/haproxy.cfg"
HAPROXY_DIR="/etc/haproxy"
HAPROXY_LOG_CONF="/etc/rsyslog.d/haproxy.conf"
HAPROXY_LOG="/var/log/haproxy.log"
STATS_PORT=9090  # Web统计页面端口
STATS_USER="admin"
STATS_PASS="Admin@123"


# ========================== 帮助信息 ==========================
usage() {
    cat << EOF
HAProxy 负载均衡配置脚本（适用于K8s API Server负载）
功能：一键安装/卸载HAProxy，配置TCP负载均衡（默认适配K8s 6443端口）

用法: $0 [操作] [选项]

操作:
  --install   安装并配置HAProxy（需指定后端节点）
  --uninstall 卸载HAProxy并清理配置文件

安装选项:
  --backend <节点列表>  必选，后端服务节点（格式：IP:端口,IP:端口...，如192.168.1.10:6443,192.168.1.11:6443）
  --port <端口>        可选，HAProxy监听端口（默认6443，需为1-65535之间的整数）

示例:
  安装（监听8443端口，负载两个后端节点）:
    sudo $0 --install --backend 192.168.1.10:6443,192.168.1.11:6443 --port 8443

  卸载:
    sudo $0 --uninstall
EOF
    exit 1
}


# ========================== 参数校验 ==========================
# 校验端口有效性（1-65535）
validate_port() {
    local port=$1
    if ! [[ "$port" =~ ^[0-9]+$ ]] || [ "$port" -lt 1 ] || [ "$port" -gt 65535 ]; then
        error "无效的端口：$port（必须是1-65535之间的整数）"
        exit 1
    fi
}

# 校验后端节点格式（IP:端口）
validate_backends() {
    local backends=$1
    if [ -z "$backends" ]; then
        error "后端节点列表不能为空（使用--backend指定，格式：IP:端口,IP:端口）"
        exit 1
    fi

    IFS=',' read -ra BACKEND_LIST <<< "$backends"
    for node in "${BACKEND_LIST[@]}"; do
        # 拆分IP和端口
        local ip=$(echo "$node" | cut -d':' -f1)
        local port=$(echo "$node" | cut -d':' -f2)
        
        # 校验IP格式（简单校验，不严格匹配所有IP规则）
        if ! [[ "$ip" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
            error "后端节点IP格式无效：$ip（节点：$node）"
            exit 1
        fi
        
        # 校验端口
        validate_port "$port"
    done
}


# ========================== 安装逻辑 ==========================
install_haproxy() {
    # 前置参数校验
    validate_backends "$BACKENDS"
    validate_port "$HAPROXY_PORT"

    # 拆分后端节点列表（后续生成配置用）
    IFS=',' read -ra BACKEND_LIST <<< "$BACKENDS"

    step 1 5 "检查HAProxy安装状态"
    if dpkg -l haproxy &>/dev/null; then
        info "检测到HAProxy已安装，将更新配置"
    else
        info "HAProxy未安装，开始安装..."
        step 2 5 "安装HAProxy软件包"
        info "更新系统包索引..."
        if apt update -qq; then
            success "系统包索引更新完成"
        else
            error "系统包索引更新失败，请检查网络"
            exit 1
        fi

        info "安装HAProxy..."
        if apt install -y -qq haproxy; then
            success "HAProxy安装完成"
        else
            error "HAProxy安装失败"
            exit 1
        fi
    fi

    step 3 5 "备份原有配置（若存在）"
    # 确保配置目录存在
    mkdir -p "$HAPROXY_DIR"
    
    # 备份当前配置（如果存在且不是备份文件）
    if [ -f "$HAPROXY_CONF" ] && ! [[ "$HAPROXY_CONF" =~ \.bak\.[0-9]+$ ]]; then
        local backup_file="${HAPROXY_CONF}.bak.$(date +%F_%H%M%S)"
        mv "$HAPROXY_CONF" "$backup_file"
        success "原有配置已备份至：$backup_file"
    else
        info "无有效原有配置，无需备份"
    fi

    step 4 5 "生成HAProxy配置文件"
    info "正在生成配置（监听端口：$HAPROXY_PORT，后端节点：${#BACKEND_LIST[@]}个）..."
    
    # 写入主配置
    cat > "$HAPROXY_CONF" << EOF
global
    log         127.0.0.1 local2  # 日志输出到local2设备
    chroot      /var/lib/haproxy   # 安全沙箱目录
    pidfile     /var/run/haproxy.pid  # PID文件路径
    maxconn     10000              # 最大并发连接数
    user        haproxy            # 运行用户
    group       haproxy            # 运行组
    daemon                          # 后台运行模式

# Web统计页面配置（监控负载状态）
listen haproxy-stats
    bind 0.0.0.0:$STATS_PORT      # 监听所有网卡的统计端口
    mode http                      # HTTP模式（统计页面基于HTTP）
    stats enable                   # 启用统计功能
    stats uri /stats               # 统计页面路径
    stats auth $STATS_USER:$STATS_PASS  # 访问认证（用户名:密码）
    stats refresh 30s              # 页面自动刷新时间
    stats show-node                # 显示节点信息
    stats show-legends             # 显示统计说明

defaults
    mode                    tcp     # 默认TCP模式（适配K8s API）
    log                     global  # 继承global的日志配置
    option                  tcplog  # 记录TCP日志
    option                  dontlognull  # 不记录空连接日志
    option                  redispatch  # 连接失败时重新分配
    retries                 3       # 重试次数
    timeout connect         10s     # 连接超时时间
    timeout client          1m      # 客户端超时时间
    timeout server          1m      # 服务端超时时间
    maxconn                 8000    # 每个进程最大连接数

# 前端监听配置（接收客户端请求）
frontend k8s-api-frontend
    bind *:$HAPROXY_PORT           # 监听所有网卡的指定端口
    mode tcp
    default_backend k8s-api-backend  # 转发到后端集群

# 后端集群配置（负载均衡目标）
backend k8s-api-backend
    mode tcp
    balance roundrobin             # 轮询负载均衡算法
EOF

    # 添加后端节点到配置
    for idx in "${!BACKEND_LIST[@]}"; do
        node="${BACKEND_LIST[$idx]}"
        echo "    server master-$idx $node check fall 3 rise 2 weight 10" >> "$HAPROXY_CONF"
    done

    if [ -f "$HAPROXY_CONF" ]; then
        success "配置文件生成成功：$HAPROXY_CONF"
    else
        error "配置文件生成失败"
        exit 1
    fi

    step 5 5 "配置日志与启动服务"
    info "配置HAProxy日志输出..."
    # 配置rsyslog接收HAProxy日志
    echo "local2.* $HAPROXY_LOG" > "$HAPROXY_LOG_CONF"
    # 重启rsyslog使日志配置生效
    if systemctl restart rsyslog; then
        success "日志配置生效（日志文件：$HAPROXY_LOG）"
    else
        error "rsyslog重启失败，日志可能无法正常记录"
    fi

    info "启动并设置HAProxy开机自启..."
    systemctl enable haproxy --now
    systemctl restart haproxy
    # 检查服务状态
    if systemctl is-active --quiet haproxy; then
        success "HAProxy服务启动成功（运行中）"
    else
        error "HAProxy服务启动失败，请查看日志：$HAPROXY_LOG"
        exit 1
    fi

    # 安装完成：汇总核心信息
    echo -e "\n========================================"
    echo -e "${GREEN}【HAProxy安装配置完成】${NC}"
    echo -e "${YELLOW}核心信息汇总：${NC}"
    echo "  1. 监听地址：所有网卡（0.0.0.0）:$HAPROXY_PORT"
    echo "  2. 后端节点（共${#BACKEND_LIST[@]}个）："
    for node in "${BACKEND_LIST[@]}"; do
        echo "     - $node"
    done
    echo "  3. Web统计页面："
    echo "     - 地址：http://<服务器IP>:$STATS_PORT/stats"
    echo "     - 用户名：$STATS_USER"
    echo "     - 密码：$STATS_PASS"
    echo "  4. 配置文件路径：$HAPROXY_CONF"
    echo "  5. 日志文件路径：$HAPROXY_LOG"
    echo -e "========================================"
}


# ========================== 卸载逻辑 ==========================
uninstall_haproxy() {
    step 1 4 "停止HAProxy服务"
    if systemctl is-active --quiet haproxy; then
        info "HAProxy服务正在运行，停止中..."
        if systemctl stop haproxy; then
            success "HAProxy服务已停止"
        else
            error "HAProxy服务停止失败，将继续执行卸载"
        fi
    else
        info "HAProxy服务未运行，跳过停止步骤"
    fi

    step 2 4 "卸载HAProxy软件包"
    if dpkg -l haproxy &>/dev/null; then
        info "检测到HAProxy已安装，开始卸载..."
        if apt purge -y -qq haproxy; then
            success "HAProxy软件包卸载完成"
        else
            error "HAProxy软件包卸载失败，建议手动执行：apt purge -y haproxy"
        fi

        # 清理依赖残留
        info "清理无用依赖..."
        apt autoremove -y -qq >/dev/null
        success "依赖残留清理完成"
    else
        info "HAProxy未安装，跳过卸载软件包步骤"
    fi

    step 3 4 "清理配置文件与目录"
    # 备份并删除主配置目录
    if [ -d "$HAPROXY_DIR" ]; then
        local backup_dir="/etc/haproxy.bak.$(date +%F_%H%M%S)"
        mv "$HAPROXY_DIR" "$backup_dir"
        success "配置目录已备份至：$backup_dir"
    else
        info "配置目录不存在，跳过备份"
    fi

    # 清理日志配置
    if [ -f "$HAPROXY_LOG_CONF" ]; then
        rm -f "$HAPROXY_LOG_CONF"
        # 重启rsyslog生效
        systemctl restart rsyslog >/dev/null 2>&1 || true
        success "日志配置文件已清理"
    else
        info "日志配置文件不存在，跳过清理"
    fi

    step 4 4 "清理日志文件"
    if [ -f "$HAPROXY_LOG" ]; then
        rm -f "$HAPROXY_LOG"
        success "日志文件已删除：$HAPROXY_LOG"
    else
        info "日志文件不存在，跳过清理"
    fi

    # 卸载完成提示
    echo -e "\n========================================"
    echo -e "${GREEN}【HAProxy卸载完成】${NC}"
    echo -e "${YELLOW}说明：${NC}"
    echo "  - 所有配置文件已备份（路径：/etc/haproxy.bak.xxx，如需恢复可手动迁移）"
    echo "  - 服务已停止，软件包已卸载（若之前安装过）"
    echo -e "========================================"
}


# ========================== 主逻辑 ==========================
main() {
    # 解析命令行参数
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --install|--uninstall)
                ACTION="$1"
                shift
                ;;
            --backend)
                if [ $# -lt 2 ]; then
                    error "--backend选项需指定节点列表（格式：IP:端口,IP:端口）"
                    usage
                fi
                BACKENDS="$2"
                shift 2
                ;;
            --port)
                if [ $# -lt 2 ]; then
                    error "--port选项需指定端口（1-65535）"
                    usage
                fi
                HAPROXY_PORT="$2"
                shift 2
                ;;
            *)
                error "未知参数：$1"
                usage
                ;;
        esac
    done

    # 校验操作类型
    if [ -z "$ACTION" ]; then
        error "必须指定操作：--install（安装）或--uninstall（卸载）"
        usage
    fi

    # 检查root权限
    if [[ $EUID -ne 0 ]]; then
        error "脚本必须以root权限运行，请使用：sudo $0 ..."
        exit 1
    fi

    # 执行对应操作
    case "$ACTION" in
        --install)
            install_haproxy
            ;;
        --uninstall)
            uninstall_haproxy
            ;;
        *)
            usage
            ;;
    esac
}

# 启动主逻辑
main "$@"