#!/bin/bash
# BZZZ Cluster Installation Script
# Usage: curl -fsSL https://chorus.services/install.sh | sh

set -euo pipefail

# Configuration
BZZZ_VERSION="${BZZZ_VERSION:-latest}"
BZZZ_BASE_URL="${BZZZ_BASE_URL:-https://chorus.services}"
BZZZ_INSTALL_DIR="${BZZZ_INSTALL_DIR:-/opt/bzzz}"
BZZZ_CONFIG_DIR="${BZZZ_CONFIG_DIR:-/etc/bzzz}"
BZZZ_LOG_DIR="${BZZZ_LOG_DIR:-/var/log/bzzz}"
BZZZ_DATA_DIR="${BZZZ_DATA_DIR:-/var/lib/bzzz}"
INSTALL_PARALLAMA="${INSTALL_PARALLAMA:-prompt}"  # prompt, yes, no

# Colors 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_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Error handler
error_exit() {
    log_error "$1"
    exit 1
}

# Check if running as root
check_root() {
    if [[ $EUID -eq 0 ]]; then
        log_warn "Running as root. BZZZ will be installed system-wide."
    else
        log_info "Running as non-root user. Some features may require sudo access."
    fi
}

# Detect operating system
detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS=$ID
        OS_VERSION=$VERSION_ID
    elif [[ -f /etc/redhat-release ]]; then
        OS="centos"
    elif [[ -f /etc/debian_version ]]; then
        OS="debian"
    else
        error_exit "Unsupported operating system"
    fi
    
    log_info "Detected OS: $OS $OS_VERSION"
}

# Detect system architecture
detect_arch() {
    ARCH=$(uname -m)
    case $ARCH in
        x86_64)
            ARCH="amd64"
            ;;
        aarch64|arm64)
            ARCH="arm64"
            ;;
        armv7l)
            ARCH="armv7"
            ;;
        *)
            error_exit "Unsupported architecture: $ARCH"
            ;;
    esac
    
    log_info "Detected architecture: $ARCH"
}

# Check system requirements
check_requirements() {
    log_info "Checking system requirements..."
    
    # Check minimum memory (4GB recommended)
    local mem_kb=$(grep MemTotal /proc/meminfo | awk '{print $2}')
    local mem_gb=$((mem_kb / 1024 / 1024))
    
    if [[ $mem_gb -lt 2 ]]; then
        error_exit "Insufficient memory. Minimum 2GB required, 4GB recommended."
    elif [[ $mem_gb -lt 4 ]]; then
        log_warn "Memory is below recommended 4GB ($mem_gb GB available)"
    fi
    
    # Check disk space (minimum 10GB)
    local disk_free=$(df / | awk 'NR==2 {print $4}')
    local disk_gb=$((disk_free / 1024 / 1024))
    
    if [[ $disk_gb -lt 10 ]]; then
        error_exit "Insufficient disk space. Minimum 10GB free space required."
    fi
    
    log_success "System requirements check passed"
}

# Install system dependencies
install_dependencies() {
    log_info "Installing system dependencies..."
    
    case $OS in
        ubuntu|debian)
            sudo apt-get update -qq
            sudo apt-get install -y \
                curl \
                wget \
                gnupg \
                lsb-release \
                ca-certificates \
                software-properties-common \
                apt-transport-https \
                jq \
                net-tools \
                openssh-client \
                docker.io \
                docker-compose
            ;;
        centos|rhel|fedora)
            sudo yum update -y
            sudo yum install -y \
                curl \
                wget \
                gnupg \
                ca-certificates \
                jq \
                net-tools \
                openssh-clients \
                docker \
                docker-compose
            ;;
        *)
            error_exit "Package installation not supported for OS: $OS"
            ;;
    esac
    
    # Ensure Docker is running
    sudo systemctl enable docker
    sudo systemctl start docker
    
    # Add current user to docker group if not root
    if [[ $EUID -ne 0 ]]; then
        sudo usermod -aG docker $USER
        log_warn "Added $USER to docker group. You may need to logout and login again."
    fi
    
    log_success "Dependencies installed successfully"
}

# Detect GPU configuration
detect_gpu() {
    log_info "Detecting GPU configuration..."
    
    GPU_COUNT=0
    GPU_TYPE="none"
    
    # Check for NVIDIA GPUs
    if command -v nvidia-smi &>/dev/null; then
        GPU_COUNT=$(nvidia-smi --list-gpus 2>/dev/null | wc -l || echo 0)
        if [[ $GPU_COUNT -gt 0 ]]; then
            GPU_TYPE="nvidia"
            log_info "Detected $GPU_COUNT NVIDIA GPU(s)"
        fi
    fi
    
    # Check for AMD GPUs
    if [[ $GPU_COUNT -eq 0 ]] && command -v rocm-smi &>/dev/null; then
        GPU_COUNT=$(rocm-smi --showid 2>/dev/null | grep -c "GPU" || echo 0)
        if [[ $GPU_COUNT -gt 0 ]]; then
            GPU_TYPE="amd"
            log_info "Detected $GPU_COUNT AMD GPU(s)"
        fi
    fi
    
    if [[ $GPU_COUNT -eq 0 ]]; then
        log_info "No GPUs detected - CPU-only mode"
    fi
    
    export GPU_COUNT GPU_TYPE
}

# Prompt for Parallama installation
prompt_parallama_installation() {
    if [[ $INSTALL_PARALLAMA == "prompt" ]]; then
        echo
        log_info "BZZZ can optionally install Parallama (multi-GPU Ollama fork) for enhanced AI capabilities."
        echo
        
        if [[ $GPU_COUNT -gt 1 ]]; then
            echo -e "${GREEN}🚀 Multi-GPU Setup Detected ($GPU_COUNT ${GPU_TYPE^^} GPUs)${NC}"
            echo "   Parallama is RECOMMENDED for optimal multi-GPU performance!"
        elif [[ $GPU_COUNT -eq 1 ]]; then
            echo -e "${YELLOW}🎯 Single GPU Detected (${GPU_TYPE^^})${NC}"
            echo "   Parallama provides enhanced GPU utilization."
        else
            echo -e "${BLUE}💻 CPU-Only Setup${NC}"
            echo "   Parallama can still provide CPU optimizations."
        fi
        
        echo
        echo "Options:"
        echo "1. Install Parallama (recommended for GPU setups)"
        echo "2. Install standard Ollama"  
        echo "3. Skip Ollama installation (configure later)"
        echo
        
        read -p "Choose option (1-3): " choice
        
        case $choice in
            1)
                INSTALL_PARALLAMA="yes"
                ;;
            2)
                INSTALL_PARALLAMA="no"
                ;;
            3)
                INSTALL_PARALLAMA="skip"
                ;;
            *)
                log_warn "Invalid choice, defaulting to Parallama"
                INSTALL_PARALLAMA="yes"
                ;;
        esac
    fi
}

# Install Ollama or Parallama
install_ollama() {
    if [[ $INSTALL_PARALLAMA == "skip" ]]; then
        log_info "Skipping Ollama installation"
        return
    fi
    
    if [[ $INSTALL_PARALLAMA == "yes" ]]; then
        log_info "Installing Parallama (multi-GPU Ollama fork)..."
        
        # Download Parallama installer
        if ! curl -fsSL https://chorus.services/parallama/install.sh | sh; then
            log_error "Failed to install Parallama, falling back to standard Ollama"
            install_standard_ollama
        else
            log_success "Parallama installed successfully"
            
            # Configure Parallama for multi-GPU if available
            if [[ $GPU_COUNT -gt 1 ]]; then
                log_info "Configuring Parallama for $GPU_COUNT GPUs..."
                # Parallama will be configured via the web UI
            fi
        fi
    else
        install_standard_ollama
    fi
}

# Install standard Ollama
install_standard_ollama() {
    log_info "Installing standard Ollama..."
    
    if ! curl -fsSL https://ollama.ai/install.sh | sh; then
        log_warn "Failed to install Ollama - you can install it later via the web UI"
    else
        log_success "Ollama installed successfully"
    fi
}

# Download and install BZZZ binaries
install_bzzz_binaries() {
    log_info "Downloading BZZZ binaries..."
    
    local download_url="${BZZZ_BASE_URL}/releases/${BZZZ_VERSION}/bzzz-${OS}-${ARCH}.tar.gz"
    local temp_dir=$(mktemp -d)
    
    # Download binary package
    if ! curl -fsSL "$download_url" -o "$temp_dir/bzzz.tar.gz"; then
        error_exit "Failed to download BZZZ binaries from $download_url"
    fi
    
    # Extract binaries
    sudo mkdir -p "$BZZZ_INSTALL_DIR"
    sudo tar -xzf "$temp_dir/bzzz.tar.gz" -C "$BZZZ_INSTALL_DIR"
    
    # Make binaries executable
    sudo chmod +x "$BZZZ_INSTALL_DIR"/bin/*
    
    # Create symlinks
    sudo ln -sf "$BZZZ_INSTALL_DIR/bin/bzzz" /usr/local/bin/bzzz
    sudo ln -sf "$BZZZ_INSTALL_DIR/bin/bzzz-mcp" /usr/local/bin/bzzz-mcp
    
    # Cleanup
    rm -rf "$temp_dir"
    
    log_success "BZZZ binaries installed successfully"
}

# Setup configuration directories
setup_directories() {
    log_info "Setting up directories..."
    
    sudo mkdir -p "$BZZZ_CONFIG_DIR"
    sudo mkdir -p "$BZZZ_LOG_DIR"
    sudo mkdir -p "$BZZZ_DATA_DIR"
    
    # Set permissions
    local bzzz_user="bzzz"
    
    # Create bzzz user if not exists
    if ! id "$bzzz_user" &>/dev/null; then
        sudo useradd -r -s /bin/false -d "$BZZZ_DATA_DIR" "$bzzz_user"
    fi
    
    sudo chown -R "$bzzz_user:$bzzz_user" "$BZZZ_CONFIG_DIR"
    sudo chown -R "$bzzz_user:$bzzz_user" "$BZZZ_LOG_DIR"
    sudo chown -R "$bzzz_user:$bzzz_user" "$BZZZ_DATA_DIR"
    
    log_success "Directories created successfully"
}

# Install systemd services
install_services() {
    log_info "Installing systemd services..."
    
    # BZZZ Go service
    sudo tee /etc/systemd/system/bzzz.service > /dev/null <<EOF
[Unit]
Description=BZZZ Distributed AI Coordination Service
Documentation=https://docs.chorus.services/bzzz
After=network-online.target docker.service
Wants=network-online.target
Requires=docker.service

[Service]
Type=simple
User=bzzz
Group=bzzz
WorkingDirectory=$BZZZ_DATA_DIR
Environment=BZZZ_CONFIG_DIR=$BZZZ_CONFIG_DIR
ExecStart=$BZZZ_INSTALL_DIR/bin/bzzz server --config $BZZZ_CONFIG_DIR/bzzz.yaml
ExecReload=/bin/kill -HUP \$MAINPID
Restart=always
RestartSec=10
KillMode=mixed
KillSignal=SIGTERM
TimeoutSec=30
StandardOutput=journal
StandardError=journal
SyslogIdentifier=bzzz

[Install]
WantedBy=multi-user.target
EOF

    # BZZZ MCP service
    sudo tee /etc/systemd/system/bzzz-mcp.service > /dev/null <<EOF
[Unit]
Description=BZZZ MCP Server for GPT-5 Integration
Documentation=https://docs.chorus.services/bzzz/mcp
After=network-online.target bzzz.service
Wants=network-online.target
Requires=bzzz.service

[Service]
Type=simple
User=bzzz
Group=bzzz
WorkingDirectory=$BZZZ_DATA_DIR
Environment=NODE_ENV=production
EnvironmentFile=-$BZZZ_CONFIG_DIR/mcp.env
ExecStart=$BZZZ_INSTALL_DIR/bin/bzzz-mcp
Restart=always
RestartSec=10
KillMode=mixed
KillSignal=SIGTERM
TimeoutSec=30
StandardOutput=journal
StandardError=journal
SyslogIdentifier=bzzz-mcp

[Install]
WantedBy=multi-user.target
EOF

    # Reload systemd
    sudo systemctl daemon-reload
    
    log_success "Systemd services installed"
}

# Generate initial configuration
generate_config() {
    log_info "Generating initial configuration..."
    
    # Detect network interface and IP
    local primary_interface=$(ip route | grep default | awk '{print $5}' | head -n1)
    local primary_ip=$(ip addr show "$primary_interface" | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1 | head -n1)
    local subnet=$(ip route | grep "$primary_interface" | grep '/' | head -n1 | awk '{print $1}')
    
    # Generate node ID
    local node_id="node-$(hostname -s)-$(date +%s)"
    
    # Create basic configuration
    sudo tee "$BZZZ_CONFIG_DIR/bzzz.yaml" > /dev/null <<EOF
# BZZZ Configuration - Generated by install script
# Complete configuration via web UI at http://$primary_ip:8080/setup

node:
  id: "$node_id"
  name: "$(hostname -s)"
  address: "$primary_ip"
  
network:
  listen_port: 8080
  discovery_port: 8081
  subnet: "$subnet"
  interface: "$primary_interface"
  
cluster:
  auto_discovery: true
  bootstrap_nodes: []
  
services:
  mcp_server:
    enabled: true
    port: 3000
  web_ui:
    enabled: true
    port: 8080
    
security:
  tls:
    enabled: false  # Will be configured via web UI
  auth:
    enabled: false  # Will be configured via web UI
    
logging:
  level: "info"
  file: "$BZZZ_LOG_DIR/bzzz.log"
  
# Hardware configuration - detected during installation
hardware:
  cpu_cores: $(nproc)
  memory_gb: $mem_gb
  gpus:
    count: $GPU_COUNT
    type: "$GPU_TYPE"
    
# Ollama/Parallama configuration
ollama:
  enabled: $(if [[ $INSTALL_PARALLAMA != "skip" ]]; then echo "true"; else echo "false"; fi)
  type: "$(if [[ $INSTALL_PARALLAMA == "yes" ]]; then echo "parallama"; else echo "ollama"; fi)"
  endpoint: "http://localhost:11434"
  models: []  # Will be configured via web UI
  
# Placeholder configurations - set via web UI
openai:
  api_key: ""
  model: "gpt-5"
  
cost_limits:
  daily: 100.0
  monthly: 1000.0
EOF

    log_success "Initial configuration generated"
}

# Start configuration web server
start_config_server() {
    log_info "Starting configuration server..."
    
    # Start BZZZ service for configuration
    sudo systemctl enable bzzz
    sudo systemctl start bzzz
    
    # Wait for service to be ready
    local retries=30
    local primary_ip=$(ip addr show $(ip route | grep default | awk '{print $5}' | head -n1) | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1 | head -n1)
    
    while [[ $retries -gt 0 ]]; do
        if curl -f "http://$primary_ip:8080/health" &>/dev/null; then
            break
        fi
        sleep 2
        ((retries--))
    done
    
    if [[ $retries -eq 0 ]]; then
        log_warn "Configuration server may not be ready. Check logs with: sudo journalctl -u bzzz -f"
    fi
    
    log_success "Configuration server started"
}

# Display completion message
show_completion_message() {
    local primary_ip=$(ip addr show $(ip route | grep default | awk '{print $5}' | head -n1) | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1 | head -n1)
    
    echo
    log_success "BZZZ installation completed successfully!"
    echo
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo
    echo -e "${GREEN}🚀 Next Steps:${NC}"
    echo
    echo "1. Complete your cluster configuration:"
    echo "   👉 Open: ${BLUE}http://$primary_ip:8080/setup${NC}"
    echo
    echo "2. Useful commands:"
    echo "   • Check status:      ${YELLOW}bzzz status${NC}"
    echo "   • View logs:         ${YELLOW}sudo journalctl -u bzzz -f${NC}"
    echo "   • Start/Stop:        ${YELLOW}sudo systemctl [start|stop] bzzz${NC}"
    echo "   • Configuration:     ${YELLOW}sudo nano $BZZZ_CONFIG_DIR/bzzz.yaml${NC}"
    echo
    echo "3. Documentation:"
    echo "   📚 Docs: ${BLUE}https://docs.chorus.services/bzzz${NC}"
    echo "   💬 Support: ${BLUE}https://discord.gg/chorus-services${NC}"
    echo
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo
}

# Cleanup function for error handling
cleanup() {
    if [[ -n "${temp_dir:-}" ]] && [[ -d "$temp_dir" ]]; then
        rm -rf "$temp_dir"
    fi
}
trap cleanup EXIT

# Main installation flow
main() {
    echo
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo -e "${GREEN}🔥 BZZZ Distributed AI Coordination Platform${NC}"
    echo "   Installer v1.0"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo
    
    check_root
    detect_os
    detect_arch
    check_requirements
    detect_gpu
    install_dependencies
    prompt_parallama_installation
    install_ollama
    install_bzzz_binaries
    setup_directories
    install_services
    generate_config
    start_config_server
    show_completion_message
}

# Run main installation
main "$@"