#!/bin/bash
# MCP-ZERO Installation Script
# This script installs MCP-ZERO and all its dependencies

set -e

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

# Constants
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEFAULT_INSTALL_DIR="$HOME/mcp-zero"

# Print colored message
print_message() {
  local color=$1
  local message=$2
  echo -e "${color}${message}${NC}"
}

# Print usage information
print_usage() {
  echo "Usage: $0 [OPTIONS]"
  echo
  echo "Options:"
  echo "  -h, --help                   Show this help message"
  echo "  --install-dir DIR            MCP-ZERO installation directory (default: $DEFAULT_INSTALL_DIR)"
  echo "  --force                      Overwrite existing installation"
  echo
  echo "Example:"
  echo "  $0 --install-dir ~/my-mcp --force"
  exit 1
}

# Parse command line arguments
parse_arguments() {
  INSTALL_DIR="$DEFAULT_INSTALL_DIR"
  FORCE=0
  
  while [[ $# -gt 0 ]]; do
    case $1 in
      -h|--help)
        print_usage
        ;;
      --install-dir)
        INSTALL_DIR="$2"
        shift 2
        ;;
      --force)
        FORCE=1
        shift
        ;;
      *)
        print_message "$RED" "Error: Unknown option $1"
        print_usage
        ;;
    esac
  done
}

# Check for required dependencies
check_dependencies() {
  print_message "$BLUE" "Checking dependencies..."
  
  MISSING_DEPS=0
  
  # Check for essential tools
  for cmd in git python3 pip3 docker; do
    if ! command -v ${cmd} &> /dev/null; then
      print_message "$RED" "Error: Required tool '${cmd}' is not installed."
      MISSING_DEPS=1
    fi
  done
  
  if [ ${MISSING_DEPS} -eq 1 ]; then
    print_message "$RED" "Please install missing dependencies and try again."
    exit 1
  fi
  
  print_message "$GREEN" "All dependencies satisfied."
}

# Install Docker if not already installed
install_docker() {
  if ! command -v docker &> /dev/null; then
    print_message "$BLUE" "Installing Docker..."
    
    # Install Docker using official script
    curl -fsSL https://get.docker.com -o get-docker.sh
    sudo sh get-docker.sh
    
    # Add current user to docker group
    sudo usermod -aG docker $USER
    
    print_message "$YELLOW" "Docker installed. You may need to log out and back in for group changes to take effect."
  else
    print_message "$GREEN" "Docker already installed."
  fi
}

# Install MCP-ZERO
install_mcp_zero() {
  print_message "$BLUE" "Installing MCP-ZERO to $INSTALL_DIR..."
  
  # Check if directory exists and is not empty
  if [ -d "$INSTALL_DIR" ] && [ "$(ls -A "$INSTALL_DIR" 2>/dev/null)" ]; then
    if [ $FORCE -eq 1 ]; then
      print_message "$YELLOW" "Removing existing MCP-ZERO installation..."
      rm -rf "$INSTALL_DIR"
    else
      print_message "$RED" "Error: $INSTALL_DIR already exists and is not empty."
      print_message "$YELLOW" "Use --force to overwrite."
      exit 1
    fi
  fi
  
  # Clone the repository
  print_message "$BLUE" "Cloning MCP-ZERO repository..."
  mkdir -p "$INSTALL_DIR"
  
  # Using current directory as source
  print_message "$YELLOW" "Using local MCP-ZERO as source..."
  SOURCE_DIR="$SCRIPT_DIR"
  
  # Copy files from source to install directory
  rsync -a --exclude '.git' --exclude 'node_modules' "$SOURCE_DIR/" "$INSTALL_DIR/"
  
  # Install Python dependencies
  print_message "$BLUE" "Installing Python dependencies..."
  if [ -f "$INSTALL_DIR/requirements.txt" ]; then
    pip3 install -r "$INSTALL_DIR/requirements.txt"
  else
    print_message "$YELLOW" "No requirements.txt found, creating one with essential packages."
    cat > "$INSTALL_DIR/requirements.txt" << EOF
numpy>=1.20.0
scipy>=1.6.0
torch>=1.9.0
transformers>=4.5.0
flask>=2.0.0
requests>=2.25.0
pytest>=6.0.0
EOF
    pip3 install -r "$INSTALL_DIR/requirements.txt"
  fi
  
  # Set up environment
  print_message "$BLUE" "Setting up environment..."
  export MCPZERO_ROOT="$INSTALL_DIR"
  
  # Add environment variables to user's bashrc
  if ! grep -q "MCPZERO_ROOT" "$HOME/.bashrc"; then
    echo "export MCPZERO_ROOT=\"$INSTALL_DIR\"" >> "$HOME/.bashrc"
    print_message "$GREEN" "Added MCPZERO_ROOT to your .bashrc"
  fi
  
  # Install burn command
  if [ -f "$INSTALL_DIR/install_burn_command.sh" ]; then
    print_message "$BLUE" "Installing burn command..."
    chmod +x "$INSTALL_DIR/install_burn_command.sh"
    (cd "$INSTALL_DIR" && ./install_burn_command.sh)
  else
    print_message "$YELLOW" "No burn command installer found, skipping."
  fi
  
  # Run initialization script if exists
  if [ -f "$INSTALL_DIR/init.sh" ]; then
    print_message "$BLUE" "Running initialization script..."
    chmod +x "$INSTALL_DIR/init.sh"
    (cd "$INSTALL_DIR" && ./init.sh)
  fi
  
  print_message "$GREEN" "MCP-ZERO installation complete!"
}

# Main execution
main() {
  print_message "$BLUE" "MCP-ZERO Installation Script"
  print_message "$BLUE" "==========================="
  
  # Parse command line arguments
  parse_arguments "$@"
  
  # Check dependencies
  check_dependencies
  
  # Install Docker
  install_docker
  
  # Install MCP-ZERO
  install_mcp_zero
  
  print_message "$GREEN" "Installation completed successfully!"
  print_message "$YELLOW" "MCP-ZERO installation: $INSTALL_DIR"
  print_message "$YELLOW" ""
  print_message "$YELLOW" "Next steps:"
  print_message "$YELLOW" "1. You may need to restart your terminal or run 'source ~/.bashrc'"
  print_message "$YELLOW" "2. Run 'burn <folder_name>' to create a burn folder"
  print_message "$YELLOW" "3. Or use the burn_folder.sh script directly with a name parameter"
}

main "$@"
