#!/bin/bash

# ================= 配置区域 =================
VERSION="1.1.15"
DATA_DIR="/var/lib/rustdesk-server"
BACKUP_DIR="/tmp/rustdesk_backup_$(date +%s)"
# ===========================================

echo "========================================================"
echo "   RustDesk Server 全能修复脚本 (安装/转移/修复/启动)"
echo "========================================================"

# --- 1. 智能查找并备份/转移现有的密钥 ---
echo ">>> [1/7] 正在扫描并搜集现有的密钥文件..."
mkdir -p "$BACKUP_DIR"

# 定义查找函数
save_files() {
    local src_dir=$1
    if [ -f "${src_dir}/id_ed25519.pub" ]; then
        echo "    发现密钥在: ${src_dir}"
        cp "${src_dir}/id_ed25519" "$BACKUP_DIR/" 2>/dev/null
        cp "${src_dir}/id_ed25519.pub" "$BACKUP_DIR/" 2>/dev/null
        cp "${src_dir}/db_v2.sqlite3" "$BACKUP_DIR/" 2>/dev/null
        # 标记为已找到
        FOUND_KEYS=1
    fi
}

# A. 检查官方目录 (是否有残留)
save_files "$DATA_DIR"
# B. 检查 Root 目录 (你刚才手动运行产生的地方)
save_files "/root"
# C. 检查当前目录
save_files "$(pwd)"

if [ "$FOUND_KEYS" == "1" ]; then
    echo "✅ 已成功提取密钥和数据，稍后会自动归位。"
else
    echo "⚠️ 未扫描到现有密钥，安装后将生成新的。"
fi

# --- 2. 彻底卸载与清理 ---
echo ">>> [2/7] 正在清理旧环境..."
systemctl stop rustdesk-hbbs rustdesk-hbbr 2>/dev/null
systemctl disable rustdesk-hbbs rustdesk-hbbr 2>/dev/null
dpkg -P rustdesk-server-hbbs rustdesk-server-hbbr 2>/dev/null
rm -rf "$DATA_DIR"  # 已经备份过了，这里彻底清空保证环境纯净

# --- 3. 环境准备与 IP 获取 ---
echo ">>> [3/7] 正在准备基础环境..."
apt-get update -q
apt-get install -y curl wget lsof ufw dnsutils

echo "    正在获取公网 IP..."
HOST_IP=$(curl -s4 --connect-timeout 5 ifconfig.me)
[ -z "$HOST_IP" ] && HOST_IP=$(curl -s4 --connect-timeout 5 ip.sb)

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

# --- 4. 下载与安装 ---
echo ">>> [4/7] 正在下载并安装..."
URL_HBBS="https://cloudreve.vps3344521.xyz/f/W9tx/rustdesk-server-hbbs_${VERSION}_amd64.deb"
URL_HBBR="https://cloudreve.vps3344521.xyz/f/j3fJ/rustdesk-server-hbbr_${VERSION}_amd64.deb"

# 下载 (带重试逻辑)
wget -O hbbs.deb "$URL_HBBS" || wget -O hbbs.deb "https://github.com/rustdesk/rustdesk-server/releases/download/${VERSION}/rustdesk-server-hbbs_${VERSION}_amd64.deb"
wget -O hbbr.deb "$URL_HBBR" || wget -O hbbr.deb "https://github.com/rustdesk/rustdesk-server/releases/download/${VERSION}/rustdesk-server-hbbr_${VERSION}_amd64.deb"

dpkg -i hbbs.deb hbbr.deb
apt-get install -f -y
rm -f hbbs.deb hbbr.deb

# --- 5. 密钥归位与权限修复 (核心步骤) ---
echo ">>> [5/7] 正在将密钥转移至官方目录..."
mkdir -p "$DATA_DIR"

if [ "$FOUND_KEYS" == "1" ]; then
    cp -f "$BACKUP_DIR/"* "$DATA_DIR/"
    echo "✅ 密钥已成功迁移回 $DATA_DIR"
    # 清理刚才手动在 /root 生成的垃圾文件，保持系统整洁
    rm -f /root/id_ed25519 /root/id_ed25519.pub /root/db_v2.sqlite3
else
    echo "    无旧密钥，系统将自动生成。"
fi

# 统一修复权限 (非常重要，否则服务起不来)
chown -R root:root "$DATA_DIR"
chmod 755 "$DATA_DIR"

# --- 6. 配置服务 (锁定官方路径) ---
echo ">>> [6/7] 正在配置系统服务..."
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

# 放行端口
if command -v ufw > /dev/null; then
    ufw allow 21115:21119/tcp
    ufw allow 21116/udp
fi

# --- 7. 启动与验证 ---
echo ">>> [7/7] 正在启动服务..."
systemctl daemon-reload
systemctl enable --now rustdesk-hbbs rustdesk-hbbr

echo "------------------------------------------------"
echo "✅ 修复完成！密钥已锁定在官方目录: $DATA_DIR"
echo "⏳ 正在读取密钥..."

PUB_KEY="${DATA_DIR}/id_ed25519.pub"
# 循环等待确保文件生成/存在
for i in {1..10}; do
    if [ -f "$PUB_KEY" ]; then
        break
    fi
    sleep 1
done

if [ -f "$PUB_KEY" ]; then
    echo "------------------------------------------------"
    echo "🎉 你的 Key (公钥) 为:"
    echo ""
    echo -e "\033[32m$(cat "$PUB_KEY")\033[0m"
    echo ""
    echo "------------------------------------------------"
else
    echo "❌ 警告：服务似乎启动失败，未检测到 Key 文件。"
    echo "请运行 journalctl -u rustdesk-hbbs -n 20 查看错误日志。"
fi