class ElasticAPM::Metrics::Registry

def collect
  sets.each_value.each_with_object([]) do |set, arr|
    samples = set.collect
    next unless samples
    arr.concat(samples)
  end
end

def collect_and_send
  return unless @config.recording?
  metricsets = collect
  metricsets.compact!
  metricsets.each do |m|
    callback.call(m)
  end
end

def define_sets
  # Only set the @sets once, in case we stop
  # and start again.
  return unless @sets.nil?
  sets = {
    system: CpuMemSet,
    vm: VMSet,
    breakdown: BreakdownSet,
    transaction: TransactionSet
  }
  if defined?(JVMSet)
    debug "Enabling JVM metrics collection"
    sets[:jvm] = JVMSet
  end
  @sets = sets.each_with_object({}) do |(key, kls), _sets_|
    debug "Adding metrics collector '#{kls}'"
    _sets_[key] = kls.new(config)
  end
end

def get(key)
  sets.fetch(key)
end

def handle_forking!
  # Note that ideally we would be able to check if the @timer_task died
  # and restart it. You can't simply check @timer_task.running? because
  # it will only return the state of the TimerTask, not whether the
  # internal thread used to manage the execution interval has died.
  # This is a limitation of the Concurrent::TimerTask object.
  # Therefore, our only option when forked is to stop and start.
  # ~estolfo
  stop
  start
end

def initialize(config, &block)
  @config = config
  @callback = block
  @sets = nil
end

def running?
  !!@running
end

def start
  unless config.collect_metrics?
    debug 'Skipping metrics'
    return
  end
  define_sets
  debug 'Starting metrics'
  @timer_task = Concurrent::TimerTask.execute(
    run_now: true,
    execution_interval: config.metrics_interval
  ) do
    begin
      debug 'Collecting metrics'
      collect_and_send
      true
    rescue StandardError => e
      error 'Error while collecting metrics: %e', e.inspect
      debug { e.backtrace.join("\n") }
      false
    end
  end
  @running = true
end

def stop
  return unless running?
  debug 'Stopping metrics'
  @timer_task.shutdown
  @running = false
end