#!/opt/puppetlabs/puppet/bin/ruby

require_relative 'pe_metrics'

include PuppetX::Puppetlabs::PuppetMetricsCollector

# Wrap code in a main method so we can rescue errors
def main
  Puppet.initialize_settings
  metrics_collector_setup

  output_file = Time.now.utc.strftime('%Y%m%dT%H%M%SZ') + '.json'

  @hosts.each do |host|
    timestamp = Time.now
    dataset = { 'timestamp' => timestamp.utc.iso8601, 'servers' => {} }
    hostkey = host.tr('.', '-')
    dataset['servers'][hostkey] = { @metrics_type => {} }
    url = "https://#{host}:#{@port}"

    # Puma based services use a different status endpoint and have an additional GC endpoint
    if ['ace', 'bolt'].include?(@metrics_type)
      status_url = URI("#{url}/admin/status")

      gc_url = URI("#{url}/admin/gc_stat")
      gc_status_output = get_endpoint(gc_url)
    else
      status_url = URI("#{url}/status/v1/services?level=debug")
    end

    status_output = get_endpoint(status_url)
    dataset['servers'][hostkey][@metrics_type] = status_output
    dataset['servers'][hostkey][@metrics_type]['gc_stats'] = gc_status_output if gc_status_output

    unless @additional_metrics.empty?
      metrics_array = retrieve_additional_metrics(URI("#{url}/metrics/v2/read"), @metrics_type, @additional_metrics)

      metrics_array.each do |metric_hash|
        metric_name = metric_hash['name']
        metric_data = metric_hash['data']
        dataset['servers'][hostkey][@metrics_type][metric_name] = metric_data
      end
    end

    dataset['servers'][hostkey][@metrics_type]['error'] = @errors
    dataset['servers'][hostkey][@metrics_type]['error_count'] = @errors.count
    dataset['servers'][hostkey][@metrics_type]['api-query-start'] = timestamp.utc.iso8601
    dataset['servers'][hostkey][@metrics_type]['api-query-duration'] = Time.now - timestamp

    filtered_dataset = filter_metrics(dataset, @excludes)
    json_dataset = JSON.generate(filtered_dataset)

    unless @output_dir.nil?
      Dir.chdir(@output_dir) do
        Dir.mkdir(host) unless File.exist?(host)
        File.open(File.join(host, output_file), 'w') do |file|
          file.write(json_dataset)
        end
      end
    end

    next unless @print
    # Use puts so that data gatherd from each host is written to stdout
    # with a newline separator. This enables parsing of multiple datasets
    # as the output stream is in line-delimited JSON form.
    STDOUT.puts(json_dataset)
  end
rescue StandardError => e
  STDERR.puts "Failed to collect #{@metrics_type} metrics: #{e.message}"
  STDERR.puts e.backtrace
  nil
end

if __FILE__ == $PROGRAM_NAME
  main
end
