#!/bin/bash

# 颜色定义
GREEN="\033[32m"
RED="\033[31m"
YELLOW="\033[33m"
PLAIN="\033[0m"

echo -e "${GREEN}=============================================${PLAIN}"
echo -e "${GREEN}      RustDesk Server (S6版) 一键部署脚本      ${PLAIN}"
echo -e "${GREEN}=============================================${PLAIN}"

# 1. 检查是否为 Root 用户
if [[ $EUID -ne 0 ]]; then
    echo -e "${RED}错误：请使用 root 用户运行此脚本！${PLAIN}"
    exit 1
fi

# 2. 检查 Docker 是否安装
if ! command -v docker &> /dev/null; then
    echo -e "${RED}未检测到 Docker，请先安装 Docker 和 Docker Compose！${PLAIN}"
    echo -e "你可以尝试运行：curl -fsSL https://get.docker.com | bash"
    exit 1
fi

# 3. 设置安装目录（恢复为原始默认路径）
INSTALL_DIR="/data/rustdesk"
echo -e "${YELLOW}当前安装目录: ${INSTALL_DIR}${PLAIN}"

# 创建目录
mkdir -p "${INSTALL_DIR}/data"
mkdir -p "${INSTALL_DIR}/api"

# 4. 获取用户输入
read -p "请输入服务器的公网 IP 或解析好的域名 (必填): " HOST_IP
if [[ -z "$HOST_IP" ]]; then
    echo -e "${RED}错误：必须输入 IP 或域名！${PLAIN}"
    exit 1
fi

# --- 交互选择：是否需要账号密码登录 ---
echo -e "${YELLOW}是否开启强制登录模式 (MUST_LOGIN)？${PLAIN}"
echo -e "1) 需要 (必须注册并登录账号才能连接)"
echo -e "2) 不需要 (直接使用 ID 和 Key 即可连接 - 默认)"
read -p "请选择 [1/2]: " LOGIN_CHOICE

if [[ "$LOGIN_CHOICE" == "1" ]]; then
    MUST_LOGIN_VAL="Y"
    echo -e "${GREEN}配置已设为：必须登录${PLAIN}"
else
    MUST_LOGIN_VAL="N"
    echo -e "${GREEN}配置已设为：无需登录（公共模式）${PLAIN}"
fi
# --------------------------------------

# 进入目录
cd "$INSTALL_DIR" || exit

# 5. 生成配置文件 (以容器名命名)
echo -e "${YELLOW}正在生成配置文件 rustdesk.yml...${PLAIN}"

cat > rustdesk.yml <<EOF
version: '3'

networks:
  rustdesk-net:
    external: false

services:
  rustdesk:
    container_name: rustdesk-server
    ports:
      - 21114:21114
      - 21115:21115
      - 21116:21116
      - 21116:21116/udp
      - 21117:21117
      - 21118:21118
      - 21119:21119
    image: lejianwen/rustdesk-server-s6:latest
    environment:
      - RELAY=${HOST_IP}
      # 动态配置登录限制
      - MUST_LOGIN=${MUST_LOGIN_VAL}
      - SINGLE_BANDWIDTH=16
      - TOTAL_BANDWIDTH=100
      - ENCRYPTED_ONLY=1
      - TZ=Asia/Shanghai
    volumes:
      - ${INSTALL_DIR}/data:/data
      - ${INSTALL_DIR}/api:/app/data
    networks:
      - rustdesk-net
    restart: unless-stopped
EOF

# 6. 启动容器
echo -e "${YELLOW}正在拉取镜像并启动服务...${PLAIN}"
docker compose -f rustdesk.yml pull
docker compose -f rustdesk.yml up -d

# 7. 检查状态并获取 Key
if [ $? -eq 0 ]; then
    echo -e "${GREEN}服务启动成功！${PLAIN}"
    echo -e "${YELLOW}正在等待密钥生成...${PLAIN}"
    sleep 5
    
    PUB_KEY_FILE="${INSTALL_DIR}/data/id_ed25519.pub"
    
    if [ -f "$PUB_KEY_FILE" ]; then
        PUB_KEY=$(cat "$PUB_KEY_FILE")
        echo -e "${GREEN}=============================================${PLAIN}"
        echo -e "   RustDesk Server 部署完成"
        echo -e "${GREEN}=============================================${PLAIN}"
        echo -e "ID/中继服务器:  ${GREEN}${HOST_IP}${PLAIN}"
        echo -e "API 服务器:     ${GREEN}http://${HOST_IP}:21114${PLAIN}"
        echo -e "Key (公钥):"
        echo -e "${YELLOW}${PUB_KEY}${PLAIN}"
        echo -e "${GREEN}=============================================${PLAIN}"
        if [[ "$MUST_LOGIN_VAL" == "Y" ]]; then
            echo -e "注意：当前为【登录模式】，请先在 API 页面注册账号。"
        fi
    else
        echo -e "${RED}未找到公钥，请手动查看目录 ${INSTALL_DIR}/data${PLAIN}"
    fi
else
    echo -e "${RED}服务启动失败！${PLAIN}"
fi