#!/bin/bash
#
# VPS初始化一键脚本
# 整合了系统更新、登录安全设置、系统清理、Docker安装、防火墙设置、时区设置、
# 内存优化、Fail2ban安装和BBR加速
# 使用方法: chmod +x vps_init.sh && ./vps_init.sh
#

# ===========================================
# 用户设置区域 - 根据需要修改
# ===========================================
# NEW_PASSWORD 已弃用 - 密码将在交互式设置中由用户输入
NEW_SSH_PORT="4399"               # SSH新端口号（默认值）
TIMEZONE="Asia/Shanghai"          # 时区设置
SWAP_SIZE=1024                    # 交换分区大小(MB，默认值)
SWAPPINESS=10                     # 交换倾向默认值(0-100，越大越倾向使用交换)
# ===========================================

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # 恢复默认颜色

# 确保脚本以root权限运行
if [ "$(id -u)" -ne 0 ]; then
    echo -e "${RED}错误: 必须以root用户运行此脚本!${NC}"
    exit 1
fi

# 显示欢迎信息
echo -e "${GREEN}=============================================${NC}"
echo -e "${BLUE}        VPS初始化一键脚本开始执行             ${NC}"
echo -e "${GREEN}=============================================${NC}"
echo ""

# ===========================================
# 交互式设置选项
# ===========================================
echo -e "${BLUE}进行交互式设置...${NC}"

# 询问是否修改SSH端口
echo -e "${YELLOW}1. 是否修改SSH端口? 当前设置为: ${NEW_SSH_PORT}${NC}"
while true; do
    read -p "修改SSH端口? (y/n, 默认n): " CHANGE_SSH_PORT
    # 设置默认值为否，用户按Enter就不修改
    CHANGE_SSH_PORT=${CHANGE_SSH_PORT:-n}
    if [[ "$CHANGE_SSH_PORT" =~ ^[Yy]$ ]] || [[ "$CHANGE_SSH_PORT" =~ ^[Nn]$ ]]; then
        break
    else
        echo -e "${RED}无效的输入，请输入 y 或 n${NC}"
    fi
done

if [[ "$CHANGE_SSH_PORT" =~ ^[Yy]$ ]]; then
    echo -e "${GREEN}SSH端口将被修改为: ${NEW_SSH_PORT}${NC}"
else
    echo -e "${GREEN}保持SSH端口不变: $NEW_SSH_PORT${NC}"
fi

# 询问是否修改密码
echo -e "${YELLOW}2. 是否修改root密码?${NC}"
while true; do
    read -p "修改root密码? (y/n, 默认n): " CHANGE_PASSWORD
    # 设置默认值为否，用户按Enter就不修改
    CHANGE_PASSWORD=${CHANGE_PASSWORD:-n}
    if [[ "$CHANGE_PASSWORD" =~ ^[Yy]$ ]] || [[ "$CHANGE_PASSWORD" =~ ^[Nn]$ ]]; then
        break
    else
        echo -e "${RED}无效的输入，请输入 y 或 n${NC}"
    fi
done

if [[ "$CHANGE_PASSWORD" =~ ^[Yy]$ ]]; then
    # 让用户输入新密码
    while true; do
        read -sp "请输入新的root密码: " NEW_PASSWORD
        echo ""
        if [ -n "$NEW_PASSWORD" ]; then
            # 要求用户确认密码
            read -sp "请再次输入密码以确认: " NEW_PASSWORD_CONFIRM
            echo ""
            if [ "$NEW_PASSWORD" = "$NEW_PASSWORD_CONFIRM" ]; then
                echo -e "${GREEN}密码已确认，将在后续步骤中修改${NC}"
                break
            else
                echo -e "${RED}两次输入的密码不一致，请重新输入${NC}"
            fi
        else
            echo -e "${RED}密码不能为空，请重新输入${NC}"
        fi
    done
else
    echo -e "${GREEN}保持root密码不变${NC}"
fi

# 询问是否修改主机名
echo -e "${YELLOW}3. 是否修改主机名?${NC}"
while true; do
    read -p "修改主机名? (y/n, 默认n): " CHANGE_HOSTNAME
    # 设置默认值为否，用户按Enter就不修改
    CHANGE_HOSTNAME=${CHANGE_HOSTNAME:-n}
    if [[ "$CHANGE_HOSTNAME" =~ ^[Yy]$ ]] || [[ "$CHANGE_HOSTNAME" =~ ^[Nn]$ ]]; then
        break
    else
        echo -e "${RED}无效的输入，请输入 y 或 n${NC}"
    fi
done

if [[ "$CHANGE_HOSTNAME" =~ ^[Yy]$ ]]; then
    # 显示当前主机名
    CURRENT_HOSTNAME=$(hostname)
    echo -e "${YELLOW}当前主机名: ${CURRENT_HOSTNAME}${NC}"
    
    # 让用户输入新主机名
    read -p "请输入新的主机名: " NEW_HOSTNAME
    if [ -n "$NEW_HOSTNAME" ]; then
        CHANGE_HOSTNAME_FLAG=true
        echo -e "${GREEN}主机名将被修改为: ${NEW_HOSTNAME}${NC}"
    else
        CHANGE_HOSTNAME_FLAG=false
        echo -e "${RED}主机名不能为空，将保持不变${NC}"
    fi
else
    CHANGE_HOSTNAME_FLAG=false
    echo -e "${GREEN}保持主机名不变${NC}"
fi

# 询问交换分区大小(SWAP_SIZE)
echo -e "${YELLOW}4. 设置交换分区大小 (MB)? 当前默认: ${SWAP_SIZE}MB${NC}"
echo -e "${YELLOW}   常用值: 1024 / 2048 (建议为物理内存的1-2倍)${NC}"
while true; do
    read -p "请输入交换分区大小 (MB, 默认${SWAP_SIZE}): " INPUT_SWAP_SIZE
    # 用户直接回车则使用默认值
    INPUT_SWAP_SIZE=${INPUT_SWAP_SIZE:-$SWAP_SIZE}
    # 校验：必须为正整数(MB)
    if [[ "$INPUT_SWAP_SIZE" =~ ^[0-9]+$ ]] && [ "$INPUT_SWAP_SIZE" -gt 0 ]; then
        SWAP_SIZE=$INPUT_SWAP_SIZE
        echo -e "${GREEN}交换分区大小将设置为: ${SWAP_SIZE}MB${NC}"
        break
    else
        echo -e "${RED}无效的输入，请输入大于0的整数(单位MB)${NC}"
    fi
done

# 询问交换倾向(swappiness)
echo -e "${YELLOW}5. 设置交换倾向 swappiness (0-100)?${NC}"
echo -e "${YELLOW}   值越大越倾向使用交换(swap)，越小越倾向保留物理内存${NC}"
while true; do
    read -p "请输入 swappiness 值 (0-100, 默认${SWAPPINESS}): " INPUT_SWAPPINESS
    # 用户直接回车则使用默认值
    INPUT_SWAPPINESS=${INPUT_SWAPPINESS:-$SWAPPINESS}
    # 校验：必须为0-100之间的整数
    if [[ "$INPUT_SWAPPINESS" =~ ^[0-9]+$ ]] && [ "$INPUT_SWAPPINESS" -ge 0 ] && [ "$INPUT_SWAPPINESS" -le 100 ]; then
        SWAPPINESS=$INPUT_SWAPPINESS
        echo -e "${GREEN}swappiness 将设置为: ${SWAPPINESS}${NC}"
        break
    else
        echo -e "${RED}无效的输入，请输入 0 到 100 之间的整数${NC}"
    fi
done

echo -e "${BLUE}交互式设置完成${NC}"
echo ""

# 记录开始时间
START_TIME=$(date +%s)

# 检查系统类型
if [ -f /etc/debian_version ]; then
    OS_TYPE="debian"
    echo -e "${GREEN}检测到Debian/Ubuntu系统${NC}"
else
    echo -e "${YELLOW}警告: 此脚本主要为Debian/Ubuntu系统设计${NC}"
    echo -e "${YELLOW}部分功能可能在其他系统上不正常工作${NC}"
    OS_TYPE="other"
fi

# 创建日志文件
LOG_FILE="/var/log/vps_init_$(date +%Y%m%d_%H%M%S).log"
touch $LOG_FILE
echo "VPS初始化脚本开始执行: $(date)" > $LOG_FILE

# 创建摘要信息文件
SUMMARY_FILE="/root/vps_init.txt"
TEMP_HISTORY_FILE="/tmp/vps_init_history_$$.txt"

# 定义日志函数
log() {
    echo -e "$1" | tee -a $LOG_FILE
}

# 定义摘要信息函数（同时输出到屏幕、日志文件和摘要文件）
log_summary() {
    echo -e "$1" | tee -a $LOG_FILE | sed 's/\x1b\[[0-9;]*m//g' >> $SUMMARY_FILE
}

# 定义错误处理函数
handle_error() {
    local exit_code=$?
    local line_no=$1
    if [ $exit_code -ne 0 ]; then
        log "${RED}错误: 在第 $line_no 行发生错误，退出代码: $exit_code${NC}"
        log "${RED}请检查日志文件: $LOG_FILE${NC}"
    fi
}

# 如果摘要文件已存在，提取历史端口和密码信息
if [ -f "$SUMMARY_FILE" ]; then
    log "${YELLOW}检测到已存在的配置文件，正在提取历史信息...${NC}"

    # 创建临时文件保存历史信息
    echo "历史配置信息（从上次运行保留）" > $TEMP_HISTORY_FILE
    echo "-------------------------------------------------------" >> $TEMP_HISTORY_FILE

    # 提取所有SSH端口历史记录
    SSH_PORTS=$(grep -E "SSH端口已更改为:" "$SUMMARY_FILE" | sed 's/.*SSH端口已更改为: //' | sed 's/[^0-9]//g' || true)

    # 提取所有密码历史记录
    PASSWORDS=$(grep -E "实际密码:" "$SUMMARY_FILE" | sed 's/.*实际密码: //' || true)

    # 写入历史SSH端口
    if [ -n "$SSH_PORTS" ]; then
        echo "历史SSH端口:" >> $TEMP_HISTORY_FILE
        echo "$SSH_PORTS" | while IFS= read -r port; do
            if [ -n "$port" ]; then
                echo "  - $port" >> $TEMP_HISTORY_FILE
            fi
        done
    fi

    # 写入历史密码
    if [ -n "$PASSWORDS" ]; then
        echo "历史密码:" >> $TEMP_HISTORY_FILE
        echo "$PASSWORDS" | while IFS= read -r pwd; do
            if [ -n "$pwd" ]; then
                echo "  - $pwd" >> $TEMP_HISTORY_FILE
            fi
        done
    fi

    echo "-------------------------------------------------------" >> $TEMP_HISTORY_FILE
    echo "" >> $TEMP_HISTORY_FILE
fi

# 创建新的摘要文件
touch $SUMMARY_FILE
echo "VPS初始化摘要信息" > $SUMMARY_FILE
echo "本次生成时间: $(date)" >> $SUMMARY_FILE
echo "=======================================================" >> $SUMMARY_FILE
echo "" >> $SUMMARY_FILE

# 如果有历史信息文件，追加到摘要文件
if [ -f "$TEMP_HISTORY_FILE" ]; then
    cat $TEMP_HISTORY_FILE >> $SUMMARY_FILE
    rm -f $TEMP_HISTORY_FILE
fi

# 设置错误跟踪
trap 'handle_error $LINENO' ERR

# ===========================================
# 安全更新软件源：自动检测并禁用失效的软件源
# （例如 Debian 版本归档后失效的 *-backports 源）后重试
# ===========================================
apt_update_safe() {
    local tmp retry=0
    tmp=$(mktemp)
    while :; do
        # 执行更新，忽略退出码（失效的源会让 apt 返回非0）
        apt-get update > "$tmp" 2>&1 || true
        cat "$tmp" >> "$LOG_FILE"
        cat "$tmp"

        # 没有检测到失效源，更新成功
        if ! grep -qiE "no longer has a Release file|404 +Not Found|Failed to fetch" "$tmp"; then
            rm -f "$tmp"
            return 0
        fi

        # 已尝试修复过一次，仍有问题则忽略并继续，避免死循环
        if [ "$retry" -ge 1 ]; then
            log "${YELLOW}部分软件源仍不可用，将忽略失效源并继续${NC}"
            rm -f "$tmp"
            return 0
        fi

        log "${YELLOW}检测到失效的软件源，正在自动禁用...${NC}"

        # 从 apt 输出中提取失效的套件名（位于 "<url> <suite> Release" 中）
        local suites suite
        suites=$(grep -oE "The repository '[^']+ Release'" "$tmp" \
            | sed -E "s/The repository '(.*) Release'/\1/" \
            | awk '{print $NF}' | sort -u)

        if [ -z "$suites" ]; then
            log "${YELLOW}无法解析失效套件名，将忽略失效源并继续${NC}"
            rm -f "$tmp"
            return 0
        fi

        for suite in $suites; do
            [ -n "$suite" ] || continue
            log "${YELLOW}  - 禁用失效套件: ${suite}${NC}"
            for f in /etc/apt/sources.list $(ls /etc/apt/sources.list.d/*.list 2>/dev/null); do
                [ -f "$f" ] || continue
                # 注释掉引用该失效套件的未注释行
                sed -i -E "s|^([^#].*[[:space:]]${suite}[[:space:]].*)$|#\1|g" "$f"
            done
        done
        retry=$((retry + 1))
    done
}

# ===========================================
# 1. 系统更新
# ===========================================
log "${BLUE}[1/10] 系统更新开始...${NC}"

# 保存当前的sources.list作为备份
if [ -f "/etc/apt/sources.list" ]; then
    cp /etc/apt/sources.list /etc/apt/sources.list.bak
    log "${GREEN}备份了软件源配置文件${NC}"
fi

# 更新系统包
if [ "$OS_TYPE" = "debian" ]; then
    apt_update_safe
    DEBIAN_FRONTEND=noninteractive apt full-upgrade -y || log "${RED}系统升级失败${NC}"
    apt install -y wget curl sudo vim git ufw net-tools htop iftop || log "${RED}安装基础软件包失败${NC}"
    log "${GREEN}系统更新完成，安装了常用工具${NC}"
else
    log "${YELLOW}非Debian系统，跳过标准更新流程${NC}"
fi

# ===========================================
# 2. 主机名设置（如果用户选择了修改）
# ===========================================
if [ "$CHANGE_HOSTNAME_FLAG" = true ]; then
    log "${BLUE}[2/10] 设置主机名...${NC}"
    
    # 备份当前主机名配置
    cp /etc/hostname /etc/hostname.bak
    cp /etc/hosts /etc/hosts.bak
    
    # 修改主机名
    echo "$NEW_HOSTNAME" > /etc/hostname
    hostname "$NEW_HOSTNAME"
    
    # 更新hosts文件
    sed -i "s/127.0.1.1.*/127.0.1.1\t$NEW_HOSTNAME/g" /etc/hosts
    
    # 检查是否修改成功
    CURRENT_HOSTNAME=$(hostname)
    if [ "$CURRENT_HOSTNAME" = "$NEW_HOSTNAME" ]; then
        log "${GREEN}主机名已成功修改为: $NEW_HOSTNAME${NC}"
    else
        log "${RED}主机名修改失败，当前名称: $CURRENT_HOSTNAME${NC}"
    fi
    
    log "${GREEN}主机名设置完成${NC}"
else
    log "${YELLOW}[2/10] 跳过主机名设置...${NC}"
fi

# ===========================================
# 3. 登录安全设置
# ===========================================
log "${BLUE}[3/10] 设置登录安全...${NC}"

# 根据用户选择修改root密码
if [[ "$CHANGE_PASSWORD" =~ ^[Yy]$ ]]; then
    echo "root:$NEW_PASSWORD" | chpasswd
    if [ $? -eq 0 ]; then
        log "${GREEN}Root密码修改成功${NC}"
    else
        log "${RED}Root密码修改失败${NC}"
    fi
else
    log "${YELLOW}根据用户选择，保持root密码不变${NC}"
fi

# 根据用户选择修改SSH端口
if [[ "$CHANGE_SSH_PORT" =~ ^[Yy]$ ]]; then
    # 备份SSH配置文件
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak
    log "${GREEN}SSH配置已备份${NC}"

    # 修改SSH配置
    sed -i 's/^#\?PermitRootLogin.*/PermitRootLogin yes/g' /etc/ssh/sshd_config
    sed -i 's/^#\?PasswordAuthentication.*/PasswordAuthentication yes/g' /etc/ssh/sshd_config
    sed -i 's/#Port/Port/' /etc/ssh/sshd_config
    sed -i "s/Port [0-9]\+/Port $NEW_SSH_PORT/g" /etc/ssh/sshd_config

    # 读取修改后的SSH端口以确认更改
    NEW_PORT_CONFIGURED=$(grep -P "^Port\s+\d+" /etc/ssh/sshd_config | awk '{print $2}')
    if [ "$NEW_PORT_CONFIGURED" = "$NEW_SSH_PORT" ]; then
        log "${GREEN}SSH端口已修改为: $NEW_SSH_PORT${NC}"
    else
        log "${RED}SSH端口修改失败，当前设置: $NEW_PORT_CONFIGURED${NC}"
        # 尝试使用另一种方法修改
        echo "Port $NEW_SSH_PORT" >> /etc/ssh/sshd_config
        log "${YELLOW}尝试使用备选方法添加端口设置${NC}"
    fi

    # 重启SSH服务（Debian的服务名为ssh，部分系统为sshd，依次尝试）
    if systemctl restart sshd 2>/dev/null || systemctl restart ssh 2>/dev/null; then
        log "${GREEN}SSH服务重启成功${NC}"
    else
        log "${RED}SSH服务重启失败${NC}"
        # 尝试使用service命令
        service sshd restart || service ssh restart || log "${RED}无法重启SSH服务，请手动检查${NC}"
    fi

    # 检查SSH服务状态
    systemctl status sshd --no-pager || service sshd status || service ssh status
    log "${GREEN}SSH配置更改完成${NC}"
    log "${YELLOW}注意：新的SSH连接端口为 $NEW_SSH_PORT${NC}"
else
    log "${YELLOW}根据用户选择，保持SSH端口不变${NC}"
    
    # 即使不修改端口，仍然应该确保其他SSH安全设置
    sed -i 's/^#\?PermitRootLogin.*/PermitRootLogin yes/g' /etc/ssh/sshd_config
    sed -i 's/^#\?PasswordAuthentication.*/PasswordAuthentication yes/g' /etc/ssh/sshd_config
    
    log "${GREEN}SSH基本安全设置完成${NC}"
fi

# ===========================================
# 4. 系统清理
# ===========================================
log "${BLUE}[4/10] 系统清理开始...${NC}"

# 清理不需要的软件包
if [ "$OS_TYPE" = "debian" ]; then
    apt autoremove --purge -y
    apt clean -y
    apt autoclean -y
    apt remove --purge $(dpkg -l | awk '/^rc/ {print $2}') -y 2>/dev/null || log "${YELLOW}没有需要清理的软件包配置${NC}"
    
    # 清理日志
    journalctl --rotate
    journalctl --vacuum-time=1d
    journalctl --vacuum-size=50M
    log "${GREEN}系统日志已清理${NC}"
    
    # 清理旧内核(保留当前运行的内核)
    apt remove --purge $(dpkg -l | awk '/^ii linux-(image|headers)-[^ ]+/{print $2}' | grep -v $(uname -r | sed 's/-.*//') | xargs) -y 2>/dev/null || log "${YELLOW}没有可清理的旧内核${NC}"
    
    log "${GREEN}系统清理完成${NC}"
else
    log "${YELLOW}非Debian系统，跳过系统清理流程${NC}"
fi

# ===========================================
# 5. Docker安装
# ===========================================
log "${BLUE}[5/10] Docker安装开始...${NC}"

# 检查Docker是否已安装
if command -v docker &> /dev/null; then
    log "${GREEN}Docker已经安装，版本信息:${NC}"
    docker --version
else
    # 安装Docker
    if [ "$OS_TYPE" = "debian" ]; then
        # 先确保软件源可用（官方脚本内部也会执行 apt update）
        apt_update_safe

        # 使用官方安装脚本
        curl -fsSL https://get.docker.com -o get-docker.sh
        sh get-docker.sh || log "${RED}Docker安装脚本执行出现问题${NC}"
        rm -f get-docker.sh

        # 校验Docker是否真正安装成功
        if command -v docker &> /dev/null; then
            systemctl enable docker 2>/dev/null || log "${YELLOW}无法设置Docker开机自启${NC}"
            systemctl start docker 2>/dev/null || true

            # 安装Docker Compose
            if ! command -v docker-compose &> /dev/null; then
                curl -L "https://github.com/docker/compose/releases/latest/download/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
                chmod +x /usr/local/bin/docker-compose
                log "${GREEN}Docker Compose 安装完成${NC}"
            fi

            log "${GREEN}Docker安装完成，版本信息:${NC}"
            docker --version
            docker-compose --version 2>/dev/null || docker compose version 2>/dev/null || true
        else
            log "${RED}Docker安装失败，请检查网络或软件源后重试，可手动执行: curl -fsSL https://get.docker.com | sh${NC}"
        fi
    else
        log "${YELLOW}非Debian系统，请手动安装Docker${NC}"
    fi
fi

# ===========================================
# 6. 防火墙设置
# ===========================================
log "${BLUE}[6/10] 防火墙设置开始...${NC}"

# 安装UFW
if [ "$OS_TYPE" = "debian" ]; then
    apt_update_safe
    apt install -y ufw net-tools lsof || log "${RED}安装防火墙相关软件包(ufw/net-tools/lsof)失败${NC}"

    # 确保防火墙默认策略
    ufw default deny incoming
    ufw default allow outgoing
    
    # 获取当前SSH端口（如果有多个SSH端口，获取所有）
    CURRENT_SSH_PORT=$(grep -P "^Port\s+\d+" /etc/ssh/sshd_config | awk '{print $2}')
    if [ -z "$CURRENT_SSH_PORT" ]; then
        # 如果没找到，使用默认端口22
        CURRENT_SSH_PORT="22"
    fi
    
    # 总是添加新配置的SSH端口（防止被锁在系统之外）
    log "${GREEN}允许SSH端口 $NEW_SSH_PORT${NC}"
    ufw allow $NEW_SSH_PORT/tcp comment 'New SSH Port'
    
    # 如果当前SSH端口与新端口不同，添加当前SSH端口作为备份
    if [ "$CURRENT_SSH_PORT" != "$NEW_SSH_PORT" ]; then
        log "${GREEN}允许当前SSH端口 $CURRENT_SSH_PORT (备份)${NC}"
        ufw allow $CURRENT_SSH_PORT/tcp comment 'Current SSH Port (Backup)'
    fi
    
    # 添加基本Web服务端口
    log "${GREEN}允许HTTP/HTTPS端口${NC}"
    ufw allow 80/tcp comment 'HTTP'
    ufw allow 443/tcp comment 'HTTPS'
    
    # 检测活跃的网络连接和正在监听的端口
    log "${YELLOW}检测当前活跃的服务端口...${NC}"
    
    # 使用netstat查找监听的TCP端口（排除仅绑定本地回环127.x/::1的端口，例如SSH的X11转发6010）
    LISTENING_PORTS=$(netstat -tlnp 2>/dev/null | awk '/LISTEN/ {print $4}' | grep -vE '^(127\.|::1:)' | awk -F: '{print $NF}' | sort -n | uniq)
    
    # 使用lsof作为备选方法（同样排除本地回环地址）
    if [ -z "$LISTENING_PORTS" ] && command -v lsof &> /dev/null; then
        LISTENING_PORTS=$(lsof -iTCP -sTCP:LISTEN -P -n 2>/dev/null | awk 'NR>1 {print $9}' | grep -vE '^(127\.|\[::1\])' | awk -F: '{print $NF}' | sort -n | uniq)
    fi
    
    # 如果仍然为空，提示手动检查
    if [ -z "$LISTENING_PORTS" ]; then
        log "${YELLOW}未检测到活跃端口，只开放SSH、HTTP和HTTPS端口${NC}"
    else
        log "${GREEN}检测到以下活跃端口:${NC}"
        for PORT in $LISTENING_PORTS; do
            # 跳过SSH端口(已经添加过)，以及常见的本地服务端口
            if [[ "$PORT" != "$NEW_SSH_PORT" && "$PORT" != "$CURRENT_SSH_PORT" && 
                  "$PORT" != "80" && "$PORT" != "443" && 
                  "$PORT" -lt "65535" && "$PORT" -gt "1024" ]]; then
                
                # 尝试找出服务名称
                SERVICE=""
                if command -v lsof &> /dev/null; then
                    SERVICE=$(lsof -i:$PORT -sTCP:LISTEN 2>/dev/null | grep -v "COMMAND" | awk '{print $1}' | head -1)
                fi
                if [ -z "$SERVICE" ]; then
                    SERVICE=$(netstat -tlnp 2>/dev/null | grep ":$PORT" | awk '{print $7}' | cut -d"/" -f2 | head -1)
                fi
                
                if [ -n "$SERVICE" ]; then
                    COMMENT="Service: $SERVICE"
                else
                    COMMENT="Unknown Service"
                fi
                
                log "${GREEN}允许端口 $PORT/tcp ($COMMENT)${NC}"
                ufw allow $PORT/tcp comment "$COMMENT"
            fi
        done
    fi
    
    # 删除询问用户是否手动开放端口的部分
    log "${GREEN}已自动开放SSH、HTTP、HTTPS端口和检测到的活跃服务端口${NC}"
    
    # 启用防火墙
    if ! ufw status | grep -q "Status: active"; then
        log "${YELLOW}启用UFW防火墙...${NC}"
        echo "y" | ufw enable || log "${RED}UFW启用失败${NC}"
    else
        log "${GREEN}UFW防火墙已启用${NC}"
    fi
    
    # 显示防火墙状态
    ufw status numbered | tee -a $LOG_FILE
    log "${GREEN}防火墙设置完成，已使用最小化原则开放端口${NC}"
else
    log "${YELLOW}非Debian系统，请手动配置防火墙${NC}"
fi

# ===========================================
# 7. 时区设置
# ===========================================
log "${BLUE}[7/10] 时区设置开始...${NC}"

# 设置时区
timedatectl set-timezone $TIMEZONE
if [ $? -eq 0 ]; then
    CURRENT_TZ=$(timedatectl show --property=Timezone --value)
    log "${GREEN}时区设置为: $CURRENT_TZ${NC}"
else
    log "${RED}时区设置失败${NC}"
fi

# ===========================================
# 8. 内存优化 - 添加交换空间
# ===========================================
log "${BLUE}[8/10] 内存优化开始...${NC}"

# 获取当前所有交换空间信息
CURRENT_SWAP_TOTAL=$(free -m | awk '/Swap:/ {print $2}')
CURRENT_SWAP_TOTAL=${CURRENT_SWAP_TOTAL:-0}
log "${YELLOW}当前系统交换空间总大小: ${CURRENT_SWAP_TOTAL}MB${NC}"

# 检查是否存在交换空间且大小与设定值接近（允许误差，mkswap头部会占用少量空间，free显示通常比设定值少几MB）
SWAP_DIFF=$(( CURRENT_SWAP_TOTAL - SWAP_SIZE ))
if [ "${SWAP_DIFF#-}" -le 5 ]; then
    log "${GREEN}当前交换空间大小(${CURRENT_SWAP_TOTAL}MB)与设定值一致，无需修改${NC}"
    # 显示交换空间信息
    free -h | tee -a $LOG_FILE
else
    # 交换空间大小不符或不存在，需要重建
    # 关键修复：旧版直接对正在使用的交换执行 swapoff，会在内存紧张时把换出页
    # 强行搬回内存，触发 OOM Killer 杀死 swapoff(exit 137)，进而导致后续
    # rm/dd/mkswap/swapon 连锁失败。
    # 解决方案：先创建并启用一个新的交换文件作为“缓冲”，再关闭旧交换，
    # 全程保证始终有足够的交换空间容纳换出页，避免 OOM。
    if [ "$CURRENT_SWAP_TOTAL" -gt "0" ]; then
        log "${YELLOW}系统已有交换空间但大小不符(${CURRENT_SWAP_TOTAL}MB)，准备安全重建...${NC}"
    else
        log "${YELLOW}系统未配置交换空间，准备创建...${NC}"
    fi

    # 备份fstab
    cp /etc/fstab /etc/fstab.bak
    log "${GREEN}备份了/etc/fstab文件${NC}"

    # 尽量释放可回收内存，降低需要搬回的换出页数量
    sync
    echo 3 > /proc/sys/vm/drop_caches 2>/dev/null || true

    # 工具函数：创建指定大小(MB)的交换文件并启用
    make_swapfile() {
        local path="$1" size="$2"
        rm -f "$path" 2>/dev/null || true
        # 优先使用 fallocate(更快)，失败则退回 dd
        if ! fallocate -l "${size}M" "$path" 2>/dev/null; then
            dd if=/dev/zero of="$path" bs=1M count="$size" status=progress
        fi
        chmod 600 "$path"
        mkswap "$path" >/dev/null
        swapon "$path"
    }

    if swapon --show=NAME --noheadings 2>/dev/null | grep -qx "/swapfile"; then
        # /swapfile 正在使用：先用临时缓冲文件接管换出页，再安全重建 /swapfile
        log "${YELLOW}检测到 /swapfile 正在使用，建立临时缓冲交换以安全重建...${NC}"
        make_swapfile /swapfile.tmp "$SWAP_SIZE"

        # 关闭旧的 /swapfile（换出页会迁移到 /swapfile.tmp，从而避免 OOM）
        log "${YELLOW}关闭旧的 /swapfile...${NC}"
        if swapoff /swapfile; then
            rm -f /swapfile
        else
            log "${RED}关闭旧 /swapfile 失败，已保留临时缓冲交换，请手动检查${NC}"
        fi

        # 关闭其它可能存在的旧交换设备/文件
        for DEVICE in $(swapon --show=NAME --noheadings 2>/dev/null); do
            [ "$DEVICE" = "/swapfile.tmp" ] && continue
            log "${YELLOW}关闭其它交换空间: $DEVICE${NC}"
            if swapoff "$DEVICE" 2>/dev/null; then
                [ -f "$DEVICE" ] && rm -f "$DEVICE"
            fi
        done

        # 创建最终的 /swapfile（临时缓冲里的换出页可迁移到这里）
        log "${GREEN}创建${SWAP_SIZE}MB大小的交换文件 /swapfile...${NC}"
        make_swapfile /swapfile "$SWAP_SIZE"

        # 关闭并删除临时缓冲（其页面会迁移到新的 /swapfile，仍不会 OOM）
        log "${YELLOW}清理临时缓冲交换...${NC}"
        swapoff /swapfile.tmp 2>/dev/null || true
        rm -f /swapfile.tmp
    else
        # /swapfile 未使用：先创建新 /swapfile，再关闭其它旧交换
        log "${GREEN}创建${SWAP_SIZE}MB大小的交换文件 /swapfile...${NC}"
        make_swapfile /swapfile "$SWAP_SIZE"
        for DEVICE in $(swapon --show=NAME --noheadings 2>/dev/null); do
            [ "$DEVICE" = "/swapfile" ] && continue
            log "${YELLOW}关闭其它交换空间: $DEVICE${NC}"
            if swapoff "$DEVICE" 2>/dev/null; then
                [ -f "$DEVICE" ] && rm -f "$DEVICE"
            fi
        done
    fi

    # 从fstab移除旧的 swap 条目(匹配 swap 文件系统类型或 /swapfile 路径)，再写入新条目
    sed -i '/[[:space:]]swap[[:space:]]/d; /swapfile/d' /etc/fstab
    if ! grep -q "/swapfile" /etc/fstab; then
        echo "/swapfile swap swap defaults 0 0" >> /etc/fstab
        log "${GREEN}已添加到fstab，开机将自动挂载${NC}"
    fi

    # 显示交换空间信息
    log "${GREEN}交换分区配置完成，当前内存和交换空间状态:${NC}"
    free -h | tee -a $LOG_FILE
fi

# 调整swappiness参数(控制系统对交换空间的使用倾向)
# 放在 if/else 之外：无论是否重建交换文件，都按用户设置应用 swappiness
echo "vm.swappiness=${SWAPPINESS}" > /etc/sysctl.d/99-swappiness.conf
sysctl -p /etc/sysctl.d/99-swappiness.conf
log "${GREEN}swappiness 已设置为: ${SWAPPINESS}${NC}"

# ===========================================
# 9. Fail2ban安装和配置
# ===========================================
log "${BLUE}[9/10] Fail2ban安装开始...${NC}"

if [ "$OS_TYPE" = "debian" ]; then
    # 安装Fail2ban
    apt_update_safe
    apt install -y fail2ban || log "${RED}安装Fail2ban失败${NC}"

    # 校验是否真正安装成功
    if ! command -v fail2ban-client &> /dev/null && ! dpkg -l fail2ban 2>/dev/null | grep -q "^ii"; then
        log "${RED}Fail2ban未成功安装，跳过其配置（请检查软件源后重试）${NC}"
    else
        systemctl enable fail2ban 2>/dev/null || log "${YELLOW}无法设置Fail2ban开机自启${NC}"

        # 配置Fail2ban
        mkdir -p /etc/fail2ban/jail.d
        if [ -f /etc/fail2ban/jail.conf ]; then
            cp /etc/fail2ban/jail.conf /etc/fail2ban/jail.local
        fi

        # 清理任何现有配置
        rm -rf /etc/fail2ban/jail.d/* 2>/dev/null || true

        # 创建SSH防护配置
        cat <<EOF > /etc/fail2ban/jail.d/sshd.local
[sshd]
enabled = true
mode = normal
port = $NEW_SSH_PORT
logpath = %(sshd_log)s
backend = systemd
maxretry = 5
bantime = 1h
findtime = 10m
EOF

        # 应用配置：已运行则优先 reload(不中断防护、保留已封禁IP)，未运行才 start
        if systemctl is-active --quiet fail2ban; then
            log "${YELLOW}Fail2ban正在运行，重载配置(保留已封禁IP)...${NC}"
            systemctl reload fail2ban 2>/dev/null \
                || systemctl restart fail2ban 2>/dev/null \
                || log "${RED}Fail2ban配置重载失败${NC}"
        else
            log "${YELLOW}启动Fail2ban服务...${NC}"
            systemctl start fail2ban 2>/dev/null \
                || systemctl restart fail2ban 2>/dev/null \
                || log "${RED}Fail2ban服务启动失败${NC}"
        fi

        # 等待服务启动完成
        log "${YELLOW}等待Fail2ban服务完全启动...${NC}"
        sleep 5

        # 检查服务状态
        if systemctl is-active --quiet fail2ban; then
            log "${GREEN}Fail2ban服务已成功启动${NC}"

            # 显示Fail2ban状态（使用错误处理避免脚本终止）
            log "${YELLOW}获取Fail2ban状态信息...${NC}"

            # 尝试获取fail2ban状态，忽略可能的错误
            fail2ban-client status >/dev/null 2>&1 || log "${YELLOW}无法获取fail2ban综合状态，但这不影响功能${NC}"

            # 尝试获取sshd监狱状态
            if fail2ban-client status sshd >/dev/null 2>&1; then
                log "${GREEN}SSH防护已成功配置${NC}"
                # 只有在前面成功的情况下才显示详细信息
                fail2ban-client status sshd
            else
                log "${YELLOW}无法获取SSH监狱状态，这可能是因为服务刚刚启动或配置需要更多时间生效${NC}"
                log "${YELLOW}如果在重启后仍有问题，请检查 /var/log/fail2ban.log${NC}"
            fi

            # 显示服务状态
            systemctl status fail2ban --no-pager || true
        else
            log "${RED}Fail2ban服务启动失败，请检查错误日志${NC}"
            log "${YELLOW}尝试查看Fail2ban日志获取错误详情:${NC}"
            tail -n 20 /var/log/fail2ban.log 2>/dev/null || log "${RED}无法读取Fail2ban日志${NC}"
        fi

        log "${GREEN}Fail2ban安装和配置完成${NC}"
        log "${YELLOW}如果出现临时错误，服务器重启后通常会正常工作${NC}"
    fi
else
    log "${YELLOW}非Debian系统，请手动安装Fail2ban${NC}"
fi

# ===========================================
# 10. BBR加速配置
# ===========================================
log "${BLUE}[10/10] BBR配置开始...${NC}"

# 检查BBR是否已启用
if sysctl net.ipv4.tcp_congestion_control | grep -q "bbr"; then
    log "${GREEN}BBR已经启用${NC}"
else
    log "${YELLOW}配置BBR...${NC}"
    # 添加BBR配置
    echo "net.core.default_qdisc=fq" >> /etc/sysctl.conf
    echo "net.ipv4.tcp_congestion_control=bbr" >> /etc/sysctl.conf
    
    # 应用设置
    sysctl -p
    
    # 验证设置
    if sysctl net.ipv4.tcp_congestion_control | grep -q "bbr"; then
        log "${GREEN}BBR启用成功${NC}"
    else
        log "${RED}BBR启用失败${NC}"
    fi
fi

# 显示可用的拥塞控制算法
log "${GREEN}当前系统支持的TCP拥塞控制算法:${NC}"
sysctl net.ipv4.tcp_available_congestion_control

# 验证模块是否加载
lsmod | grep bbr

# ===========================================
# 完成处理
# ===========================================
END_TIME=$(date +%s)
DURATION=$((END_TIME - START_TIME))
MINUTES=$((DURATION / 60))
SECONDS=$((DURATION % 60))

log_summary ""
log_summary "本次配置信息"
log_summary "======================================================="
log_summary "${GREEN}VPS初始化完成！用时: ${MINUTES}分${SECONDS}秒${NC}"
log_summary "${GREEN}=======================================================${NC}"
log_summary "${YELLOW}重要提示:${NC}"

# 根据用户选择显示相应的提示信息
TIP_COUNT=1

# 如果用户选择了修改SSH端口，显示端口信息
if [[ "$CHANGE_SSH_PORT" =~ ^[Yy]$ ]]; then
    log_summary "${YELLOW}$TIP_COUNT. SSH端口已更改为: ${NEW_SSH_PORT}${NC}"
    TIP_COUNT=$((TIP_COUNT + 1))
fi

# 如果用户选择了修改root密码，显示密码信息
if [[ "$CHANGE_PASSWORD" =~ ^[Yy]$ ]]; then
    log_summary "${YELLOW}$TIP_COUNT. root密码已更改 (长度: ${#NEW_PASSWORD} 字符)${NC}"
    # 将实际密码单独写入摘要文件（不输出到屏幕）
    echo "   实际密码: ${NEW_PASSWORD}" | sed 's/\x1b\[[0-9;]*m//g' >> $SUMMARY_FILE
    TIP_COUNT=$((TIP_COUNT + 1))
fi

# 如果用户选择了修改主机名，显示主机名信息
if [ "$CHANGE_HOSTNAME_FLAG" = true ]; then
    log_summary "${YELLOW}$TIP_COUNT. 主机名已更改为: ${NEW_HOSTNAME}${NC}"
    TIP_COUNT=$((TIP_COUNT + 1))
fi

# 始终显示防火墙和日志文件信息
log_summary "${YELLOW}$TIP_COUNT. 防火墙已启用，只开放了必要端口${NC}"
TIP_COUNT=$((TIP_COUNT + 1))
log_summary "${YELLOW}$TIP_COUNT. 日志文件保存在: ${LOG_FILE}${NC}"

# 如果已启用BBR，显示BBR信息
if sysctl net.ipv4.tcp_congestion_control 2>/dev/null | grep -q "bbr"; then
    TIP_COUNT=$((TIP_COUNT + 1))
    log_summary "${YELLOW}$TIP_COUNT. BBR加速已成功启用${NC}"
fi

# 如果配置了交换空间，显示交换空间信息
if [ "$CURRENT_SWAP_TOTAL" -gt "0" ]; then
    TIP_COUNT=$((TIP_COUNT + 1))
    log_summary "${YELLOW}$TIP_COUNT. 交换空间大小: $(free -m | grep "Swap:" | awk '{print $2}')MB${NC}"
fi

log_summary "${GREEN}=======================================================${NC}"
log_summary "${BLUE}建议您现在重启服务器以应用所有更改${NC}"
log_summary "${GREEN}=======================================================${NC}"
log_summary ""
log_summary "摘要信息已保存到: ${SUMMARY_FILE}"

# 提示用户是否立即重启
while true; do
    read -p "是否立即重启服务器？(y/n, 默认n): " REBOOT_NOW
    # 设置默认值为否
    REBOOT_NOW=${REBOOT_NOW:-n}
    if [[ "$REBOOT_NOW" =~ ^[Yy]$ ]] || [[ "$REBOOT_NOW" =~ ^[Nn]$ ]]; then
        break
    else
        echo -e "${RED}无效的输入，请输入 y 或 n${NC}"
    fi
done

if [[ "$REBOOT_NOW" =~ ^[Yy]$ ]]; then
    log "${GREEN}服务器将在5秒后重启...${NC}"
    sleep 5
    reboot
else
    log "${YELLOW}请稍后手动重启服务器以应用所有更改${NC}"
fi