#!/bin/bash

# ================= 核心配置区域 =================
# 1. 设置版本号 (修改此处即可控制下载版本)
VERSION="1.1.15"

# 2. 官方标准数据目录 (不要改动，这是 RustDesk 的默认位置)
DATA_DIR="/var/lib/rustdesk-server"

# 3. 下载链接 (自动适配上方版本号)
# 主源：你的私有源
URL_HBBS_PRIMARY="https://cloudreve.vps3344521.xyz/f/W9tx/rustdesk-server-hbbs_${VERSION}_amd64.deb"
URL_HBBR_PRIMARY="https://cloudreve.vps3344521.xyz/f/j3fJ/rustdesk-server-hbbr_${VERSION}_amd64.deb"
# 备用源：GitHub (防止私有源挂掉)
URL_HBBS_FALLBACK="https://github.com/rustdesk/rustdesk-server/releases/download/${VERSION}/rustdesk-server-hbbs_${VERSION}_amd64.deb"
URL_HBBR_FALLBACK="https://github.com/rustdesk/rustdesk-server/releases/download/${VERSION}/rustdesk-server-hbbr_${VERSION}_amd64.deb"
# ===============================================

# --- 功能: 检测安装并询问是否卸载重装 ---
echo "正在检查系统状态..."

if systemctl list-unit-files | grep -q "rustdesk-hbbs.service"; then
    echo "------------------------------------------------"
    echo "⚠️  检测到系统已安装 RustDesk Server！"
    echo "------------------------------------------------"
    read -p "是否卸载当前版本并重新安装？(输入 y 确认重装，输入 n 退出): " choice
    case "$choice" in 
      y|Y ) 
        echo ">>> 正在停止服务..."
        systemctl stop rustdesk-hbbs rustdesk-hbbr
        systemctl disable rustdesk-hbbs rustdesk-hbbr
        
        echo ">>> 正在卸载程序..."
        dpkg -P rustdesk-server-hbbs rustdesk-server-hbbr
        apt-get purge -y rustdesk-server-hbbs rustdesk-server-hbbr
        
        # 询问是否清空数据
        echo "------------------------------------------------"
        echo "注意：官方数据目录为 $DATA_DIR"
        read -p "是否清空旧的 Key 和数据库？(y/n) [推荐 y 以防冲突]: " del_data
        if [[ "$del_data" == "y" || "$del_data" == "Y" ]]; then
            rm -rf "$DATA_DIR"
            echo ">>> 旧数据已清除。"
        else
            echo ">>> 旧数据已保留。"
        fi
        echo ">>> 卸载完成，准备重新安装..."
        sleep 2
        ;;
      * ) 
        echo ">>> 操作已取消，脚本退出。"
        exit 0
        ;;
    esac
fi

# ================== 安装流程 ==================

# 1. 基础环境修复与依赖安装
echo ">>> [1/6] 正在配置基础环境..."
apt-get update -q
# 增加 ufw 用于防火墙管理，dnsutils 用于网络查询
apt-get install -y curl wget lsof ufw dnsutils

# 2. 智能获取公网 IP (双重保险)
echo ">>> [2/6] 正在获取公网 IP..."
HOST_IP=$(curl -s4 --connect-timeout 5 ifconfig.me)
if [ -z "$HOST_IP" ]; then
    echo "    备用方案启动：尝试通过 ip.sb 获取..."
    HOST_IP=$(curl -s4 --connect-timeout 5 ip.sb)
fi

if [ -z "$HOST_IP" ]; then
    echo "❌ 错误：无法获取公网 IP，请检查网络连接。"
    exit 1
fi
echo "    获取成功: $HOST_IP"

# 3. 下载安装包 (自动重试)
echo ">>> [3/6] 正在下载安装包 (版本: $VERSION)..."

# 下载函数
download_pkg() {
    local name=$1
    local url_pri=$2
    local url_sec=$3
    
    echo "    正在下载 $name..."
    wget -O "$name.deb" "$url_pri"
    
    # 检查文件是否有效（大于 1KB）
    if [ ! -s "$name.deb" ] || [ $(stat -c%s "$name.deb") -lt 1000 ]; then
        echo "    ⚠️ 主源下载失败，切换备用源..."
        rm -f "$name.deb"
        wget -O "$name.deb" "$url_sec"
    fi
    
    if [ ! -s "$name.deb" ]; then
        echo "❌ 错误：$name 下载彻底失败，请检查链接。"
        exit 1
    fi
}

download_pkg "hbbs" "$URL_HBBS_PRIMARY" "$URL_HBBS_FALLBACK"
download_pkg "hbbr" "$URL_HBBR_PRIMARY" "$URL_HBBR_FALLBACK"

# 4. 安装
echo ">>> [4/6] 正在安装软件包..."
dpkg -i hbbs.deb hbbr.deb
apt-get install -f -y
# 自动清理安装包
rm -f hbbs.deb hbbr.deb

# 5. 服务配置 (锁定官方路径，但修正启动参数)
echo ">>> [5/6] 正在配置系统服务..."

# 确保目录存在且权限正确
mkdir -p "$DATA_DIR"
chown -R root:root "$DATA_DIR"

# 写入服务文件 (覆盖到 /etc，优先级高于默认安装的 /lib)
# 使用 Standard Output 和 Error 重定向，方便排错
cat > /etc/systemd/system/rustdesk-hbbs.service <<EOF
[Unit]
Description=RustDesk ID Server
After=network.target

[Service]
Type=simple
LimitNOFILE=1000000
ExecStart=/usr/bin/hbbs -r ${HOST_IP}:21117 -k _
WorkingDirectory=${DATA_DIR}
Restart=always
User=root
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF

cat > /etc/systemd/system/rustdesk-hbbr.service <<EOF
[Unit]
Description=RustDesk Relay Server
After=network.target

[Service]
Type=simple
LimitNOFILE=1000000
ExecStart=/usr/bin/hbbr -k _
WorkingDirectory=${DATA_DIR}
Restart=always
User=root
RestartSec=10

[Install]
WantedBy=multi-user.target
EOF

# 6. 防火墙配置 (新增优化)
echo ">>> [6/6] 正在优化防火墙端口..."
if command -v ufw > /dev/null; then
    ufw allow 21115:21119/tcp
    ufw allow 21116/udp
    echo "    UFW 防火墙规则已添加。"
else
    echo "    未检测到 UFW，跳过防火墙配置。"
fi

# 启动服务
systemctl daemon-reload
systemctl enable --now rustdesk-hbbs rustdesk-hbbr

# ================= 结果展示与 Key 等待 =================
echo "------------------------------------------------"
echo "✅ RustDesk 安装并启动完成！"
echo "   安装位置: 官方默认 (/usr/bin)"
echo "   数据位置: $DATA_DIR"
echo "------------------------------------------------"
echo "⏳ 正在等待 Key 文件生成 (最多等待 30 秒)..."

PUB_KEY="${DATA_DIR}/id_ed25519.pub"
MAX_RETRIES=30
COUNT=0

# 循环检测 Key 是否生成
while [ ! -f "$PUB_KEY" ]; do
    sleep 1
    ((COUNT++))
    if [ "$COUNT" -ge "$MAX_RETRIES" ]; then
        echo "⚠️ 等待超时！Key 文件尚未生成。"
        echo "请稍后手动查看: cat $PUB_KEY"
        exit 0
    fi
    echo -ne "."
done
echo ""

# 成功获取
echo "------------------------------------------------"
echo "🎉 你的 Key (公钥) 为:"
echo ""
echo -e "\033[32m$(cat "$PUB_KEY")\033[0m"
echo ""
echo "------------------------------------------------"