#!/bin/bash
# Optimized startup script for the FastAPI application

set -e

# Default configuration
DEFAULT_HOST="0.0.0.0"
DEFAULT_PORT="8001"
DEFAULT_API_WORKERS="4"
DEFAULT_PROFILE="balanced"
DEFAULT_LOG_LEVEL="info"
DEFAULT_MAX_RESTARTS="3"
DEFAULT_CHECK_INTERVAL="5"

# Parse command-line parameters
HOST=${HOST:-$DEFAULT_HOST}
PORT=${PORT:-$DEFAULT_PORT}
API_WORKERS=${API_WORKERS:-$DEFAULT_API_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}

# Color output
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 optimized unified startup script"
    echo "=========================================================="
    echo
}

print_config() {
    print_color $GREEN "Startup configuration:"
    echo "- API server: http://$HOST:$PORT"
    echo "- API worker processes: $API_WORKERS"
    echo "- Performance profile: $PROFILE"
    echo "- Log level: $LOG_LEVEL"
    echo "- Maximum restarts: $MAX_RESTARTS"
    echo "- Health check interval: ${CHECK_INTERVAL} seconds"
    echo
}

setup_environment() {
    print_color $YELLOW "Setting environment variables..."

    # Set environment variables based on the profile
    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

    # General optimizations
    export PYTHONUNBUFFERED=1
    export PYTHONDONTWRITEBYTECODE=1

    print_color $GREEN "Environment variables set successfully"
}

create_directories() {
    print_color $YELLOW "Creating project directories..."

    directories=(
        "projects/data"
        "projects/uploads"
        "projects/robot"
    )

    for dir in "${directories[@]}"; do
        mkdir -p "$dir"
    done

    print_color $GREEN "Project directories created successfully"
}

check_dependencies() {
    print_color $YELLOW "Checking dependencies..."

    # Check the Python command
    if ! command -v python3 &> /dev/null; then
        print_color $RED "Error: python3 command not found"
        exit 1
    fi

    # Check required packages
    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 "Error: missing required packages: ${missing_packages[*]}"
        print_color $YELLOW "Please run: pip install ${missing_packages[*]}"
        exit 1
    fi

    # Check optional packages
    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 "Notice: missing optional optimization packages: ${optional_missing[*]}"
        print_color $YELLOW "Recommended: pip install -r requirements_optimization.txt"
    fi

    print_color $GREEN "Dependency check completed"
}

start_services() {
    print_color $YELLOW "Starting services..."

    # Start the API server
    print_color $BLUE "Starting FastAPI server..."
    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 server PID: $API_PID"

    echo
    print_color $GREEN "All services started successfully!"
    print_color $GREEN "API server: http://$HOST:$PORT"
    echo "Press Ctrl+C to stop all services"
    echo
}

monitor_services() {
    local restart_counts=(0)  # API

    while true; do
        # Check the API server
        if ! kill -0 $API_PID 2>/dev/null; then
            print_color $RED "API server stopped unexpectedly"

            if [ ${restart_counts[0]} -lt $MAX_RESTARTS ]; then
                print_color $YELLOW "Restarting API server (${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 server restarted successfully, PID: $API_PID"
            else
                print_color $RED "API server restart limit reached, stopping all services"
                break
            fi
        fi

        # Wait for the next check interval
        sleep $CHECK_INTERVAL
    done
}

cleanup() {
    echo
    print_color $YELLOW "Stopping all services..."

    # Stop the API server
    if [ ! -z "$API_PID" ] && kill -0 $API_PID 2>/dev/null; then
        print_color $BLUE "Stopping API server (PID: $API_PID)..."
        kill $API_PID 2>/dev/null || true

        # Wait for graceful shutdown
        local count=0
        while kill -0 $API_PID 2>/dev/null && [ $count -lt 10 ]; do
            sleep 1
            count=$((count + 1))
        done

        # Force terminate if it is still running
        if kill -0 $API_PID 2>/dev/null; then
            print_color $RED "Force stopping API server..."
            kill -9 $API_PID 2>/dev/null || true
        fi
    fi

    print_color $GREEN "All services have been stopped"
    exit 0
}

# Main function
main() {
    print_header

    # Parse parameters
    if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
        echo "Usage: $0 [options]"
        echo
        echo "Environment variable options:"
        echo "  HOST               API bind host address (default: $DEFAULT_HOST)"
        echo "  PORT               API bind port (default: $DEFAULT_PORT)"
        echo "  API_WORKERS        Number of API worker processes (default: $DEFAULT_API_WORKERS)"
        echo "  PROFILE            Performance profile: low_memory, balanced, high_performance (default: $DEFAULT_PROFILE)"
        echo "  LOG_LEVEL          Log level: debug, info, warning, error (default: $DEFAULT_LOG_LEVEL)"
        echo "  MAX_RESTARTS       Maximum restart count (default: $DEFAULT_MAX_RESTARTS)"
        echo "  CHECK_INTERVAL     Health check interval in seconds (default: $DEFAULT_CHECK_INTERVAL)"
        echo
        echo "Examples:"
        echo "  PROFILE=high_performance API_WORKERS=8 $0"
        echo "  PORT=8080 API_WORKERS=4 $0"
        exit 0
    fi

    print_config
    check_dependencies
    setup_environment
    create_directories
    start_services

    # Set signal handlers
    trap cleanup SIGINT SIGTERM

    # Monitor services
    monitor_services
}

# Run the main function
main "$@"