#!/bin/bash

# AI Knowledge Distillation Platform - Quick Start Script
# منصة تقطير المعرفة للذكاء الاصطناعي - سكريبت البدء السريع

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# Unicode symbols
CHECK="✅"
CROSS="❌"
WARNING="⚠️"
INFO="ℹ️"
ROCKET="🚀"
GEAR="🔧"
MEMORY="💾"
CPU="🖥️"

echo -e "${PURPLE}================================================${NC}"
echo -e "${PURPLE}    AI Knowledge Distillation Platform${NC}"
echo -e "${PURPLE}    منصة تقطير المعرفة للذكاء الاصطناعي${NC}"
echo -e "${PURPLE}================================================${NC}"
echo ""

# Function to print colored output
print_status() {
    echo -e "${GREEN}${CHECK}${NC} $1"
}

print_error() {
    echo -e "${RED}${CROSS}${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}${WARNING}${NC} $1"
}

print_info() {
    echo -e "${BLUE}${INFO}${NC} $1"
}

# Check if Python is installed
check_python() {
    if command -v python3 &> /dev/null; then
        PYTHON_VERSION=$(python3 --version | cut -d' ' -f2)
        print_status "Python $PYTHON_VERSION found"
        return 0
    else
        print_error "Python 3 not found. Please install Python 3.9 or higher."
        return 1
    fi
}

# Check system requirements
check_system() {
    print_info "Checking system requirements..."
    
    # Check memory
    if command -v free &> /dev/null; then
        TOTAL_MEM=$(free -g | awk '/^Mem:/{print $2}')
        if [ "$TOTAL_MEM" -ge 4 ]; then
            print_status "Memory: ${TOTAL_MEM}GB (sufficient)"
        else
            print_warning "Memory: ${TOTAL_MEM}GB (minimum 4GB recommended)"
        fi
    fi
    
    # Check CPU cores
    if command -v nproc &> /dev/null; then
        CPU_CORES=$(nproc)
        print_status "CPU cores: $CPU_CORES"
    fi
    
    # Check disk space
    DISK_SPACE=$(df -h . | awk 'NR==2{print $4}')
    print_status "Available disk space: $DISK_SPACE"
}

# Create necessary directories
create_directories() {
    print_info "Creating necessary directories..."
    
    directories=(
        "cache"
        "cache/datasets"
        "cache/transformers"
        "cache/medical_datasets"
        "database"
        "logs"
        "models"
        "backups"
        "uploads"
        "temp"
    )
    
    for dir in "${directories[@]}"; do
        if [ ! -d "$dir" ]; then
            mkdir -p "$dir"
            print_status "Created directory: $dir"
        fi
    done
}

# Install dependencies
install_dependencies() {
    print_info "Checking dependencies..."
    
    if [ ! -f "requirements.txt" ]; then
        print_error "requirements.txt not found!"
        return 1
    fi
    
    # Check if virtual environment exists
    if [ ! -d "venv" ]; then
        print_info "Creating virtual environment..."
        python3 -m venv venv
        print_status "Virtual environment created"
    fi
    
    # Activate virtual environment
    source venv/bin/activate
    
    # Upgrade pip
    print_info "Upgrading pip..."
    pip install --upgrade pip
    
    # Install dependencies
    print_info "Installing dependencies..."
    pip install -r requirements.txt
    
    print_status "Dependencies installed"
}

# Set environment variables
set_environment() {
    print_info "Setting environment variables..."
    
    # CPU optimization
    export OMP_NUM_THREADS=$(nproc)
    export MKL_NUM_THREADS=$(nproc)
    export NUMEXPR_NUM_THREADS=$(nproc)
    export OPENBLAS_NUM_THREADS=$(nproc)
    
    # Memory optimization
    export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
    export TOKENIZERS_PARALLELISM=false
    
    # Disable GPU (force CPU-only)
    export CUDA_VISIBLE_DEVICES=""
    
    # Cache directories
    export HF_DATASETS_CACHE=./cache/datasets
    export TRANSFORMERS_CACHE=./cache/transformers
    export HF_HOME=./cache/huggingface
    
    print_status "Environment variables set"
}

# Start the application
start_application() {
    print_info "Starting application..."
    
    # Check which runner to use
    if [ -f "run_optimized.py" ]; then
        print_status "Using optimized runner"
        python run_optimized.py
    elif [ -f "app.py" ]; then
        print_status "Using standard runner"
        python app.py
    else
        print_error "No application file found!"
        return 1
    fi
}

# Main execution
main() {
    echo -e "${CYAN}${ROCKET} Starting setup process...${NC}"
    echo ""
    
    # Check Python
    if ! check_python; then
        exit 1
    fi
    
    # Check system
    check_system
    echo ""
    
    # Create directories
    create_directories
    echo ""
    
    # Install dependencies
    if [ "$1" != "--skip-install" ]; then
        install_dependencies
        echo ""
    else
        print_info "Skipping dependency installation"
        # Still activate venv if it exists
        if [ -d "venv" ]; then
            source venv/bin/activate
        fi
    fi
    
    # Set environment
    set_environment
    echo ""

    # Setup tokens
    if [ -f "setup_tokens.py" ]; then
        print_info "Setting up Hugging Face tokens..."
        python setup_tokens.py
        echo ""
    fi

    # Final status
    echo -e "${GREEN}${CHECK} Setup completed successfully!${NC}"
    echo ""
    echo -e "${CYAN}${GEAR} System Information:${NC}"
    echo -e "  ${MEMORY} Memory optimization: Enabled"
    echo -e "  ${CPU} CPU threads: $OMP_NUM_THREADS"
    echo -e "  🔒 Security: Token encryption enabled"
    echo -e "  🏥 Medical AI: Supported"
    echo ""
    echo -e "${YELLOW}${ROCKET} Starting AI Knowledge Distillation Platform...${NC}"
    echo -e "${BLUE}🌐 Access the application at: http://localhost:8000${NC}"
    echo -e "${BLUE}🔑 Token management: http://localhost:8000/tokens${NC}"
    echo -e "${BLUE}🏥 Medical datasets: http://localhost:8000/medical-datasets${NC}"
    echo ""
    echo -e "${PURPLE}================================================${NC}"
    
    # Start application
    start_application
}

# Handle script arguments
case "$1" in
    --help|-h)
        echo "Usage: $0 [OPTIONS]"
        echo ""
        echo "Options:"
        echo "  --help, -h          Show this help message"
        echo "  --skip-install      Skip dependency installation"
        echo "  --check-only        Only check system requirements"
        echo ""
        echo "Examples:"
        echo "  $0                  Full setup and start"
        echo "  $0 --skip-install   Start without installing dependencies"
        echo "  $0 --check-only     Check system requirements only"
        exit 0
        ;;
    --check-only)
        check_python
        check_system
        exit 0
        ;;
    *)
        main "$@"
        ;;
esac