#!/usr/bin/env ruby
Process.setproctitle($0)
# frozen_string_literal: true

require 'strscan'

Process.setproctitle($PROGRAM_NAME)

def camelize(subject)
  subject.gsub(%r{(^|[/_])[a-z]}) { |x| x.sub('/', '::').sub('_', '').upcase }
end

def underscore(subject)
  subject.split(/(?=[A-Z])/).map(&:downcase).join('_').gsub('::_', '/')
end

def constantize(subject)
  Object.const_get(subject)
end

# Split on blanks unless quoted or escaped:
# "--foo --bar=bar --baz='baz baz'\ baz" #=> ["--foo", "--bar=bar", "--baz=baz baz baz"]
def split_options(options)
  res = []

  return res unless options

  current = ''
  s = StringScanner.new(options)
  until s.eos?
    if s.scan(/\s+/)
      res << current
      current = ''
    elsif s.scan(/\\./)
      current += s.matched[1]
    elsif s.scan(/['"]/)
      match = s.matched
      loop do
        if s.scan(match)
          break
        elsif s.scan(/\\./)
          current += s.matched[1]
        else
          current += s.getch
        end
      end
    else
      current += s.getch
    end
  end

  res << current unless current.empty?

  res
end

def read_flags(argv)
  res = []

  while (arg = argv.shift)
    break if arg == '--'

    res << arg
  end

  res
end

def usage
  warn <<~USAGE
    usage: riemann-wrapper [common options] -- tool1 [tool1 options] [-- tool2 [tool2 options] ...]
           riemann-wrapper /path/to/configuration/file.yml

    Run multiple Riemann tools in a single process.  A single connection to
    riemann is maintained and shared for all tools, the connection flags should
    only be passed as common options.

    Examples:
      1. Run the fd, health and ntp tools with default options:

         riemann-wrapper -- fd -- health -- ntp

      2. Run the fd, health and ntp tools against a remote riemann server using
         TCP and tagging each event with the name of the tool that produced it:

         riemann-wrapper --host riemann.example.com --tcp -- \\
                         fd     --tag=fd     -- \\
                         health --tag=health -- \\
                         ntp    --tag=ntp

      3. Same as above example, but using a configuration file (more verbose but
         easier to handle when running riemann-wrapper manually of managing it
         with a Configuration Management system):

         cat > config.yml << EOT
         ---
         options: --host riemann.example.com --tcp
         tools:
         - name: fd
           options: --tag=fd
         - name: health
           options: --tag=health
         - name: ntp
           options: --tag=ntp
         EOT
         riemann-wrapper config.yml
  USAGE
  exit 1
end

usage if ARGV.empty?

if ARGV.size == 1
  unless File.readable?(ARGV[0])
    warn "Cannot open file for reading: #{ARGV[0]}"
    usage
  end

  require 'yaml'
  config = YAML.safe_load(File.read(ARGV[0]))

  arguments = split_options(config['options'])
  config['tools'].each do |tool|
    arguments << '--'
    arguments << tool['name']
    arguments += split_options(tool['options'])
  end

  ARGV.replace(arguments)
end

argv = ARGV.dup

common_argv = read_flags(argv)

threads = []

# Terminate the whole process is some thread fail
Thread.abort_on_exception = true

while argv.any?
  tool = argv.shift
  tool_argv = read_flags(argv)

  require "riemann/tools/#{tool}"
  tool_class = constantize(camelize("riemann/tools/#{tool}"))

  ARGV.replace(common_argv + tool_argv)
  instance = tool_class.new
  # Force evaluation of options.  This rely on ARGV and needs to be done before
  # we launch multiple threads which compete to read information from there.
  instance.options
  threads << Thread.new(instance, &:run)
end

threads.each(&:join)
