#!/bin/bash
# ZETA Components Setup Script
# Enterprise-Grade MCP-ZERO Integration with Offline-First Resilience

set -e

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

# Default paths
MCPZERO_ROOT=${MCPZERO_ROOT:-$(pwd)}
PYTHON_BIN=${PYTHON_BIN:-python3}
VIRTUAL_ENV=${VIRTUAL_ENV:-"${MCPZERO_ROOT}/.venv"}
OFFLINE_MODE=${OFFLINE_MODE:-"true"}
CONNECTION_ATTEMPTS=${CONNECTION_ATTEMPTS:-1}
RESILIENCE_MODE=${RESILIENCE_MODE:-"permanent"}

# Parse arguments
INSTALL_CUDA=true
INSTALL_TENSOR=true
INSTALL_FACE=true
INSTALL_LLM=true
VERBOSE=false

print_usage() {
    echo "Usage: $0 [options]"
    echo "Options:"
    echo "  --no-cuda         Skip ZETACUDA installation"
    echo "  --no-tensor       Skip ZETATENSOR installation"
    echo "  --no-face         Skip ZETAFACE installation" 
    echo "  --no-llm          Skip LLM training components"
    echo "  --online          Attempt online package installation"
    echo "  --verbose         Show detailed output"
    echo "  --help            Show this help message"
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --no-cuda)
            INSTALL_CUDA=false
            shift
            ;;
        --no-tensor)
            INSTALL_TENSOR=false
            shift
            ;;
        --no-face)
            INSTALL_FACE=false
            shift
            ;;
        --no-llm)
            INSTALL_LLM=false
            shift
            ;;
        --online)
            OFFLINE_MODE=false
            shift
            ;;
        --verbose)
            VERBOSE=true
            shift
            ;;
        --help)
            print_usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            print_usage
            exit 1
            ;;
    esac
done

log() {
    if [ "$VERBOSE" = true ] || [ "$2" = "important" ]; then
        echo -e "${GREEN}[ZETA SETUP]${NC} $1"
    fi
}

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

error() {
    echo -e "${RED}[ZETA ERROR]${NC} $1"
    exit 1
}

# Create directories
create_directories() {
    log "Creating directories..."
    
    # Create playground directory
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/playground/zetacuda"
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/playground/zetacuda/src"
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/playground/zetatensor"
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/playground/zetaface"
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/playground/zeta_llm"
    
    # Create data directories
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/models"
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/data"
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/cache"
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/lock"
    
    # Create log directory
    mkdir -p "${MCPZERO_ROOT}/.mcpzero/logs"
    
    log "Directories created successfully"
}

# Setup Python virtual environment with offline-first resilience
setup_python_env() {
    log "Setting up Python virtual environment..." "important"
    
    # Create virtual environment if it doesn't exist
    if [ ! -d "$VIRTUAL_ENV" ]; then
        log "Creating virtual environment at $VIRTUAL_ENV"
        $PYTHON_BIN -m venv "$VIRTUAL_ENV"
    else
        log "Using existing virtual environment at $VIRTUAL_ENV"
    fi
    
    # Activate virtual environment
    source "$VIRTUAL_ENV/bin/activate"
    
    # Install requirements with offline-first resilience
    if [ "$OFFLINE_MODE" = true ]; then
        log "Installing packages in offline mode..."
        pip install --no-index --find-links="${MCPZERO_ROOT}/mlops/wheelhouse" -r "${MCPZERO_ROOT}/mlops/requirements-zeta.txt" || {
            warn "Offline package installation failed, some features may not work"
            # Record offline installation failure
            echo "$(date +%s),setup_zeta,pip_offline_install_failed" >> "${MCPZERO_ROOT}/.mcpzero/lock/installation_failures.log"
        }
    else
        log "Attempting online package installation..."
        # First try online installation
        pip install -r "${MCPZERO_ROOT}/mlops/requirements-zeta.txt" || {
            warn "Online installation failed, falling back to offline mode"
            pip install --no-index --find-links="${MCPZERO_ROOT}/mlops/wheelhouse" -r "${MCPZERO_ROOT}/mlops/requirements-zeta.txt" || {
                warn "Offline package installation failed, some features may not work"
                # Record installation failure
                echo "$(date +%s),setup_zeta,pip_install_failed" >> "${MCPZERO_ROOT}/.mcpzero/lock/installation_failures.log"
            }
        }
    fi
    
    log "Python environment setup complete" "important"
}

# Setup ZETACUDA
setup_zetacuda() {
    if [ "$INSTALL_CUDA" = false ]; then
        log "Skipping ZETACUDA installation"
        return 0
    fi
    
    log "Setting up ZETACUDA..." "important"
    
    # Check for Rust
    if command -v rustc &> /dev/null && [ "$OFFLINE_MODE" = false ]; then
        log "Rust is installed, building ZETACUDA from source"
        
        # Copy Cargo.toml if not exists
        if [ ! -f "${MCPZERO_ROOT}/.mcpzero/playground/zetacuda/Cargo.toml" ]; then
            cat > "${MCPZERO_ROOT}/.mcpzero/playground/zetacuda/Cargo.toml" << EOL
[package]
name = "zetacuda"
version = "0.1.0"
edition = "2021"

[lib]
name = "zetacuda"
crate-type = ["cdylib", "rlib"]

[dependencies]
num = "0.4"
rayon = "1.6"
EOL
        fi
        
        # Build ZETACUDA
        cd "${MCPZERO_ROOT}/.mcpzero/playground/zetacuda"
        cargo build --release || {
            warn "Failed to build ZETACUDA, using Python fallback"
            echo "$(date +%s),setup_zeta,zetacuda_build_failed" >> "${MCPZERO_ROOT}/.mcpzero/lock/installation_failures.log"
        }
        
    else
        log "Rust not available or offline mode enabled, using Python implementation"
    fi
    
    log "ZETACUDA setup complete" "important"
}

# Setup environment variables
setup_env_vars() {
    log "Setting up environment variables..." "important"
    
    # Create .env file if it doesn't exist
    ENV_FILE="${MCPZERO_ROOT}/.mcpzero/.env"
    if [ ! -f "$ENV_FILE" ]; then
        cat > "$ENV_FILE" << EOL
# ZETA Components Environment Variables
MCPZERO_ROOT=${MCPZERO_ROOT}
MCP_ZERO_OFFLINE_FIRST=${OFFLINE_MODE}
CONNECTION_ATTEMPTS=${CONNECTION_ATTEMPTS}
RESILIENCE_MODE=${RESILIENCE_MODE}

# ZETACUDA Configuration
ZETACUDA_ENABLED=true
ZETACUDA_SIMD=true
ZETACUDA_MEMORY_LIMIT_MB=1024

# ZETATENSOR Configuration
ZETATENSOR_ENABLED=true
ZETATENSOR_DIMENSIONS=64,64

# ZETAFACE Configuration
ZETAFACE_ENABLED=true
ZETAFACE_CACHE_DIR=${MCPZERO_ROOT}/.mcpzero/cache

# ZETA LLM Configuration
ZETA_LLM_ENABLED=true
LLM_PARAMETER_COUNT=3000000
LLM_DEFAULT_BATCH_SIZE=8
LLM_LEARNING_RATE=5e-5
LLM_MAX_STEPS=1000
LLM_TRAINING_DEVICE=cpu
LLM_MEMORY_LIMIT_MB=1024

# Logging Configuration
LOG_DIR=${MCPZERO_ROOT}/.mcpzero/logs
LOG_LEVEL=INFO
EOL
    else
        log ".env file already exists, keeping existing configuration"
    fi
    
    # Create activation script
    ACTIVATE_SCRIPT="${MCPZERO_ROOT}/.mcpzero/activate_zeta.sh"
    cat > "$ACTIVATE_SCRIPT" << EOL
#!/bin/bash
# Activation script for ZETA components

# Activate Python environment
source "${VIRTUAL_ENV}/bin/activate"

# Set environment variables
export MCPZERO_ROOT=${MCPZERO_ROOT}
export PYTHONPATH=\${PYTHONPATH}:${MCPZERO_ROOT}/.mcpzero/playground
export MCP_ZERO_OFFLINE_FIRST=${OFFLINE_MODE}
export CONNECTION_ATTEMPTS=${CONNECTION_ATTEMPTS}
export RESILIENCE_MODE=${RESILIENCE_MODE}
export LOG_DIR=${MCPZERO_ROOT}/.mcpzero/logs
export LOG_LEVEL=INFO

# Load all environment variables from .env file
set -a
source "${ENV_FILE}"
set +a

echo "ZETA environment activated"
EOL

    chmod +x "$ACTIVATE_SCRIPT"
    
    log "Environment variables setup complete" "important"
}

# Create requirements file
create_requirements() {
    log "Creating requirements file..." "important"
    
    REQUIREMENTS_FILE="${MCPZERO_ROOT}/mlops/requirements-zeta.txt"
    
    if [ ! -f "$REQUIREMENTS_FILE" ]; then
        cat > "$REQUIREMENTS_FILE" << EOL
# ZETA Components Requirements
numpy>=1.20.0
scipy>=1.7.0
pyyaml>=5.4.0
tqdm>=4.62.0
matplotlib>=3.4.0
psutil>=5.8.0
requests>=2.25.0
redis>=4.0.0

# Optional dependencies
# torch>=1.9.0  # Uncomment for full LLM training support
# tensorboard>=2.6.0  # Uncomment for training visualization
EOL
        log "Requirements file created at $REQUIREMENTS_FILE"
    else
        log "Requirements file already exists, keeping existing file"
    fi
}

# Main installation function
main() {
    echo -e "${GREEN}===============================================${NC}"
    echo -e "${GREEN}  ZETA Components Setup for MCP-ZERO${NC}"
    echo -e "${GREEN}  Enterprise-Grade with Offline-First Resilience${NC}"
    echo -e "${GREEN}===============================================${NC}"
    
    log "Starting ZETA components setup..." "important"
    log "MCPZERO_ROOT: $MCPZERO_ROOT" "important"
    log "Offline mode: $OFFLINE_MODE" "important"
    
    # Create directories
    create_directories
    
    # Create requirements file
    create_requirements
    
    # Setup Python environment
    setup_python_env
    
    # Setup ZETACUDA
    if [ "$INSTALL_CUDA" = true ]; then
        setup_zetacuda
    fi
    
    # Setup environment variables
    setup_env_vars
    
    echo -e "${GREEN}===============================================${NC}"
    echo -e "${GREEN}  ZETA Components Setup Complete${NC}"
    echo -e "${GREEN}===============================================${NC}"
    echo ""
    echo -e "To activate the ZETA environment, run:"
    echo -e "  ${YELLOW}source ${MCPZERO_ROOT}/.mcpzero/activate_zeta.sh${NC}"
    echo ""
    echo -e "Quick start commands:"
    echo -e "  ${YELLOW}python -m zeta_integration${NC} - Test ZETA integration"
    echo -e "  ${YELLOW}python -m zetaface.zetaface${NC} - Test ZETAFACE API"
    echo ""
}

# Run main function
main
