#!/bin/bash
# 优化版启动脚本 - 整合 FastAPI 应用和队列消费者

set -e

# 默认配置
DEFAULT_HOST="0.0.0.0"
DEFAULT_PORT="8001"
DEFAULT_API_WORKERS="4"
DEFAULT_QUEUE_WORKERS="2"
DEFAULT_PROFILE="balanced"
DEFAULT_LOG_LEVEL="info"
DEFAULT_MAX_RESTARTS="3"
DEFAULT_CHECK_INTERVAL="5"

# 解析命令行参数
HOST=${HOST:-$DEFAULT_HOST}
PORT=${PORT:-$DEFAULT_PORT}
API_WORKERS=${API_WORKERS:-$DEFAULT_API_WORKERS}
QUEUE_WORKERS=${QUEUE_WORKERS:-$DEFAULT_QUEUE_WORKERS}
PROFILE=${PROFILE:-$DEFAULT_PROFILE}
LOG_LEVEL=${LOG_LEVEL:-$DEFAULT_LOG_LEVEL}
MAX_RESTARTS=${MAX_RESTARTS:-$DEFAULT_MAX_RESTARTS}
CHECK_INTERVAL=${CHECK_INTERVAL:-$DEFAULT_CHECK_INTERVAL}

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

print_color() {
    local color=$1
    local message=$2
    echo -e "${color}${message}${NC}"
}

print_header() {
    echo "=========================================================="
    print_color $BLUE "Qwen-Agent 优化统一启动脚本"
    echo "=========================================================="
    echo
}

print_config() {
    print_color $GREEN "启动配置："
    echo "- API服务器: http://$HOST:$PORT"
    echo "- API工作进程: $API_WORKERS"
    echo "- 队列工作线程: $QUEUE_WORKERS"
    echo "- 性能配置: $PROFILE"
    echo "- 日志级别: $LOG_LEVEL"
    echo "- 最大重启次数: $MAX_RESTARTS"
    echo "- 健康检查间隔: ${CHECK_INTERVAL}秒"
    echo
}

setup_environment() {
    print_color $YELLOW "设置环境变量..."
    
    # 根据配置文件设置环境变量
    case $PROFILE in
        "low_memory")
            export TOKENIZERS_PARALLELISM=false
            export TOOL_CACHE_MAX_SIZE=20
            ;;
        "balanced")
            export TOKENIZERS_PARALLELISM=true
            export TOKENIZERS_FAST=1
            export TOOL_CACHE_MAX_SIZE=50
            ;;
        "high_performance")
            export TOKENIZERS_PARALLELISM=true
            export TOKENIZERS_FAST=1
            export TOOL_CACHE_MAX_SIZE=100
            ;;
    esac

    # 通用优化
    export PYTHONUNBUFFERED=1
    export PYTHONDONTWRITEBYTECODE=1
    
    print_color $GREEN "环境变量设置完成"
}

create_directories() {
    print_color $YELLOW "创建项目目录..."
    
    directories=(
        "projects/queue_data"
        "projects/data"
        "projects/uploads"
        "projects/robot"
    )
    
    for dir in "${directories[@]}"; do
        mkdir -p "$dir"
    done
    
    print_color $GREEN "项目目录创建完成"
}

check_dependencies() {
    print_color $YELLOW "检查依赖..."
    
    # 检查 Python 命令
    if ! command -v python3 &> /dev/null; then
        print_color $RED "错误: 未找到 python3 命令"
        exit 1
    fi
    
    # 检查必需的包
    local missing_packages=()
    
    if ! python3 -c "import uvicorn" 2>/dev/null; then
        missing_packages+=("uvicorn")
    fi
    
    if ! python3 -c "import fastapi" 2>/dev/null; then
        missing_packages+=("fastapi")
    fi
    
    if [ ${#missing_packages[@]} -ne 0 ]; then
        print_color $RED "错误: 缺少必需的包: ${missing_packages[*]}"
        print_color $YELLOW "请运行: pip install ${missing_packages[*]}"
        exit 1
    fi
    
    # 检查可选包
    local optional_missing=()
    
    if ! python3 -c "import psutil" 2>/dev/null; then
        optional_missing+=("psutil")
    fi
    
    if ! python3 -c "import uvloop" 2>/dev/null; then
        optional_missing+=("uvloop")
    fi
    
    if [ ${#optional_missing[@]} -ne 0 ]; then
        print_color $YELLOW "提示: 缺少可选优化包: ${optional_missing[*]}"
        print_color $YELLOW "建议运行: pip install -r requirements_optimization.txt"
    fi
    
    print_color $GREEN "依赖检查完成"
}

start_services() {
    print_color $YELLOW "启动服务..."
    
    # 启动 API 服务器
    print_color $BLUE "启动 FastAPI 服务器..."
    python3 -m uvicorn fastapi_app:app \
        --host $HOST \
        --port $PORT \
        --workers $API_WORKERS \
        --log-level $LOG_LEVEL \
        --access-log \
        > api_server.log 2>&1 &
    
    API_PID=$!
    echo "API 服务器 PID: $API_PID"
    
    # 启动队列消费者
    print_color $BLUE "启动队列消费者..."
    python3 task_queue/consumer.py \
        --workers=$QUEUE_WORKERS \
        --worker-type=threads \
        > queue_consumer.log 2>&1 &
    
    CONSUMER_PID=$!
    echo "队列消费者 PID: $CONSUMER_PID"
    
    echo
    print_color $GREEN "所有服务启动成功！"
    print_color $GREEN "API 服务器: http://$HOST:$PORT"
    echo "按 Ctrl+C 停止所有服务"
    echo
}

monitor_services() {
    local restart_counts=(0 0)  # API, Consumer
    
    while true; do
        # 检查 API 服务器
        if ! kill -0 $API_PID 2>/dev/null; then
            print_color $RED "API 服务器意外停止"
            
            if [ ${restart_counts[0]} -lt $MAX_RESTARTS ]; then
                print_color $YELLOW "重启 API 服务器 (${restart_counts[0]} + 1/$MAX_RESTARTS)..."
                python3 -m uvicorn fastapi_app:app \
                    --host $HOST \
                    --port $PORT \
                    --workers $API_WORKERS \
                    --log-level $LOG_LEVEL \
                    --access-log \
                    >> api_server.log 2>&1 &
                
                API_PID=$!
                restart_counts[0]=$((restart_counts[0] + 1))
                print_color $GREEN "API 服务器重启成功，PID: $API_PID"
            else
                print_color $RED "API 服务器重启次数已达上限，停止所有服务"
                break
            fi
        fi
        
        # 检查队列消费者
        if ! kill -0 $CONSUMER_PID 2>/dev/null; then
            print_color $RED "队列消费者意外停止"
            
            if [ ${restart_counts[1]} -lt $MAX_RESTARTS ]; then
                print_color $YELLOW "重启队列消费者 (${restart_counts[1]} + 1/$MAX_RESTARTS)..."
                python3 task_queue/consumer.py \
                    --workers=$QUEUE_WORKERS \
                    --worker-type=threads \
                    >> queue_consumer.log 2>&1 &
                
                CONSUMER_PID=$!
                restart_counts[1]=$((restart_counts[1] + 1))
                print_color $GREEN "队列消费者重启成功，PID: $CONSUMER_PID"
            else
                print_color $RED "队列消费者重启次数已达上限，停止所有服务"
                break
            fi
        fi
        
        # 等待检查间隔
        sleep $CHECK_INTERVAL
    done
}

cleanup() {
    echo
    print_color $YELLOW "正在停止所有服务..."
    
    # 停止 API 服务器
    if [ ! -z "$API_PID" ] && kill -0 $API_PID 2>/dev/null; then
        print_color $BLUE "停止 API 服务器 (PID: $API_PID)..."
        kill $API_PID 2>/dev/null || true
        
        # 等待优雅关闭
        local count=0
        while kill -0 $API_PID 2>/dev/null && [ $count -lt 10 ]; do
            sleep 1
            count=$((count + 1))
        done
        
        # 如果还未关闭，强制终止
        if kill -0 $API_PID 2>/dev/null; then
            print_color $RED "强制终止 API 服务器..."
            kill -9 $API_PID 2>/dev/null || true
        fi
    fi
    
    # 停止队列消费者
    if [ ! -z "$CONSUMER_PID" ] && kill -0 $CONSUMER_PID 2>/dev/null; then
        print_color $BLUE "停止队列消费者 (PID: $CONSUMER_PID)..."
        kill $CONSUMER_PID 2>/dev/null || true
        
        # 等待优雅关闭
        local count=0
        while kill -0 $CONSUMER_PID 2>/dev/null && [ $count -lt 10 ]; do
            sleep 1
            count=$((count + 1))
        done
        
        # 如果还未关闭，强制终止
        if kill -0 $CONSUMER_PID 2>/dev/null; then
            print_color $RED "强制终止队列消费者..."
            kill -9 $CONSUMER_PID 2>/dev/null || true
        fi
    fi
    
    print_color $GREEN "所有服务已停止"
    exit 0
}

# 主函数
main() {
    print_header
    
    # 解析参数
    if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
        echo "用法: $0 [选项]"
        echo
        echo "环境变量选项："
        echo "  HOST               API绑定主机地址 (默认: $DEFAULT_HOST)"
        echo "  PORT               API绑定端口 (默认: $DEFAULT_PORT)"
        echo "  API_WORKERS        API工作进程数 (默认: $DEFAULT_API_WORKERS)"
        echo "  QUEUE_WORKERS      队列工作线程数 (默认: $DEFAULT_QUEUE_WORKERS)"
        echo "  PROFILE            性能配置文件: low_memory, balanced, high_performance (默认: $DEFAULT_PROFILE)"
        echo "  LOG_LEVEL          日志级别: debug, info, warning, error (默认: $DEFAULT_LOG_LEVEL)"
        echo "  MAX_RESTARTS       最大重启次数 (默认: $DEFAULT_MAX_RESTARTS)"
        echo "  CHECK_INTERVAL     健康检查间隔秒数 (默认: $DEFAULT_CHECK_INTERVAL)"
        echo
        echo "示例："
        echo "  PROFILE=high_performance API_WORKERS=8 $0"
        echo "  PORT=8080 QUEUE_WORKERS=4 $0"
        exit 0
    fi
    
    print_config
    check_dependencies
    setup_environment
    create_directories
    start_services
    
    # 设置信号处理
    trap cleanup SIGINT SIGTERM
    
    # 监控服务
    monitor_services
}

# 运行主函数
main "$@"