#!/bin/bash

set -e

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

# 输出颜色信息
log() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

# 检测系统IP
detect_ip() {
    log "正在检测服务器IP地址..."
    
    # 尝试多种方法获取公网IP
    PUBLIC_IP=$(curl -s --connect-timeout 5 http://ipinfo.io/ip || \
                curl -s --connect-timeout 5 http://ifconfig.me || \
                curl -s --connect-timeout 5 http://api.ipify.org || \
                echo "unknown")
    
    # 获取内网IP
    LOCAL_IP=$(ip route get 1 | awk '{print $7; exit}' 2>/dev/null || \
               hostname -I | awk '{print $1}' 2>/dev/null || \
               echo "unknown")
    
    # 显示检测结果
    if [ "$PUBLIC_IP" != "unknown" ]; then
        info "检测到公网IP: $PUBLIC_IP"
        SERVER_IP="$PUBLIC_IP"
    else
        warn "无法获取公网IP，使用内网IP: $LOCAL_IP"
        SERVER_IP="$LOCAL_IP"
    fi
}

# 生成密钥
generate_keys() {
    log "正在生成RustDesk密钥对..."
    
    # 创建目录
    mkdir -p {hbbs,hbbr,data}
    
    # 检查是否已有密钥
    if [ -f "data/key_pair" ]; then
        warn "检测到已存在的密钥对，使用现有密钥"
        # 从保存的密钥对中恢复
        cat data/key_pair | grep "私钥" | cut -d' ' -f2 | base64 -d > data/id_ed25519
        cat data/key_pair | grep "公钥" | cut -d' ' -f2 | base64 -d > data/id_ed25519.pub
    else
        # 生成新密钥
        openssl genrsa -out data/id_ed25519 2048 2>/dev/null
        openssl rsa -in data/id_ed25519 -pubout -out data/id_ed25519.pub 2>/dev/null
        
        # 保存密钥对信息
        echo "私钥: $(cat data/id_ed25519 | base64 -w 0)" > data/key_pair
        echo "公钥: $(cat data/id_ed25519.pub | base64 -w 0)" >> data/key_pair
        log "新的密钥对生成完成"
    fi
    
    # 复制密钥到服务目录
    cp data/id_ed25519* hbbs/
    cp data/id_ed25519* hbbr/
    
    # 获取公钥用于显示
    PUBLIC_KEY=$(cat data/id_ed25519.pub | base64 -w 0)
}

# 生成Docker Compose配置
generate_docker_compose() {
    log "生成Docker Compose配置文件..."
    
    cat > docker-compose.yml << EOF
version: '3.8'

services:
  hbbs:
    image: lejianwen/rustdesk-server-s6
    container_name: rustdesk-hbbs
    restart: unless-stopped
    ports:
      - "21115:21115"
      - "21116:21116"
      - "21116:21116/udp"
      - "21117:21117"
      - "21118:21118"
      - "21119:21119"
    volumes:
      - ./hbbs:/root
      - ./data:/data
    command: hbbs -r ${SERVER_IP}:21117
    networks:
      - rustdesk-net

  hbbr:
    image: lejianwen/rustdesk-server-s6
    container_name: rustdesk-hbbr
    restart: unless-stopped
    ports:
      - "21117:21117"
      - "21118:21118"
      - "21119:21119"
    volumes:
      - ./hbbr:/root
      - ./data:/data
    command: hbbr
    networks:
      - rustdesk-net

networks:
  rustdesk-net:
    driver: bridge
EOF
}

# 生成管理脚本
generate_management_script() {
    log "生成管理脚本..."
    
    cat > manage.sh << 'EOF'
#!/bin/bash

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

# 管理功能
case "$1" in
    "start")
        docker-compose up -d
        echo -e "${GREEN}服务已启动${NC}"
        ;;
    "stop")
        docker-compose down
        echo -e "${YELLOW}服务已停止${NC}"
        ;;
    "restart")
        docker-compose restart
        echo -e "${GREEN}服务已重启${NC}"
        ;;
    "status")
        docker-compose ps
        ;;
    "logs")
        docker-compose logs -f
        ;;
    "update")
        docker-compose pull
        docker-compose down
        docker-compose up -d
        echo -e "${GREEN}服务已更新${NC}"
        ;;
    "backup")
        BACKUP_DIR="backup/$(date +%Y%m%d_%H%M%S)"
        mkdir -p $BACKUP_DIR
        cp -r data $BACKUP_DIR/
        echo -e "${GREEN}密钥已备份到: $BACKUP_DIR${NC}"
        ;;
    "key-info")
        if [ -f "data/key_pair" ]; then
            echo -e "${BLUE}=== 密钥信息 ===${NC}"
            echo "公钥 (Base64): $(cat data/key_pair | grep '公钥' | cut -d' ' -f2)"
            echo -e "${BLUE}===============${NC}"
        else
            echo -e "${RED}没有找到密钥对${NC}"
        fi
        ;;
    "export-key")
        if [ -f "data/key_pair" ]; then
            echo -e "${YELLOW}=== 导出密钥 ===${NC}"
            cat data/key_pair
            echo -e "${YELLOW}===============${NC}"
        else
            echo -e "${RED}没有找到密钥对${NC}"
        fi
        ;;
    "import-key")
        if [ -z "$2" ]; then
            echo -e "${RED}用法: ./manage.sh import-key <base64_private_key>${NC}"
            exit 1
        fi
        
        echo "$2" | base64 -d > data/id_ed25519
        openssl rsa -in data/id_ed25519 -pubout -out data/id_ed25519.pub 2>/dev/null
        
        echo "私钥: $(cat data/id_ed25519 | base64 -w 0)" > data/key_pair
        echo "公钥: $(cat data/id_ed25519.pub | base64 -w 0)" >> data/key_pair
        
        cp data/id_ed25519* hbbs/
        cp data/id_ed25519* hbbr/
        
        docker-compose restart
        echo -e "${GREEN}密钥已导入并应用${NC}"
        ;;
    *)
        echo -e "${BLUE}RustDesk 服务管理脚本${NC}"
        echo "用法: $0 [命令]"
        echo ""
        echo "命令:"
        echo "  start       启动服务"
        echo "  stop        停止服务"
        echo "  restart     重启服务"
        echo "  status      查看状态"
        echo "  logs        查看日志"
        echo "  update      更新服务"
        echo "  backup      备份密钥"
        echo "  key-info    查看密钥信息"
        echo "  export-key  导出密钥"
        echo "  import-key  导入密钥 (需要Base64格式私钥)"
        ;;
esac
EOF

    chmod +x manage.sh
}

# 显示部署信息
show_deployment_info() {
    echo ""
    echo -e "${GREEN}===============================================${NC}"
    echo -e "${GREEN}          RustDesk Server 部署完成！          ${NC}"
    echo -e "${GREEN}===============================================${NC}"
    echo ""
    echo -e "${BLUE}服务器信息:${NC}"
    echo "  ID 服务器: $SERVER_IP"
    echo "  中继服务器: $SERVER_IP:21117"
    echo ""
    echo -e "${BLUE}密钥信息:${NC}"
    echo "  公钥 (Base64): $(cat data/key_pair | grep '公钥' | cut -d' ' -f2)"
    echo ""
    echo -e "${BLUE}端口信息:${NC}"
    echo "  HBBS 端口: 21115-21119"
    echo "  HBBR 端口: 21117-21119"
    echo ""
    echo -e "${BLUE}管理命令:${NC}"
    echo "  启动服务: ./manage.sh start"
    echo "  停止服务: ./manage.sh stop"
    echo "  查看状态: ./manage.sh status"
    echo "  查看日志: ./manage.sh logs"
    echo "  备份密钥: ./manage.sh backup"
    echo "  导出密钥: ./manage.sh export-key"
    echo ""
    echo -e "${YELLOW}重要提示:${NC}"
    echo "  1. 请确保防火墙开放端口 21115-21119"
    echo "  2. 备份 data/ 目录下的密钥文件"
    echo "  3. 在其他服务器部署时使用相同密钥保证客户端兼容"
    echo -e "${GREEN}===============================================${NC}"
}

# 检查防火墙
check_firewall() {
    log "检查防火墙状态..."
    
    # 检查 ufw
    if command -v ufw >/dev/null 2>&1 && ufw status | grep -q "active"; then
        warn "检测到 ufw 防火墙，请确保已开放端口 21115-21119"
        echo "运行以下命令开放端口:"
        echo "  sudo ufw allow 21115:21119/tcp"
        echo "  sudo ufw allow 21116/udp"
    fi
    
    # 检查 firewalld
    if command -v firewall-cmd >/dev/null 2>&1 && firewall-cmd --state >/dev/null 2>&1; then
        warn "检测到 firewalld 防火墙，请确保已开放端口 21115-21119"
        echo "运行以下命令开放端口:"
        echo "  sudo firewall-cmd --permanent --add-port=21115-21119/tcp"
        echo "  sudo firewall-cmd --permanent --add-port=21116/udp"
        echo "  sudo firewall-cmd --reload"
    fi
    
    # 检查 iptables
    if command -v iptables >/dev/null 2>&1; then
        warn "请检查 iptables 规则，确保端口 21115-21119 已开放"
    fi
}

# 主部署流程
main() {
    echo -e "${BLUE}===============================================${NC}"
    echo -e "${BLUE}         RustDesk Server 一键部署脚本         ${NC}"
    echo -e "${BLUE}===============================================${NC}"
    
    # 检测IP
    detect_ip
    
    # 生成密钥
    generate_keys
    
    # 生成Docker配置
    generate_docker_compose
    
    # 生成管理脚本
    generate_management_script
    
    # 启动服务
    log "启动RustDesk服务..."
    docker-compose up -d
    
    # 显示部署信息
    show_deployment_info
    
    # 检查防火墙
    check_firewall
    
    echo ""
    log "部署完成！可以使用 ./manage.sh 管理服务"
}

# 运行主函数
main