#!/bin/bash
################################################################################
# PyGuardian Docker Entrypoint Script
# Handles different deployment modes and initialization
################################################################################

set -e

# Colors for logging
RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
NC='\033[0m'

# Logging function
log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" >&2
}

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

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

# Wait for service to be ready
wait_for_service() {
    local host=$1
    local port=$2
    local timeout=${3:-30}
    
    log "Waiting for $host:$port to be ready..."
    for i in $(seq 1 $timeout); do
        if timeout 1 bash -c "echo >/dev/tcp/$host/$port" 2>/dev/null; then
            success "Service $host:$port is ready"
            return 0
        fi
        sleep 1
    done
    error "Timeout waiting for $host:$port"
    return 1
}

# Initialize configuration
init_config() {
    local mode=$1
    
    log "Initializing configuration for mode: $mode"
    
    # Create config directory if not exists
    mkdir -p /opt/pyguardian/config
    
    # Copy default config if not exists
    if [[ ! -f /opt/pyguardian/config/config.yaml ]]; then
        if [[ -f /opt/pyguardian/config/config.yaml.example ]]; then
            cp /opt/pyguardian/config/config.yaml.example /opt/pyguardian/config/config.yaml
        fi
    fi
    
    # Generate auth configuration
    if [[ ! -f /opt/pyguardian/config/auth.yaml ]]; then
        log "Generating authentication configuration..."
        python3 -c "
import yaml
import secrets
import os

auth_config = {
    'authentication': {
        'enabled': True,
        'jwt_secret': secrets.token_hex(32),
        'token_expiry_minutes': 60,
        'max_agents': 100
    },
    'encryption': {
        'algorithm': 'AES-256-GCM',
        'key_derivation': 'PBKDF2',
        'iterations': 100000
    }
}

with open('/opt/pyguardian/config/auth.yaml', 'w') as f:
    yaml.dump(auth_config, f, default_flow_style=False)
    
print('✅ Authentication configuration generated')
"
    fi
    
    # Set permissions
    chmod 600 /opt/pyguardian/config/*.yaml 2>/dev/null || true
    
    success "Configuration initialized for $mode mode"
}

# Initialize database
init_database() {
    log "Initializing database..."
    
    python3 -c "
import asyncio
import sys
sys.path.insert(0, '/opt/pyguardian/src')

from storage import Storage

async def init_db():
    storage = Storage('/opt/pyguardian/data/pyguardian.db')
    await storage.init_database()
    print('✅ Database initialized successfully')

if __name__ == '__main__':
    asyncio.run(init_db())
"
    
    success "Database initialization completed"
}

# Setup monitoring
setup_monitoring() {
    log "Setting up system monitoring..."
    
    # Create monitoring script
    cat > /opt/pyguardian/monitor.py << 'EOF'
#!/usr/bin/env python3
import psutil
import json
import sys

def get_system_info():
    return {
        'cpu_percent': psutil.cpu_percent(interval=1),
        'memory_percent': psutil.virtual_memory().percent,
        'disk_percent': psutil.disk_usage('/').percent,
        'load_avg': psutil.getloadavg(),
        'boot_time': psutil.boot_time()
    }

if __name__ == '__main__':
    try:
        info = get_system_info()
        print(json.dumps(info, indent=2))
        sys.exit(0)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)
EOF
    
    chmod +x /opt/pyguardian/monitor.py
    success "Monitoring setup completed"
}

# Start controller mode
start_controller() {
    log "Starting PyGuardian Controller..."
    
    init_config "controller"
    init_database
    setup_monitoring
    
    # Validate configuration
    if [[ -z "${TELEGRAM_BOT_TOKEN:-}" ]]; then
        warn "TELEGRAM_BOT_TOKEN not set - Telegram notifications disabled"
    fi
    
    if [[ -z "${CLUSTER_SECRET:-}" ]]; then
        warn "CLUSTER_SECRET not set - using generated secret"
        export CLUSTER_SECRET=$(openssl rand -hex 32)
    fi
    
    log "Starting controller with API on port ${PYGUARDIAN_API_PORT:-8443}"
    exec python3 main.py --mode controller
}

# Start agent mode
start_agent() {
    log "Starting PyGuardian Agent..."
    
    init_config "agent"
    setup_monitoring
    
    # Validate required environment variables
    if [[ -z "${CONTROLLER_HOST:-}" ]]; then
        error "CONTROLLER_HOST environment variable is required for agent mode"
        exit 1
    fi
    
    if [[ -z "${CLUSTER_SECRET:-}" ]]; then
        error "CLUSTER_SECRET environment variable is required for agent mode"
        exit 1
    fi
    
    # Wait for controller to be ready
    wait_for_service "${CONTROLLER_HOST}" "${CONTROLLER_PORT:-8443}" 60
    
    log "Starting agent connecting to ${CONTROLLER_HOST}:${CONTROLLER_PORT:-8443}"
    exec python3 main.py --mode agent --controller "${CONTROLLER_HOST}"
}

# Start standalone mode
start_standalone() {
    log "Starting PyGuardian Standalone..."
    
    init_config "standalone"
    init_database
    setup_monitoring
    
    if [[ -z "${TELEGRAM_BOT_TOKEN:-}" ]]; then
        warn "TELEGRAM_BOT_TOKEN not set - Telegram notifications disabled"
    fi
    
    log "Starting standalone mode with API on port ${PYGUARDIAN_API_PORT:-8443}"
    exec python3 main.py --mode standalone
}

# Development mode
start_development() {
    log "Starting PyGuardian Development Mode..."
    
    init_config "development" 
    init_database
    setup_monitoring
    
    # Start Jupyter lab in background if requested
    if [[ "${START_JUPYTER:-false}" == "true" ]]; then
        log "Starting Jupyter Lab on port 8888..."
        nohup jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root &
    fi
    
    log "Development environment ready"
    log "API will be available on port ${PYGUARDIAN_API_PORT:-8443}"
    log "Jupyter Lab: http://localhost:8888 (if enabled)"
    
    exec python3 main.py --mode standalone --debug
}

# Handle signals for graceful shutdown
handle_signal() {
    log "Received shutdown signal, stopping PyGuardian..."
    kill -TERM "$child" 2>/dev/null || true
    wait "$child"
    success "PyGuardian stopped gracefully"
    exit 0
}

trap handle_signal SIGTERM SIGINT

# Main execution
main() {
    log "=== PyGuardian Docker Container Starting ==="
    log "Mode: ${1:-standalone}"
    log "Python: $(python3 --version)"
    log "User: $(whoami)"
    log "Working directory: $(pwd)"
    
    case "${1:-standalone}" in
        "controller")
            start_controller
            ;;
        "agent")
            start_agent
            ;;
        "standalone")
            start_standalone
            ;;
        "development"|"dev")
            start_development
            ;;
        *)
            error "Unknown mode: $1"
            error "Available modes: controller, agent, standalone, development"
            exit 1
            ;;
    esac
}

# Run main function with all arguments
main "$@" &
child=$!
wait "$child"