#!/bin/bash

# Secure AI Agents Suite - Deployment and Setup Script
# This script sets up the environment for HuggingFace Spaces deployment

set -e  # Exit on any error

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
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"
}

# Check if running in HuggingFace Spaces
check_spaces_environment() {
    log_info "Checking HuggingFace Spaces environment..."
    
    if [[ -z "${HF_TOKEN}" ]] && [[ -z "${HUGGINGFACE_HUB_TOKEN}" ]]; then
        log_warning "No HuggingFace token detected. Running in local development mode."
        export SPACES_ENVIRONMENT="local"
    else
        log_success "HuggingFace Spaces environment detected."
        export SPACES_ENVIRONMENT="spaces"
    fi
}

# Check system requirements
check_system_requirements() {
    log_info "Checking system requirements..."
    
    # Check Python version
    python_version=$(python3 --version 2>/dev/null | cut -d' ' -f2)
    if [[ -z "${python_version}" ]]; then
        log_error "Python 3 is not installed or not in PATH"
        exit 1
    fi
    
    log_success "Python version: ${python_version}"
    
    # Check memory
    if command -v free &> /dev/null; then
        memory_gb=$(free -g | awk '/^Mem:/{print $2}')
        if [[ ${memory_gb} -lt 2 ]]; then
            log_warning "Low memory detected (${memory_gb}GB). Recommended: 4GB+"
        else
            log_success "Sufficient memory: ${memory_gb}GB"
        fi
    fi
    
    # Check disk space
    disk_space=$(df -BG . | awk 'NR==2{print $4}' | sed 's/G//')
    if [[ ${disk_space} -lt 5 ]]; then
        log_warning "Low disk space detected (${disk_space}GB). Recommended: 10GB+"
    else
        log_success "Sufficient disk space: ${disk_space}GB"
    fi
}

# Create necessary directories
create_directories() {
    log_info "Creating necessary directories..."
    
    directories=(
        "/tmp/logs"
        "/tmp/cache"
        "/tmp/models"
        "/tmp/data"
        "./logs"
        "./cache"
        "./models"
        "./data"
        "./docs/images"
        "./tests"
        "./scripts"
    )
    
    for dir in "${directories[@]}"; do
        if [[ ! -d "${dir}" ]]; then
            mkdir -p "${dir}"
            log_success "Created directory: ${dir}"
        fi
    done
}

# Install dependencies
install_dependencies() {
    log_info "Installing dependencies..."
    
    # Upgrade pip first
    python3 -m pip install --upgrade pip setuptools wheel
    
    # Install requirements
    if [[ -f "requirements.txt" ]]; then
        log_info "Installing from requirements.txt..."
        pip install -r requirements.txt
    else
        log_warning "requirements.txt not found, installing basic dependencies..."
        pip install gradio==3.50.2 pandas numpy plotly
    fi
    
    # Install additional Spaces-specific dependencies
    pip install spaces==0.19.4 huggingface-hub==0.19.4 safetensors==0.4.1
    
    log_success "Dependencies installed successfully"
}

# Configure environment variables
configure_environment() {
    log_info "Configuring environment variables..."
    
    # Create .env file if it doesn't exist
    if [[ ! -f ".env" ]]; then
        cat > .env << EOF
# Secure AI Agents Suite - Environment Configuration
APP_NAME=Secure AI Agents Suite
APP_VERSION=2.0.0
APP_ENV=production

# Logging Configuration
LOG_LEVEL=INFO
DEBUG=false

# Performance Settings
MAX_CONCURRENT_REQUESTS=5
CACHE_TTL=3600
MODEL_CACHE_SIZE=1000
METRICS_RETENTION_DAYS=7

# Security Settings
ENABLE_RATE_LIMITING=true
MAX_REQUESTS_PER_MINUTE=100
ENABLE_INPUT_VALIDATION=true
ENABLE_AUDIT_LOGGING=true

# Spaces Configuration
HOST=0.0.0.0
PORT=7860
HF_TRANSFER=true

# Model Settings
MODEL_CACHE_DIR=/tmp/cache
HF_HOME=/tmp/cache
TRANSFORMERS_CACHE=/tmp/cache

# Monitoring
ENABLE_PERFORMANCE_MONITORING=true
HEALTH_CHECK_INTERVAL=30
EOF
        log_success "Created .env file with default configuration"
    else
        log_info ".env file already exists, skipping creation"
    fi
    
    # Load environment variables
    if [[ -f ".env" ]]; then
        source .env
        log_success "Environment variables loaded from .env"
    fi
}

# Setup performance optimizations
setup_performance_optimizations() {
    log_info "Setting up performance optimizations..."
    
    # Set Python optimization flags
    export PYTHONUNBUFFERED=1
    export PYTHONDONTWRITEBYTECODE=1
    export PYTHONOPTIMIZE=1
    
    # Set pip optimization flags
    export PIP_NO_CACHE_DIR=1
    export PIP_DISABLE_PIP_VERSION_CHECK=1
    
    # Set HuggingFace optimization flags
    export HF_HUB_ENABLE_HF_TRANSFER=1
    
    log_success "Performance optimizations configured"
}

# Create systemd service file (for local deployment)
create_systemd_service() {
    if [[ "${SPACES_ENVIRONMENT}" == "local" ]] && command -v systemctl &> /dev/null; then
        log_info "Creating systemd service file..."
        
        service_file="/tmp/secure-ai-agents-suite.service"
        cat > "${service_file}" << EOF
[Unit]
Description=Secure AI Agents Suite
After=network.target

[Service]
Type=simple
User=appuser
WorkingDirectory=$(pwd)
Environment=PYTHONPATH=$(pwd)
ExecStart=$(which python3) app.py
Restart=always
RestartSec=3

[Install]
WantedBy=multi-user.target
EOF
        
        log_success "Systemd service file created at ${service_file}"
        log_info "To install the service, run: sudo cp ${service_file} /etc/systemd/system/"
        log_info "Then enable with: sudo systemctl enable secure-ai-agents-suite"
        log_info "And start with: sudo systemctl start secure-ai-agents-suite"
    fi
}

# Setup health checks
setup_health_checks() {
    log_info "Setting up health checks..."
    
    # Create health check script
    cat > health_check.sh << 'EOF'
#!/bin/bash
# Health check script for Secure AI Agents Suite

curl -f http://localhost:7860/health > /dev/null 2>&1
if [ $? -eq 0 ]; then
    echo "Health check passed"
    exit 0
else
    echo "Health check failed"
    exit 1
fi
EOF
    
    chmod +x health_check.sh
    log_success "Health check script created"
    
    # Add to crontab if available
    if command -v crontab &> /dev/null; then
        (crontab -l 2>/dev/null; echo "*/5 * * * * $(pwd)/health_check.sh") | crontab -
        log_success "Health check cron job added (runs every 5 minutes)"
    fi
}

# Create monitoring script
create_monitoring_script() {
    log_info "Creating monitoring script..."
    
    cat > monitor.sh << 'EOF'
#!/bin/bash
# Monitoring script for Secure AI Agents Suite

check_process() {
    if pgrep -f "python.*app.py" > /dev/null; then
        echo "✓ Application process is running"
    else
        echo "✗ Application process is not running"
        return 1
    fi
}

check_port() {
    if netstat -tuln 2>/dev/null | grep -q ":7860"; then
        echo "✓ Port 7860 is listening"
    else
        echo "✗ Port 7860 is not listening"
        return 1
    fi
}

check_memory() {
    memory_usage=$(free | awk '/Mem/ {printf "%.0f", $3/$2 * 100}')
    if [[ ${memory_usage} -lt 80 ]]; then
        echo "✓ Memory usage is healthy (${memory_usage}%)"
    else
        echo "⚠ Memory usage is high (${memory_usage}%)"
    fi
}

check_disk() {
    disk_usage=$(df . | awk 'NR==2 {print $5}' | sed 's/%//')
    if [[ ${disk_usage} -lt 80 ]]; then
        echo "✓ Disk usage is healthy (${disk_usage}%)"
    else
        echo "⚠ Disk usage is high (${disk_usage}%)"
    fi
}

echo "=== Secure AI Agents Suite Monitoring Report ==="
echo "Timestamp: $(date)"
echo

check_process
check_port
check_memory
check_disk

echo
echo "=== System Resources ==="
free -h
echo
df -h .
echo
echo "=== Recent Logs ==="
if [[ -f "/tmp/spaces.log" ]]; then
    tail -10 /tmp/spaces.log
else
    echo "No logs found"
fi
EOF
    
    chmod +x monitor.sh
    log_success "Monitoring script created"
}

# Run tests
run_tests() {
    log_info "Running basic tests..."
    
    # Test Python imports
    python3 -c "
import gradio
import pandas as np
import plotly
from autonomous_engine_fixed import RefactoredAutonomousAgent
print('✓ All imports successful')
" || {
        log_error "Import test failed"
        return 1
    }
    
    log_success "Basic tests passed"
}

# Main setup function
main() {
    echo "=========================================="
    echo "Secure AI Agents Suite - Setup Script"
    echo "=========================================="
    echo
    
    check_spaces_environment
    check_system_requirements
    create_directories
    install_dependencies
    configure_environment
    setup_performance_optimizations
    create_systemd_service
    setup_health_checks
    create_monitoring_script
    run_tests
    
    echo
    echo "=========================================="
    log_success "Setup completed successfully!"
    echo "=========================================="
    echo
    
    echo "Next steps:"
    echo "1. Review configuration in .env file"
    echo "2. Run: python app.py"
    echo "3. Open: http://localhost:7860"
    echo "4. Monitor with: ./monitor.sh"
    echo
    
    if [[ "${SPACES_ENVIRONMENT}" == "local" ]]; then
        echo "For production deployment on Spaces:"
        echo "1. Push to GitHub repository"
        echo "2. Create new Space on HuggingFace"
        echo "3. Connect repository to Space"
        echo "4. Deploy automatically"
        echo
    fi
    
    log_info "Setup script completed at $(date)"
}

# Script entry point
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi