#!/bin/bash

# CHORUS Distributed AI System Installer
# One-command installation script for the CHORUS cluster
# 
# Usage:
#   curl -fsSL https://chorus.services/install.sh | sh
#   
# Or with options:
#   curl -fsSL https://chorus.services/install.sh | sh -s -- --coordinator --models "llama3.2,qwen2.5:7b"

set -e

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

# ASCII Art Banner
print_banner() {
    echo -e "${PURPLE}"
    cat << 'EOF'
     ██████╗██╗  ██╗ ██████╗ ██████╗ ██╗   ██╗███████╗
    ██╔════╝██║  ██║██╔═══██╗██╔══██╗██║   ██║██╔════╝
    ██║     ███████║██║   ██║██████╔╝██║   ██║███████╗
    ██║     ██╔══██║██║   ██║██╔══██╗██║   ██║╚════██║
    ╚██████╗██║  ██║╚██████╔╝██║  ██║╚██████╔╝███████║
     ╚═════╝╚═╝  ╚═╝ ╚═════╝ ╚═╝  ╚═╝ ╚═════╝ ╚══════╝
                                                        
    🎵 Distributed AI Orchestration Platform
EOF
    echo -e "${NC}"
    echo -e "${CYAN}    A living knowledge fabric for collaborative AI agents${NC}"
    echo ""
}

# Configuration
CHORUS_VERSION="latest"
CHORUS_USER="$USER"
CHORUS_HOME="$HOME/chorus"
BZZZ_DIR="$CHORUS_HOME/project-queues/active/BZZZ"
WHOOSH_DIR="$CHORUS_HOME/project-queues/active/WHOOSH"
GITEA_URL="https://gitea.deepblack.cloud"
REGISTRY_URL="registry.home.deepblack.cloud"

# Installation flags
IS_COORDINATOR=false
INSTALL_MODELS=""
SKIP_DEPENDENCIES=false
QUIET_MODE=false

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --coordinator)
            IS_COORDINATOR=true
            shift
            ;;
        --models)
            INSTALL_MODELS="$2"
            shift 2
            ;;
        --skip-deps)
            SKIP_DEPENDENCIES=true
            shift
            ;;
        --quiet)
            QUIET_MODE=true
            shift
            ;;
        --help)
            echo "CHORUS Installation Script"
            echo ""
            echo "Usage: $0 [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --coordinator     Install as cluster coordinator node"
            echo "  --models MODEL    Comma-separated list of models to install"
            echo "  --skip-deps       Skip dependency installation"
            echo "  --quiet           Minimal output"
            echo "  --help            Show this help"
            echo ""
            echo "Examples:"
            echo "  $0 --coordinator --models \"llama3.2,qwen2.5:7b\""
            echo "  $0 --models \"codellama:7b\""
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

# Logging functions
log_info() {
    if [[ "$QUIET_MODE" != "true" ]]; then
        echo -e "${GREEN}✓${NC} $1"
    fi
}

log_warn() {
    echo -e "${YELLOW}⚠${NC} $1"
}

log_error() {
    echo -e "${RED}✗${NC} $1"
}

log_step() {
    if [[ "$QUIET_MODE" != "true" ]]; then
        echo -e "${BLUE}▶${NC} $1"
    fi
}

# System detection
detect_system() {
    log_step "Detecting system..."
    
    if [[ "$OSTYPE" == "linux-gnu"* ]]; then
        if command -v apt-get >/dev/null 2>&1; then
            PACKAGE_MANAGER="apt"
            INSTALL_CMD="sudo apt-get update && sudo apt-get install -y"
        elif command -v yum >/dev/null 2>&1; then
            PACKAGE_MANAGER="yum"
            INSTALL_CMD="sudo yum install -y"
        elif command -v pacman >/dev/null 2>&1; then
            PACKAGE_MANAGER="pacman"
            INSTALL_CMD="sudo pacman -S --noconfirm"
        else
            log_error "Unsupported Linux distribution"
            exit 1
        fi
    elif [[ "$OSTYPE" == "darwin"* ]]; then
        PACKAGE_MANAGER="brew"
        INSTALL_CMD="brew install"
    else
        log_error "Unsupported operating system: $OSTYPE"
        exit 1
    fi
    
    log_info "Detected system: $OSTYPE with $PACKAGE_MANAGER"
}

# Check prerequisites
check_prerequisites() {
    log_step "Checking prerequisites..."
    
    local missing_deps=()
    
    # Check for required commands
    for cmd in curl git; do
        if ! command -v $cmd >/dev/null 2>&1; then
            missing_deps+=($cmd)
        fi
    done
    
    if [[ ${#missing_deps[@]} -gt 0 ]] && [[ "$SKIP_DEPENDENCIES" != "true" ]]; then
        log_step "Installing missing dependencies: ${missing_deps[*]}"
        case $PACKAGE_MANAGER in
            apt)
                $INSTALL_CMD ${missing_deps[*]}
                ;;
            yum)
                $INSTALL_CMD ${missing_deps[*]}
                ;;
            pacman)
                $INSTALL_CMD ${missing_deps[*]}
                ;;
            brew)
                brew install ${missing_deps[*]}
                ;;
        esac
    fi
    
    log_info "Prerequisites satisfied"
}


# Install Ollama
install_ollama() {
    if command -v ollama >/dev/null 2>&1; then
        log_info "Ollama already installed"
        return
    fi
    
    log_step "Installing Ollama..."
    curl -fsSL https://ollama.com/install.sh | sh
    
    # Start Ollama service
    sudo systemctl enable ollama || true
    sudo systemctl start ollama || true
    
    log_info "Ollama installed and started"
}

# Setup CHORUS directory structure
setup_chorus_structure() {
    log_step "Setting up CHORUS directory structure..."
    
    mkdir -p "$CHORUS_HOME"/{docs,business/secrets,project-queues/active}
    mkdir -p "$HOME/.chorus"/{logs,config,cache}
    
    log_info "Directory structure created"
}

# Download CHORUS binaries
download_binaries() {
    log_step "Downloading CHORUS binaries..."
    
    # Create BZZZ directory structure
    mkdir -p "$BZZZ_DIR"/{config,logs}
    cd "$BZZZ_DIR"
    
    # Detect system architecture
    local arch=""
    case $(uname -m) in
        x86_64) arch="amd64" ;;
        aarch64|arm64) arch="arm64" ;;
        armv7l) arch="armv7" ;;
        *) 
            log_error "Unsupported architecture: $(uname -m)"
            exit 1
            ;;
    esac
    
    local os="linux"
    if [[ "$OSTYPE" == "darwin"* ]]; then
        os="darwin"
    fi
    
    # Download BZZZ binary
    local bzzz_url="https://chorus.services/releases/bzzz/latest/bzzz-${os}-${arch}"
    log_step "Downloading BZZZ binary for ${os}-${arch}..."
    
    if curl -fsSL -o bzzz "$bzzz_url"; then
        chmod +x bzzz
        log_info "BZZZ binary downloaded successfully"
    else
        log_error "Failed to download BZZZ binary from $bzzz_url"
        log_info "Falling back to minimal installation..."
        # Create a placeholder binary that will be updated later
        echo '#!/bin/bash' > bzzz
        echo 'echo "BZZZ binary not available for this platform"' >> bzzz
        chmod +x bzzz
    fi
    
    # Download systemd service file
    local service_url="https://chorus.services/releases/bzzz/latest/bzzz.service"
    if curl -fsSL -o bzzz.service "$service_url"; then
        log_info "BZZZ service file downloaded"
    else
        log_warn "Failed to download service file, creating default..."
        create_default_service_file
    fi
    
    # Download install script
    local install_script_url="https://chorus.services/releases/bzzz/latest/install-service.sh"
    if curl -fsSL -o install-service.sh "$install_script_url"; then
        chmod +x install-service.sh
        log_info "Install script downloaded"
    else
        log_warn "Failed to download install script, creating default..."
        create_default_install_script
    fi
    
    # Download WHOOSH if coordinator
    if [[ "$IS_COORDINATOR" == "true" ]]; then
        log_step "Downloading WHOOSH dashboard..."
        mkdir -p "$WHOOSH_DIR"
        
        local whoosh_url="https://chorus.services/releases/whoosh/latest/whoosh-docker-compose.tar.gz"
        if curl -fsSL "$whoosh_url" | tar -xzf - -C "$WHOOSH_DIR" 2>/dev/null; then
            log_info "WHOOSH dashboard downloaded"
        else
            log_warn "Failed to download WHOOSH dashboard"
        fi
    fi
    
    log_info "Binary downloads completed"
}

# Create default service file
create_default_service_file() {
    cat > bzzz.service << 'EOF'
[Unit]
Description=BZZZ P2P Task Coordination System
Documentation=https://chorus.services/docs/bzzz
After=network.target
Wants=network.target

[Service]
Type=simple
User=PLACEHOLDER_USER
Group=PLACEHOLDER_USER
WorkingDirectory=PLACEHOLDER_BZZZ_DIR
ExecStart=PLACEHOLDER_BZZZ_DIR/bzzz
Restart=always
RestartSec=10
KillMode=mixed
KillSignal=SIGTERM
TimeoutStopSec=30

# Environment variables
Environment=HOME=PLACEHOLDER_HOME
Environment=USER=PLACEHOLDER_USER

# Logging
StandardOutput=journal
StandardError=journal
SyslogIdentifier=bzzz

# Security settings
NoNewPrivileges=true
PrivateTmp=true
ProtectSystem=strict
ProtectHome=false
ReadWritePaths=PLACEHOLDER_BZZZ_DIR

# Resource limits
LimitNOFILE=65536
LimitNPROC=4096

[Install]
WantedBy=multi-user.target
EOF
    
    # Replace placeholders
    sed -i "s|PLACEHOLDER_USER|$USER|g" bzzz.service
    sed -i "s|PLACEHOLDER_HOME|$HOME|g" bzzz.service
    sed -i "s|PLACEHOLDER_BZZZ_DIR|$BZZZ_DIR|g" bzzz.service
    
    log_info "Default service file created"
}

# Create default install script
create_default_install_script() {
    cat > install-service.sh << 'EOF'
#!/bin/bash

# BZZZ Service Installation Script
# Installs BZZZ as a systemd service

set -e

echo "🐝 Installing BZZZ P2P Task Coordination Service..."

# Check if running as root or with sudo
if [ "$EUID" -ne 0 ]; then
    echo "❌ This script must be run as root or with sudo"
    exit 1
fi

# Define paths
BZZZ_DIR="$(pwd)"
SERVICE_FILE="$BZZZ_DIR/bzzz.service"
SYSTEMD_DIR="/etc/systemd/system"

# Check if BZZZ binary exists
if [ ! -f "$BZZZ_DIR/bzzz" ]; then
    echo "❌ BZZZ binary not found at $BZZZ_DIR/bzzz"
    exit 1
fi

# Make binary executable
chmod +x "$BZZZ_DIR/bzzz"
echo "✅ Made BZZZ binary executable"

# Copy service file to systemd directory
cp "$SERVICE_FILE" "$SYSTEMD_DIR/bzzz.service"
echo "✅ Copied service file to $SYSTEMD_DIR/bzzz.service"

# Set proper permissions
chmod 644 "$SYSTEMD_DIR/bzzz.service"
echo "✅ Set service file permissions"

# Reload systemd daemon
systemctl daemon-reload
echo "✅ Reloaded systemd daemon"

# Enable service to start on boot
systemctl enable bzzz.service
echo "✅ Enabled BZZZ service for auto-start"

# Start the service
systemctl start bzzz.service
echo "✅ Started BZZZ service"

echo ""
echo "🎉 BZZZ P2P Task Coordination Service installed successfully!"
EOF
    
    chmod +x install-service.sh
    log_info "Default install script created"
}

# Configure BZZZ
configure_bzzz() {
    log_step "Configuring BZZZ..."
    
    # Interactive configuration
    echo ""
    echo -e "${CYAN}BZZZ Node Configuration${NC}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━"
    
    read -p "Node hostname [$(hostname)]: " node_hostname
    node_hostname=${node_hostname:-$(hostname)}
    
    read -p "Node role [worker]: " node_role
    node_role=${node_role:-worker}
    
    if [[ "$IS_COORDINATOR" == "true" ]]; then
        node_role="coordinator"
        echo "Setting role to coordinator (--coordinator flag provided)"
    fi
    
    read -p "API port [8080]: " api_port
    api_port=${api_port:-8080}
    
    read -p "P2P port [4001]: " p2p_port
    p2p_port=${p2p_port:-4001}
    
    # Generate configuration
    cat > "$BZZZ_DIR/config/bzzz.json" << EOF
{
  "node": {
    "id": "$node_hostname",
    "role": "$node_role"
  },
  "api": {
    "host": "0.0.0.0",
    "port": $api_port
  },
  "p2p": {
    "port": $p2p_port,
    "discovery": {
      "enabled": true,
      "bootstrap_nodes": []
    }
  },
  "coordinator": $([ "$IS_COORDINATOR" == "true" ] && echo "true" || echo "false"),
  "logging": {
    "level": "info",
    "file": "$HOME/.chorus/logs/bzzz.log"
  },
  "storage": {
    "path": "$HOME/.chorus/data"
  }
}
EOF
    
    log_info "BZZZ configured"
}

# Install BZZZ service
install_bzzz_service() {
    log_step "Installing BZZZ as systemd service..."
    
    cd "$BZZZ_DIR"
    
    # Check if BZZZ binary exists
    if [[ ! -f "bzzz" ]]; then
        log_error "BZZZ binary not found - cannot install service"
        return 1
    fi
    
    # Check if install script exists
    if [[ ! -f "install-service.sh" ]]; then
        log_error "install-service.sh not found - cannot install service"
        return 1
    fi
    
    # Make sure install script is executable
    chmod +x install-service.sh
    
    # Install the service
    if sudo ./install-service.sh; then
        log_info "BZZZ service installed and started"
        return 0
    else
        log_error "Failed to install BZZZ service"
        return 1
    fi
}

# Install AI models
install_models() {
    if [[ -z "$INSTALL_MODELS" ]]; then
        return
    fi
    
    log_step "Installing AI models..."
    
    IFS=',' read -ra MODELS <<< "$INSTALL_MODELS"
    for model in "${MODELS[@]}"; do
        model=$(echo "$model" | xargs) # trim whitespace
        log_step "Pulling model: $model"
        ollama pull "$model"
        log_info "Model $model installed"
    done
}

# Setup coordinator-specific components
setup_coordinator() {
    if [[ "$IS_COORDINATOR" != "true" ]]; then
        return
    fi
    
    log_step "Setting up coordinator components..."
    
    # Install Docker for WHOOSH
    if ! command -v docker >/dev/null 2>&1; then
        curl -fsSL https://get.docker.com | sh
        sudo usermod -aG docker "$USER"
        log_info "Docker installed"
    fi
    
    # Generate Age encryption keys
    if ! command -v age >/dev/null 2>&1; then
        # Install age
        wget -O /tmp/age.tar.gz https://github.com/FiloSottile/age/releases/download/v1.1.1/age-v1.1.1-linux-amd64.tar.gz
        tar -xzf /tmp/age.tar.gz -C /tmp
        sudo cp /tmp/age/age* /usr/local/bin/
        sudo chmod +x /usr/local/bin/age*
        rm -rf /tmp/age*
    fi
    
    # Generate keys
    mkdir -p "$CHORUS_HOME/business/secrets"
    age-keygen -o "$CHORUS_HOME/business/secrets/age-key.txt"
    
    log_info "Coordinator setup completed"
}

# Show completion message
show_completion() {
    echo ""
    echo -e "${GREEN}🎉 CHORUS Installation Complete!${NC}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    echo -e "${WHITE}Services Status:${NC}"
    local bzzz_status=$(systemctl is-active bzzz 2>/dev/null || echo 'inactive')
    echo "  • BZZZ Agent: $bzzz_status"
    echo "  • Ollama: $(systemctl is-active ollama 2>/dev/null || echo 'inactive')"
    
    if [[ "$bzzz_status" == "inactive" ]]; then
        echo ""
        echo -e "${YELLOW}⚠ BZZZ Service Issues:${NC}"
        echo "  The BZZZ service may not be running."
        echo "  This could be due to configuration or binary compatibility issues."
        echo ""
        echo -e "${WHITE}Troubleshooting:${NC}"
        echo "  1. Check logs: sudo journalctl -u bzzz -f"
        echo "  2. Test binary: $BZZZ_DIR/bzzz --version"
        echo "  3. Restart service: sudo systemctl restart bzzz"
        echo "  4. Check config: $BZZZ_DIR/config/bzzz.json"
    fi
    echo ""
    echo -e "${WHITE}Useful Commands:${NC}"
    echo "  sudo systemctl status bzzz     - Check BZZZ status"
    echo "  sudo journalctl -u bzzz -f     - Follow BZZZ logs"
    echo "  ollama list                    - List installed models"
    echo "  curl http://localhost:8080/health - Test BZZZ API"
    echo ""
    if [[ "$IS_COORDINATOR" == "true" ]]; then
        echo -e "${WHITE}Coordinator Setup:${NC}"
        echo "  • Age encryption keys: $CHORUS_HOME/business/secrets/age-key.txt"
        echo "  • WHOOSH dashboard will be available after Docker deployment"
        echo ""
    fi
    echo -e "${WHITE}Documentation:${NC}"
    echo "  • CHORUS Docs: $CHORUS_HOME/docs/"
    echo "  • BZZZ Config: $BZZZ_DIR/config/bzzz.json"
    echo ""
    echo -e "${CYAN}Next Steps:${NC}"
    if [[ "$IS_COORDINATOR" == "true" ]]; then
        echo "  1. Deploy WHOOSH dashboard: cd $WHOOSH_DIR && docker-compose up -d"
        echo "  2. Access dashboard at: http://$(hostname):3001"
        echo "  3. Add worker nodes by running this script on other machines"
    else
        echo "  1. Configure coordinator endpoint in BZZZ config"
        echo "  2. Restart BZZZ: sudo systemctl restart bzzz"
        echo "  3. Check connection: curl http://localhost:8080/health"
    fi
    echo ""
}

# Error handling
cleanup_on_error() {
    log_error "Installation failed. Cleaning up..."
    # Cleanup logic here
    exit 1
}

# Disable strict error handling for controlled failure management
set +e

# Main installation flow
main() {
    print_banner
    
    echo -e "${WHITE}Starting CHORUS installation...${NC}"
    echo ""
    
    detect_system
    check_prerequisites
    install_ollama
    setup_chorus_structure
    download_binaries
    configure_bzzz
    install_bzzz_service
    
    install_models
    setup_coordinator
    
    show_completion
}

# Run main installation
main "$@"