#!/bin/bash
set -e
PVE_WORKERS=6          # 想改数字改这里
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'

echo -e "${YELLOW}=== 一键宿主机内存优化 ===${NC}"

# -------- 参数询问 --------
read -rp "磁盘 swap 大小（GB）[默认 8]: " disk_gb
disk_gb=${disk_gb:-8}
read -rp "是否同时开启 zram（50% RAM，上限 64 GB）？[Y/n]: " zram_yn
zram_yn=${zram_yn:-Y}

# -------- 计算预期效果 --------
total_kb=$(grep MemTotal /proc/meminfo | awk '{print $2}')
zram_kb=$(( total_kb / 2 ))
[[ $zram_kb -gt 67108864 ]] && zram_kb=67108864
zram_gb=$(( zram_kb / 1024 / 1024 ))
total_swap_gb=$(( zram_gb + disk_gb ))
total_h=$(free -h | awk '/Mem:/{print $2}')

# -------- 完整预期输出 --------
echo -e "${GREEN}━━━━━━━━━  预 期  效  果  ━━━━━━━━━${NC}"
echo -e "  物理内存 : $total_h"
[[ "$zram_yn" =~ [Yy] ]] && echo -e "  zram 容量 : ${zram_gb} GB  （50 % RAM，上限 64 GB）"
echo -e "  磁盘 swap : ${disk_gb} GB"
[[ "$zram_yn" =~ [Yy] ]] && echo -e "  总 swap   : ${total_swap_gb} GB"
echo -e "  KSM 状态  : 即将开启"
echo -e "  PVE 代理  : workers = ${PVE_WORKERS} （即将锁定）"
echo -e "${GREEN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo ""

# -------- 确认执行 --------
read -rp "确认执行？输入 y 开始，其他退出: " confirm
[[ "$confirm" != "y" ]] && echo "已取消。" && exit 0

# -------- 1. 开启 KSM --------
echo -e "${YELLOW}>>> 1. 开启 KSM${NC}"
echo 1 > /sys/kernel/mm/ksm/run
echo 1000 > /sys/kernel/mm/ksm/pages_to_scan
echo 20   > /sys/kernel/mm/ksm/sleep_millisecs
grep -q "^echo 1 > /sys/kernel/mm/ksm/run" /etc/rc.local 2>/dev/null || \
  (echo "echo 1 > /sys/kernel/mm/ksm/run" >> /etc/rc.local && chmod +x /etc/rc.local)

# -------- 2. 锁定 PVE 代理 workers --------
echo -e "${YELLOW}>>> 2. 锁定 PVE 代理 workers = ${PVE_WORKERS}${NC}"
grep -q ^PVEPROXY_WORKERS /etc/default/pveproxy 2>/dev/null && \
  sed -i "s/^PVEPROXY_WORKERS=.*/PVEPROXY_WORKERS=${PVE_WORKERS}/" /etc/default/pveproxy || \
  echo "PVEPROXY_WORKERS=${PVE_WORKERS}" >> /etc/default/pveproxy
systemctl restart pveproxy

# -------- 3. 磁盘 swap --------
echo -e "${YELLOW}>>> 3. 挂载 ${disk_gb} GB 磁盘 swap${NC}"
disk_swap=$(lsblk -n -b -o NAME,SIZE | awk -v db=$disk_gb '$2>=db*1024*1024*1024 && $2<=db*1.1*1024*1024*1024 {print "/dev/"$1}' | head -n1)
[[ -z "$disk_swap" ]] && disk_swap=/dev/sda5
swapoff -a 2>/dev/null || true
swapon -p 10 "$disk_swap"

# -------- 4. 可选 zram --------
if [[ "$zram_yn" =~ [Yy] ]]; then
    echo -e "${YELLOW}>>> 4. 创建 zram（${zram_gb} GB）${NC}"
    rmmod zram 2>/dev/null || true
    modprobe zram num_devices=1
    echo $(( zram_kb * 1024 )) > /sys/block/zram0/disksize
    echo lz4 > /sys/block/zram0/comp_algorithm
    mkswap -L zram0 /dev/zram0
    swapon -p 100 /dev/zram0

    # 开机服务
    cat >/etc/systemd/system/zram-swap.service <<EOF
[Unit]
Description=Auto zram (50% RAM max 64G) + ${disk_gb}G disk swap
After=local-fs.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/local/sbin/zram-swap-start
ExecStop=/usr/local/sbin/zram-swap-stop

[Install]
WantedBy=multi-user.target
EOF

    cat >/usr/local/sbin/zram-swap-start <<EOF
#!/bin/bash
swapoff -a
rmmod zram 2>/dev/null || true
modprobe zram num_devices=1
echo $(( zram_kb * 1024 )) > /sys/block/zram0/disksize
echo lz4 > /sys/block/zram0/comp_algorithm
mkswap -L zram0 /dev/zram0
swapon -p 100 /dev/zram0
swapon -p 10 $disk_swap
EOF

    cat >/usr/local/sbin/zram-swap-stop <<EOF
#!/bin/bash
swapoff /dev/zram0 2>/dev/null || true
echo 1 > /sys/block/zram0/reset 2>/dev/null || true
swapoff $disk_swap
EOF
    chmod +x /usr/local/sbin/zram-swap-{start,stop}
    systemctl daemon-reload
    systemctl enable --now zram-swap.service
fi

# -------- 结果 --------
echo -e "${GREEN}>>> 完成！当前内存/交换概览：${NC}"
free -h
swapon -s
[[ "$zram_yn" =~ [Yy] ]] && zramctl