#!/bin/bash
# TradingAgents-CN 一键部署脚本
# 适用于 Ubuntu 20.04+ 系统

set -e  # 遇到错误立即退出

echo "=========================================="
echo "🚀 TradingAgents-CN 一键部署脚本"
echo "=========================================="

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

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

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

# 检查是否为root用户
check_root() {
    if [[ $EUID -eq 0 ]]; then
        log_error "请不要使用root用户运行此脚本"
        exit 1
    fi
}

# 检查系统
check_system() {
    log_info "检查系统环境..."
    
    if [[ ! -f /etc/os-release ]]; then
        log_error "无法识别的操作系统"
        exit 1
    fi
    
    . /etc/os-release
    if [[ $ID != "ubuntu" ]]; then
        log_warning "此脚本专为Ubuntu优化，其他系统可能需要手动调整"
    fi
    
    log_success "系统检查完成"
}

# 更新系统
update_system() {
    log_info "更新系统包..."
    sudo apt update && sudo apt upgrade -y
    sudo apt install -y git curl wget vim htop unzip software-properties-common
    log_success "系统更新完成"
}

# 安装Python 3.11
install_python() {
    log_info "安装Python 3.11..."
    
    if command -v python3.11 &> /dev/null; then
        log_success "Python 3.11 已安装"
        return
    fi
    
    sudo add-apt-repository ppa:deadsnakes/ppa -y
    sudo apt update
    sudo apt install -y python3.11 python3.11-venv python3.11-pip python3.11-dev
    
    log_success "Python 3.11 安装完成"
}

# 安装Docker
install_docker() {
    log_info "安装Docker..."
    
    if command -v docker &> /dev/null; then
        log_success "Docker 已安装"
        return
    fi
    
    curl -fsSL https://get.docker.com -o get-docker.sh
    sudo sh get-docker.sh
    sudo usermod -aG docker $USER
    rm get-docker.sh
    
    # 安装Docker Compose
    sudo curl -L "https://github.com/docker/compose/releases/latest/download/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
    sudo chmod +x /usr/local/bin/docker-compose
    
    log_success "Docker 安装完成"
    log_warning "请重新登录以使Docker权限生效，然后重新运行此脚本"
}

# 安装Nginx
install_nginx() {
    log_info "安装Nginx..."
    
    if command -v nginx &> /dev/null; then
        log_success "Nginx 已安装"
        return
    fi
    
    sudo apt install -y nginx
    sudo systemctl enable nginx
    sudo systemctl start nginx
    
    log_success "Nginx 安装完成"
}

# 克隆项目
clone_project() {
    log_info "克隆TradingAgents-CN项目..."
    
    PROJECT_DIR="$HOME/TradingAgents-CN"
    
    if [[ -d $PROJECT_DIR ]]; then
        log_warning "项目目录已存在，更新代码..."
        cd $PROJECT_DIR
        git pull origin main
    else
        git clone https://github.com/hsliuping/TradingAgents-CN.git $PROJECT_DIR
        cd $PROJECT_DIR
    fi
    
    log_success "项目代码准备完成"
}

# 创建虚拟环境
setup_venv() {
    log_info "创建Python虚拟环境..."
    
    cd $HOME/TradingAgents-CN
    
    if [[ ! -d "env" ]]; then
        python3.11 -m venv env
    fi
    
    source env/bin/activate
    pip install --upgrade pip
    pip install -r requirements.txt
    
    # 安装数据库支持（可选）
    if [[ -f requirements_db.txt ]]; then
        pip install -r requirements_db.txt
    fi
    
    log_success "虚拟环境创建完成"
}

# 配置环境变量
setup_env() {
    log_info "配置环境变量..."
    
    cd $HOME/TradingAgents-CN
    
    if [[ ! -f .env ]]; then
        if [[ -f .env.example ]]; then
            cp .env.example .env
        else
            cat > .env << EOF
# API密钥配置 (请替换为真实密钥)
DASHSCOPE_API_KEY=your_dashscope_api_key_here
FINNHUB_API_KEY=your_finnhub_api_key_here
GOOGLE_API_KEY=your_google_api_key_here

# 数据库配置
MONGODB_ENABLED=true
REDIS_ENABLED=true
MONGODB_HOST=localhost
MONGODB_PORT=27017
REDIS_HOST=localhost
REDIS_PORT=6379

# Web服务配置
STREAMLIT_SERVER_PORT=8501
STREAMLIT_SERVER_ADDRESS=0.0.0.0
EOF
        fi
        
        log_warning "请编辑 $HOME/TradingAgents-CN/.env 文件，配置真实的API密钥"
    else
        log_success "环境变量文件已存在"
    fi
}

# 启动数据库服务
start_databases() {
    log_info "启动数据库服务..."
    
    cd $HOME/TradingAgents-CN
    
    # 检查Docker是否在运行
    if ! docker info &> /dev/null; then
        log_error "Docker 未运行，请启动Docker服务"
        return 1
    fi
    
    # 启动数据库
    docker-compose up -d mongodb redis redis-commander
    
    # 等待服务启动
    log_info "等待数据库服务启动..."
    sleep 10
    
    # 检查服务状态
    if docker-compose ps | grep -q "Up"; then
        log_success "数据库服务启动成功"
    else
        log_error "数据库服务启动失败"
        return 1
    fi
}

# 配置Nginx
setup_nginx() {
    log_info "配置Nginx反向代理..."
    
    # 获取服务器IP
    SERVER_IP=$(curl -s ifconfig.me || echo "localhost")
    
    # 创建Nginx配置
    sudo tee /etc/nginx/sites-available/tradingagents > /dev/null << EOF
server {
    listen 80;
    server_name $SERVER_IP localhost;

    location / {
        proxy_pass http://127.0.0.1:8501;
        proxy_set_header Host \$host;
        proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto \$scheme;
        
        # WebSocket支持
        proxy_http_version 1.1;
        proxy_set_header Upgrade \$http_upgrade;
        proxy_set_header Connection "upgrade";
        
        # 超时设置
        proxy_connect_timeout 60s;
        proxy_send_timeout 60s;
        proxy_read_timeout 60s;
    }
}
EOF
    
    # 启用站点
    sudo ln -sf /etc/nginx/sites-available/tradingagents /etc/nginx/sites-enabled/
    sudo rm -f /etc/nginx/sites-enabled/default
    
    # 测试配置
    if sudo nginx -t; then
        sudo systemctl restart nginx
        log_success "Nginx 配置完成"
    else
        log_error "Nginx 配置错误"
        return 1
    fi
}

# 创建系统服务
setup_service() {
    log_info "创建系统服务..."
    
    USER_NAME=$(whoami)
    PROJECT_DIR="$HOME/TradingAgents-CN"
    
    sudo tee /etc/systemd/system/tradingagents.service > /dev/null << EOF
[Unit]
Description=TradingAgents-CN Web Application
After=network.target docker.service
Requires=docker.service

[Service]
Type=simple
User=$USER_NAME
WorkingDirectory=$PROJECT_DIR
Environment=PATH=$PROJECT_DIR/env/bin
ExecStart=$PROJECT_DIR/env/bin/python -m streamlit run web/app.py --server.port 8501 --server.address 0.0.0.0
Restart=always
RestartSec=10
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
EOF
    
    # 重新加载systemd并启动服务
    sudo systemctl daemon-reload
    sudo systemctl enable tradingagents
    sudo systemctl start tradingagents
    
    log_success "系统服务创建完成"
}

# 配置防火墙
setup_firewall() {
    log_info "配置防火墙..."
    
    if command -v ufw &> /dev/null; then
        sudo ufw --force enable
        sudo ufw allow ssh
        sudo ufw allow 80
        sudo ufw allow 443
        sudo ufw allow 8501
        log_success "防火墙配置完成"
    else
        log_warning "UFW 未安装，跳过防火墙配置"
    fi
}

# 检查服务状态
check_services() {
    log_info "检查服务状态..."
    
    echo "=========================================="
    echo "📊 服务状态检查"
    echo "=========================================="
    
    # 检查Docker服务
    echo "🐳 Docker服务:"
    docker-compose ps
    
    echo ""
    echo "🌐 TradingAgents服务:"
    sudo systemctl status tradingagents --no-pager -l
    
    echo ""
    echo "🔧 Nginx服务:"
    sudo systemctl status nginx --no-pager -l
}

# 显示访问信息
show_access_info() {
    SERVER_IP=$(curl -s ifconfig.me || echo "localhost")
    
    echo ""
    echo "=========================================="
    echo "🎉 部署完成！"
    echo "=========================================="
    echo ""
    echo "📱 Web访问地址:"
    echo "   - HTTP: http://$SERVER_IP"
    echo "   - 本地: http://localhost"
    echo ""
    echo "🔧 管理界面:"
    echo "   - Redis Commander: http://$SERVER_IP:8081"
    echo "   - 直接访问Streamlit: http://$SERVER_IP:8501"
    echo ""
    echo "📋 重要信息:"
    echo "   - 项目目录: $HOME/TradingAgents-CN"
    echo "   - 配置文件: $HOME/TradingAgents-CN/.env"
    echo "   - 日志查看: sudo journalctl -u tradingagents -f"
    echo ""
    echo "⚠️  下一步操作:"
    echo "   1. 编辑 $HOME/TradingAgents-CN/.env 文件，配置真实的API密钥"
    echo "   2. 重启服务: sudo systemctl restart tradingagents"
    echo "   3. 如有域名，请配置DNS A记录指向: $SERVER_IP"
    echo ""
}

# 主函数
main() {
    check_root
    check_system
    update_system
    install_python
    install_docker
    install_nginx
    clone_project
    setup_venv
    setup_env
    start_databases
    setup_nginx
    setup_service
    setup_firewall
    
    # 等待服务启动
    log_info "等待服务完全启动..."
    sleep 15
    
    check_services
    show_access_info
}

# 错误处理
trap 'log_error "部署过程中发生错误，请检查日志"; exit 1' ERR

# 运行主函数
main "$@" 