#!/bin/bash
#======================================================================================#
# 终极通用虚拟内存（Swap）永久增加脚本 - 支持二次追加
# 特性：自动检测/永久生效/安全配置/通用兼容/支持GB输入/支持多次追加
#======================================================================================#

# --------------------------
# 颜色定义
# --------------------------
RED=$(printf '\033[0;31m'); GREEN=$(printf '\033[0;32m'); YELLOW=$(printf '\033[1;33m')
BLUE=$(printf '\033[0;34m'); PURPLE=$(printf '\033[0;35m'); CYAN=$(printf '\033[0;36m')
NC=$(printf '\033[0m')

# --------------------------
# 全局变量 (自动获取)
# --------------------------
OS_FAMILY=""
SWAP_FILE_PATH="/swapfile"
SWAP_SIZE_GB=""
SWAP_SIZE_MB=""
APPEND_MODE=false

#====================================================================================
# 核心函数区
#====================================================================================

# 检查是否以root用户运行
check_root() {
    [[ $EUID -ne 0 ]] && {
        echo -e "${RED}错误：必须使用root用户运行此脚本！${NC}"
        exit 1
    }
}

# 检测操作系统类型
detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS_FAMILY=$ID
        [[ "$ID_LIKE" =~ "rhel" ]] && OS_FAMILY="rhel"
    elif [[ -f /etc/redhat-release ]]; then
        OS_FAMILY="rhel"
    fi

    echo -e "${GREEN}检测到系统家族：${OS_FAMILY}${NC}"
}

# 获取下一个可用的swap文件名
get_next_swapfile() {
    local base_path="/swapfile"
    local next_num=1
    
    # 如果 /swapfile 不存在，直接使用
    if [ ! -f "$base_path" ]; then
        SWAP_FILE_PATH="$base_path"
        return
    fi
    
    # 寻找下一个可用的文件名
    while [ -f "${base_path}${next_num}" ]; do
        ((next_num++))
    done
    
    SWAP_FILE_PATH="${base_path}${next_num}"
}

# 检查当前交换分区状态
check_swap() {
    local swap_info=$(swapon --show=NAME,SIZE,USED --noheadings 2>/dev/null)
    local total_swap=$(free -m | awk '/Swap:/ {print $2}')
    
    echo -e "${BLUE}============================================${NC}"
    echo -e "${BLUE}       当前交换空间状态${NC}"
    echo -e "${BLUE}============================================${NC}"
    
    if [ ! -z "$swap_info" ]; then
        echo -e "${GREEN}已存在的交换空间：${NC}"
        swapon --show
        echo ""
        echo -e "${GREEN}总交换空间：${total_swap}MB ($(echo "scale=2; $total_swap/1024" | bc)GB)${NC}"
        echo ""
        
        # 询问用户是追加还是替换
        echo -e "${CYAN}请选择操作：${NC}"
        echo -e "  ${GREEN}1.${NC} 追加新的交换空间（推荐，保留现有）"
        echo -e "  ${YELLOW}2.${NC} 删除所有交换空间并重新创建"
        echo -e "  ${RED}3.${NC} 取消操作"
        read -p "请输入选项 (1-3): " choice
        
        case "$choice" in
            1)
                APPEND_MODE=true
                get_next_swapfile
                echo -e "${GREEN}将创建新的交换文件: ${SWAP_FILE_PATH}${NC}"
                ;;
            2)
                APPEND_MODE=false
                echo -e "${YELLOW}将删除所有现有交换空间...${NC}"
                remove_all_swap
                SWAP_FILE_PATH="/swapfile"
                ;;
            3)
                echo -e "${GREEN}操作已取消。${NC}"
                exit 0
                ;;
            *)
                echo -e "${RED}无效选项，默认使用追加模式。${NC}"
                APPEND_MODE=true
                get_next_swapfile
                ;;
        esac
    else
        echo -e "${YELLOW}当前没有交换空间。${NC}"
        SWAP_FILE_PATH="/swapfile"
    fi
    
    echo -e "${BLUE}============================================${NC}"
}

# 删除所有交换空间
remove_all_swap() {
    echo -e "${YELLOW}正在停止所有交换空间...${NC}"
    swapoff -a 2>/dev/null || true
    
    echo -e "${YELLOW}正在删除交换文件...${NC}"
    # 删除所有 swapfile
    for f in /swapfile /swapfile[0-9]*; do
        if [ -f "$f" ]; then
            echo -e "  删除: $f"
            rm -f "$f"
        fi
    done
    
    echo -e "${YELLOW}正在清理 /etc/fstab...${NC}"
    # 从 fstab 中删除所有 swapfile 条目
    sed -i '/swapfile/d' /etc/fstab
    
    echo -e "${GREEN}所有交换空间已清理。${NC}"
}

# 获取用户输入的交换空间大小（支持GB和浮点数）
get_swap_size() {
    echo ""
    echo -e "${CYAN}请输入您想增加的交换空间大小 (GB为单位): ${NC}"
    read -p "例如: 2.5 (表示2.5GB) -> " SWAP_SIZE_GB
    
    # 验证输入是否为有效的数字（整数或浮点数）
    if ! [[ "$SWAP_SIZE_GB" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
        echo -e "${RED}错误：请输入一个有效的数字！${NC}"
        get_swap_size
    fi
    
    # 将 GB 转换为 MB，并确保结果为整数
    SWAP_SIZE_MB=$(awk "BEGIN {printf \"%.0f\", ${SWAP_SIZE_GB} * 1024}")
    
    if [ "$SWAP_SIZE_MB" -eq 0 ]; then
        echo -e "${RED}错误：输入值太小，请输入大于0的数值！${NC}"
        get_swap_size
    fi
    
    echo -e "${GREEN}将创建 ${SWAP_SIZE_GB}GB (${SWAP_SIZE_MB}MB) 的交换空间。${NC}"
}

# 创建交换文件
create_swap_file() {
    echo -e "${CYAN}正在创建交换文件 ${SWAP_FILE_PATH}...${NC}"
    
    # 检查文件是否已存在（追加模式下不应该存在）
    if [ -f "$SWAP_FILE_PATH" ]; then
        if [ "$APPEND_MODE" = true ]; then
            echo -e "${RED}错误：文件 ${SWAP_FILE_PATH} 已存在！${NC}"
            exit 1
        else
            echo -e "${YELLOW}警告：交换文件 ${SWAP_FILE_PATH} 已存在，将删除旧文件...${NC}"
            swapoff "$SWAP_FILE_PATH" 2>/dev/null || true
            rm -f "$SWAP_FILE_PATH"
        fi
    fi
    
    # 使用 dd 或 fallocate 创建文件，fallocate 更快
    if command -v fallocate &> /dev/null; then
        fallocate -l "${SWAP_SIZE_MB}M" "$SWAP_FILE_PATH"
        if [ $? -ne 0 ]; then
            echo -e "${RED}fallocate 命令失败，尝试使用 dd...${NC}"
            dd if=/dev/zero of="$SWAP_FILE_PATH" bs=1M count="$SWAP_SIZE_MB" status=progress
        fi
    else
        echo -e "${YELLOW}未找到 fallocate，使用 dd 创建文件...${NC}"
        dd if=/dev/zero of="$SWAP_FILE_PATH" bs=1M count="$SWAP_SIZE_MB" status=progress
    fi
    
    if [ $? -ne 0 ]; then
        echo -e "${RED}错误：交换文件创建失败！${NC}"
        exit 1
    fi
    
    echo -e "${GREEN}交换文件创建成功。${NC}"
}

# 配置并启用交换文件
activate_swap() {
    echo -e "${CYAN}正在配置并启用交换文件...${NC}"
    
    # 设置文件权限，防止未授权访问
    chmod 600 "$SWAP_FILE_PATH"
    
    # 将文件标记为交换空间
    mkswap "$SWAP_FILE_PATH"
    if [ $? -ne 0 ]; then
        echo -e "${RED}错误：mkswap 失败！${NC}"
        exit 1
    fi
    
    # 启用交换空间
    swapon "$SWAP_FILE_PATH"
    if [ $? -ne 0 ]; then
        echo -e "${RED}错误：swapon 失败！${NC}"
        exit 1
    fi
    
    echo -e "${GREEN}交换空间启用成功！${NC}"
}

# 配置 swappiness（内存使用阈值）
configure_swappiness() {
    echo -e "${CYAN}正在配置内存使用阈值...${NC}"
    
    # 获取物理内存大小（MB）
    local total_memory=$(free -m | awk '/Mem:/ {print $2}')
    local memory_gb=$(echo "scale=2; $total_memory/1024" | bc)
    
    echo -e "${BLUE}检测到物理内存：${memory_gb}GB (${total_memory}MB)${NC}"
    
    # 默认设置为 50（表示内存使用 50% 时开始使用 swap）
    # swappiness 范围是 0-100
    # 0 = 尽量不使用 swap
    # 100 = 积极使用 swap
    # 50 = 内存使用 50% 时开始使用 swap
    local swappiness_value=50
    
    echo -e "${CYAN}请选择内存使用阈值：${NC}"
    echo -e "  ${GREEN}1.${NC} 保守模式 (30%) - 内存使用 70% 时开始使用 swap"
    echo -e "  ${GREEN}2.${NC} 平衡模式 (50%) - 内存使用 50% 时开始使用 swap [推荐]"
    echo -e "  ${GREEN}3.${NC} 积极模式 (70%) - 内存使用 30% 时开始使用 swap"
    echo -e "  ${GREEN}4.${NC} 自定义数值"
    read -p "请输入选项 (1-4) [默认: 2]: " swappiness_choice
    
    case "$swappiness_choice" in
        1)
            swappiness_value=30
            echo -e "${GREEN}已选择保守模式 (swappiness=30)${NC}"
            ;;
        2|"")
            swappiness_value=50
            echo -e "${GREEN}已选择平衡模式 (swappiness=50)${NC}"
            ;;
        3)
            swappiness_value=70
            echo -e "${GREEN}已选择积极模式 (swappiness=70)${NC}"
            ;;
        4)
            read -p "请输入 swappiness 值 (0-100): " custom_value
            if [[ "$custom_value" =~ ^[0-9]+$ ]] && [ "$custom_value" -ge 0 ] && [ "$custom_value" -le 100 ]; then
                swappiness_value=$custom_value
                echo -e "${GREEN}已设置 swappiness=${swappiness_value}${NC}"
            else
                echo -e "${YELLOW}输入无效，使用默认值 50${NC}"
                swappiness_value=50
            fi
            ;;
        *)
            echo -e "${YELLOW}无效选项，使用默认平衡模式 (swappiness=50)${NC}"
            swappiness_value=50
            ;;
    esac
    
    # 立即生效
    sysctl -w vm.swappiness=$swappiness_value
    
    # 永久生效
    if grep -q "^vm.swappiness" /etc/sysctl.conf; then
        sed -i "s/^vm.swappiness.*/vm.swappiness=$swappiness_value/" /etc/sysctl.conf
    else
        echo "vm.swappiness=$swappiness_value" >> /etc/sysctl.conf
    fi
    
    # 创建独立的配置文件
    cat > /etc/sysctl.d/99-swappiness.conf << EOF
# 虚拟内存使用阈值配置
# 值越低，越倾向于使用物理内存
# 值越高，越倾向于使用虚拟内存
# 50 表示当物理内存使用 50% 时开始使用 swap
vm.swappiness=$swappiness_value
EOF
    
    echo -e "${GREEN}swappiness 已设置为 ${swappiness_value}${NC}"
    echo -e "${YELLOW}说明：${NC}"
    if [ "$swappiness_value" -le 30 ]; then
        echo -e "  - 系统会尽量使用物理内存"
        echo -e "  - 只有当物理内存不足时才会使用 swap"
    elif [ "$swappiness_value" -ge 70 ]; then
        echo -e "  - 系统会积极使用 swap"
        echo -e "  - 可以释放更多物理内存给缓存使用"
    else
        echo -e "  - 平衡使用物理内存和 swap"
        echo -e "  - 当内存使用 ${swappiness_value}% 时开始使用 swap"
    fi
}

# 添加到 /etc/fstab 以实现永久性
make_permanent() {
    echo -e "${CYAN}正在将交换文件配置为永久生效...${NC}"
    
    # 检查 fstab 中是否已存在相同路径的条目，避免重复添加
    if ! grep -q "^${SWAP_FILE_PATH} " /etc/fstab; then
        echo "$SWAP_FILE_PATH none swap sw 0 0" >> /etc/fstab
        echo -e "${GREEN}已添加到 /etc/fstab${NC}"
    else
        echo -e "${YELLOW}警告：${SWAP_FILE_PATH} 已在 /etc/fstab 中${NC}"
    fi
    
    echo -e "${GREEN}配置已保存，系统重启后仍将自动加载。${NC}"
}

#====================================================================================
# 主执行流程
#====================================================================================
main() {
    check_root
    detect_os
    check_swap
    get_swap_size
    create_swap_file
    activate_swap
    make_permanent
    configure_swappiness
    
    echo -e "\n${PURPLE}============================================${NC}"
    echo -e "${BLUE}       虚拟内存配置完成！${NC}"
    echo -e "${PURPLE}============================================${NC}"
    echo -e "${GREEN}当前所有交换空间状态：${NC}"
    swapon --show
    echo ""
    echo -e "${GREEN}物理内存：$(free -h | awk '/Mem:/ {print $2}')${NC}"
    echo -e "${GREEN}总交换空间：$(free -h | awk '/Swap:/ {print $2}')${NC}"
    echo -e "${GREEN}交换空间已使用：$(free -h | awk '/Swap:/ {print $3}')${NC}"
    echo -e "${GREEN}当前 swappiness：$(cat /proc/sys/vm/swappiness)${NC}"
    echo -e "${PURPLE}============================================${NC}"
    
    if [ "$APPEND_MODE" = true ]; then
        echo -e "${YELLOW}提示：您可以使用此脚本多次追加交换空间。${NC}"
    fi
}

main