#!/usr/bin/env ruby

require 'optparse'
require 'regexp-examples'
require 'rcmd'

# This program is for executing a command on multiple nodes at the same time and accepts the following arguments.
#
# * -u, --username    | Username for the connections
# * -n, --nodes       | comma seperated list of nodes. '-' to use piped list
# * -t, --threads     | Number of threads to run
# * -c, --command     | Command to be executed
# * -q, --quiet       | Boolean flag for suppressing stdout of executed commands
# * -v, --version     | Print version string and exit


# Set default options
options = { :threads => 4, :nodes => nil, :environment => nil, :host_file => nil, :user => 'root', :password => false, :command => nil, :quiet => false, :version => false, :debug => false, :expression => nil, :database => false }

host_list = []

# Setup options
opts = OptionParser.new #:nodoc:

opts.on('-u username', '--username username', String, "Username for SSH connections") { |v| options[:user] = v }
opts.on('-n nodes', '--nodes x,y,z', Array, "Comma seperated list of nodes. use '-' for a space seperated list piped from another command") { |v| options[:nodes] = v }
opts.on('-r regex', '--regexp regex', String, "Use Regex to build host list (ruby regexp)") { |v| options[:expression] = v}
# opts.on('-d database', '--database', "Use hosts from database (~/.dbconf.yaml)") { |v| options[:database] = v}
opts.on('-t threads', '--threads threads', Integer, "Number of threads to run") { |v| options[:threads] = v }
opts.on('-c <command>', '--command <command>', String, "Quoted string containing the command to be run") { |v| options[:command] = v }
opts.on('-q', '--quiet', "Suppress stdout of commands.  stderr will still be displayed") { |v| options[:quiet] = v }
opts.on('-v', '--version', "Print what version of the command is in use") { |v| options[:version] = v}
opts.on('-D', '--debug', "Print debug information") { |v| options[:debug] = v }

# Process options
begin
  opts.parse!(ARGV)
  unless options[:version] or options[:expression]
      mandatory = [:command, :nodes]
      missing = mandatory.select{ |param| options[param].nil? }
      raise OptionParser::MissingArgument, missing.join(', ') unless missing.empty?
  end
  rescue OptionParser::ParseError => e
    $stderr.puts e
    $stderr.puts opts
    exit
end

# Process Regex for node list generation
if options[:expression]
  unless options[:database]
    options[:nodes] = Regexp.new(options[:expression]).examples(max_group_results: 500)
  end
end

# Print version and exit if version option chosen
if options[:version]
  $stdout.puts "Version: #{Rcmd::VERSION}"
  exit
end

# Populate the host list
if options[:nodes][0] == '-'
  ARGF.read.split().each do |h|
    host_list << h
  end
else
  host_list = options[:nodes]
end

# abort if host_list is empty
abort("No hosts to run command on") if host_list.nil?

# Set needed variables
Rcmd.host_list= host_list
Rcmd.nthreads= options[:threads]
Rcmd.user= options[:user]
Rcmd.quiet= options[:quiet]
Rcmd.command= options[:command]
Rcmd.debug= options[:debug]

# Execute command on provided hosts
Rcmd.run_command