#!/bin/bash
#====================================================
# Swap Builder - Swap 管理工具
# 夏日之瓜×仰望星辰工作室 出品
#====================================================

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

# 检查 root 权限
check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo -e "${RED}[错误] 此脚本需要 root 权限运行。${NC}"
        echo -e "${YELLOW}请使用：curl -sSL swap.fangqihang.cn | sudo bash${NC}"
        exit 1
    fi
}

# 检查交互环境
check_tty() {
    if [ ! -e /dev/tty ]; then
        echo -e "${RED}[错误] 无法访问 /dev/tty，请在终端中运行此脚本。${NC}"
        exit 1
    fi
}

# 打印横幅
print_banner() {
    echo -e "${CYAN}"
    echo "=================================================="
    echo "          Swap Builder - Swap 管理工具"
    echo "=================================================="
    echo "        夏日之瓜×仰望星辰工作室 出品"
    echo "=================================================="
    echo -e "${NC}"
}

# 显示当前状态
show_status() {
    echo -e "${BLUE}---------- 当前系统状态 ----------${NC}"
    echo ""

    # Swap 设备信息
    if command -v swapon &>/dev/null; then
        local swap_info=$(swapon --show 2>/dev/null)
        if [ -n "$swap_info" ]; then
            echo -e "${GREEN}[Swap 设备]${NC}"
            echo "$swap_info"
        else
            echo -e "${YELLOW}[Swap 设备] 当前未启用 Swap${NC}"
        fi
    fi

    echo ""

    # 内存信息
    if command -v free &>/dev/null; then
        echo -e "${GREEN}[内存使用]${NC}"
        free -h 2>/dev/null || free -m 2>/dev/null
    fi

    echo ""

    # swappiness
    local current_swappiness=$(cat /proc/sys/vm/swappiness 2>/dev/null)
    if [ -n "$current_swappiness" ]; then
        echo -e "${GREEN}[swappiness] 当前值：${current_swappiness}${NC}"
    fi

    echo ""
    echo -e "${BLUE}-----------------------------------${NC}"
    echo ""
}

# 设置 Swap
set_swap() {
    echo -e "${CYAN}【设置 Swap】${NC}"
    echo ""
    echo -e "请输入想要设置的 Swap 大小（单位：MB）"
    echo -e "${YELLOW}建议：物理内存的 1~2 倍，例如 2048 = 2GB${NC}"
    echo -e "${YELLOW}输入 0 可关闭 Swap${NC}"
    echo ""
    echo -n "Swap 大小 (MB): "
    read swap_size < /dev/tty
    echo ""

    # 输入校验
    if ! [[ "$swap_size" =~ ^[0-9]+$ ]]; then
        echo -e "${RED}[错误] 请输入有效的数字！${NC}"
        return 1
    fi

    # 关闭 Swap
    if [ "$swap_size" -eq 0 ]; then
        echo -e "${YELLOW}正在关闭 Swap...${NC}"
        swapoff -a 2>/dev/null

        # 从 fstab 中移除 swap 条目（备份后操作）
        if grep -q ' swap ' /etc/fstab 2>/dev/null; then
            cp /etc/fstab /etc/fstab.bak.$(date +%Y%m%d%H%M%S)
            sed -i '/ swap /d' /etc/fstab
            echo -e "${GREEN}已备份 /etc/fstab 并移除 Swap 条目${NC}"
        fi

        echo -e "${GREEN}Swap 已成功关闭！${NC}"
        return 0
    fi

    # 范围检查
    if [ "$swap_size" -lt 128 ]; then
        echo -e "${RED}[错误] Swap 大小不能小于 128MB！${NC}"
        return 1
    fi

    if [ "$swap_size" -gt 65536 ]; then
        echo -e "${RED}[错误] Swap 大小不能超过 65536MB (64GB)！${NC}"
        return 1
    fi

    # 磁盘空间检查
    local available_mb=$(df -m / 2>/dev/null | awk 'NR==2 {print $4}')
    local required_mb=$((swap_size + 1024))

    if [ -z "$available_mb" ]; then
        echo -e "${YELLOW}[警告] 无法检测磁盘空间，继续操作...${NC}"
    elif [ "$available_mb" -lt "$required_mb" ]; then
        echo -e "${RED}[错误] 磁盘空间不足！${NC}"
        echo -e "${YELLOW}可用：${available_mb}MB | 需要：${required_mb}MB（含 1GB 余量）${NC}"
        return 1
    fi

    # 确认操作
    echo -e "${YELLOW}即将创建 ${swap_size}MB 的 Swap 文件${NC}"
    echo -n "确认继续？(y/n): "
    read confirm < /dev/tty
    echo ""

    if [[ "$confirm" != "y" && "$confirm" != "Y" ]]; then
        echo -e "${YELLOW}操作已取消。${NC}"
        return 0
    fi

    # 备份 fstab
    cp /etc/fstab /etc/fstab.bak.$(date +%Y%m%d%H%M%S) 2>/dev/null
    echo -e "${GREEN}已备份 /etc/fstab${NC}"

    # 关闭现有 swap
    swapoff -a 2>/dev/null

    # 确定 swap 文件路径
    local swap_file="/swapfile"
    if [ -d "/www" ]; then
        swap_file="/www/swap"
    fi

    # 移除旧文件（如果存在）
    if [ -f "$swap_file" ]; then
        rm -f "$swap_file"
    fi

    # 创建 swap 文件
    echo -e "${YELLOW}正在创建 Swap 文件（${swap_size}MB）...${NC}"

    if command -v fallocate &>/dev/null; then
        fallocate -l ${swap_size}M "$swap_file" 2>/dev/null
        if [ $? -ne 0 ]; then
            echo -e "${YELLOW}fallocate 不可用，使用 dd 创建...${NC}"
            dd if=/dev/zero of="$swap_file" bs=1M count=$swap_size
        fi
    else
        dd if=/dev/zero of="$swap_file" bs=1M count=$swap_size
    fi

    # 检查文件是否创建成功
    if [ ! -f "$swap_file" ]; then
        echo -e "${RED}[错误] Swap 文件创建失败！${NC}"
        return 1
    fi

    # 设置权限
    chmod 600 "$swap_file"

    # 创建 swap 并启用
    mkswap "$swap_file"
    swapon "$swap_file"

    if [ $? -ne 0 ]; then
        echo -e "${RED}[错误] Swap 启用失败！${NC}"
        return 1
    fi

    # 持久化到 fstab（先移除旧条目再添加）
    sed -i '/ swap /d' /etc/fstab
    echo "$swap_file none swap sw 0 0" >> /etc/fstab

    echo ""
    echo -e "${GREEN}Swap 设置成功！${NC}"
    echo ""
    swapon --show 2>/dev/null
    echo ""
    free -h 2>/dev/null || free -m 2>/dev/null
}

# 设置 swappiness
set_swappiness() {
    echo -e "${CYAN}【设置 swappiness】${NC}"
    echo ""
    echo -e "swappiness 控制系统使用 Swap 的积极程度"
    echo -e "  ${YELLOW}0${NC}  = 尽量不使用 Swap"
    echo -e "  ${YELLOW}100${NC} = 积极使用 Swap"
    echo -e "  ${YELLOW}建议${NC}：服务器推荐 10~30"
    echo ""

    local current=$(cat /proc/sys/vm/swappiness 2>/dev/null)
    echo -e "${GREEN}当前 swappiness 值：${current}${NC}"
    echo ""
    echo -n "请输入 swappiness 值 (0-100): "
    read swappiness_val < /dev/tty
    echo ""

    # 输入校验
    if ! [[ "$swappiness_val" =~ ^[0-9]+$ ]]; then
        echo -e "${RED}[错误] 请输入有效的数字！${NC}"
        return 1
    fi

    if [ "$swappiness_val" -gt 100 ]; then
        echo -e "${RED}[错误] swappiness 值必须在 0~100 之间！${NC}"
        return 1
    fi

    # 临时生效
    sysctl -w vm.swappiness=$swappiness_val

    # 持久化
    if grep -q "^vm.swappiness" /etc/sysctl.conf 2>/dev/null; then
        sed -i "s/^vm.swappiness=.*/vm.swappiness=$swappiness_val/" /etc/sysctl.conf
    else
        echo "vm.swappiness=$swappiness_val" >> /etc/sysctl.conf
    fi

    echo ""
    echo -e "${GREEN}swappiness 已设置为 ${swappiness_val}（已持久化到 /etc/sysctl.conf）${NC}"
}

# 主菜单
main_menu() {
    while true; do
        clear 2>/dev/null
        print_banner
        show_status

        echo -e "${CYAN}请选择操作：${NC}"
        echo -e "  ${GREEN}1${NC}) 设置 Swap"
        echo -e "  ${GREEN}2${NC}) 设置优先度（swappiness）"
        echo -e "  ${GREEN}0${NC}) 退出"
        echo ""
        echo -n "请输入选项 [0-2]: "
        read choice < /dev/tty
        echo ""

        case $choice in
            1)
                set_swap
                ;;
            2)
                set_swappiness
                ;;
            0)
                echo -e "${GREEN}感谢使用 Swap Builder！${NC}"
                echo -e "${CYAN}—— 夏日之瓜×仰望星辰工作室 出品 ——${NC}"
                exit 0
                ;;
            *)
                echo -e "${RED}无效选项，请重新输入！${NC}"
                ;;
        esac

        echo ""
        echo -n "按回车键返回主菜单..."
        read < /dev/tty
    done
}

# 捕获 Ctrl+C
trap 'echo ""; echo -e "${YELLOW}操作已中断。${NC}"; exit 1' INT

# 入口
check_root
check_tty
main_menu