#!/usr/bin/env ruby
# frozen_string_literal: true

require "bundler/setup"
require "benchmark/ips"
require "tmpdir"
require "socket"
require "statsd-instrument"

def send_metrics(client)
  client.increment("StatsD.increment", 10)
  client.measure("StatsD.measure") { 1 + 1 }
  client.gauge("StatsD.gauge", 12.0, tags: ["foo:bar", "quc"])
  client.set("StatsD.set", "value", tags: { foo: "bar", baz: "quc" })
  if client.datagram_builder_class == StatsD::Instrument::DogStatsDDatagramBuilder
    client.event("StasD.event", "12345")
    client.service_check("StatsD.service_check", "ok")
  end
end

def benchmark_implementation(name, env = {})
  revision = %x(git rev-parse HEAD).rstrip
  base_revision = %x(git rev-parse origin/main).rstrip
  branch = if revision == base_revision
    "main"
  else
    %x(git rev-parse --abbrev-ref HEAD).rstrip
  end

  log_dir = "#{Dir.tmpdir}/statsd-instrument-benchmarks"
  intermediate_results_filename = File.join(log_dir, "#{File.basename($PROGRAM_NAME)}-#{name}")
  log_filename = File.join(log_dir, "#{File.basename($PROGRAM_NAME)}-#{name}.log")
  FileUtils.mkdir_p(log_dir)

  # Set up an UDP listener to which we can send StatsD packets
  receiver = UDPSocket.new
  receiver.bind("localhost", 0)

  log_file = File.open(log_filename, "w+", level: Logger::WARN)
  StatsD.logger = Logger.new(log_file)

  udp_client = StatsD::Instrument::Environment.new(ENV.to_h.merge(
    "STATSD_ADDR" => "#{receiver.addr[2]}:#{receiver.addr[1]}",
    "STATSD_IMPLEMENTATION" => "dogstatsd",
    "STATSD_ENV" => "production",
  ).merge(env)).client

  puts "===== #{name} ====="
  report = Benchmark.ips do |bench|
    bench.report("#{name} (branch: #{branch}, sha: #{revision[0, 7]})") do
      send_metrics(udp_client)
    end

    # Store the results in between runs
    bench.save!(intermediate_results_filename)
    bench.compare!
  end

  receiver.close
  udp_client.shutdown if udp_client.respond_to?(:shutdown)

  if report.entries.length == 1
    puts
    puts "To compare the performance of this revision against another revision (e.g. main),"
    puts "check out a different branch and run this benchmark script again."
  elsif ENV["KEEP_RESULTS"]
    puts
    puts "The intermediate results have been stored in #{intermediate_results_filename}"
  else
    File.unlink(intermediate_results_filename)
  end

  # log_file.close
  logs = File.read(log_filename)
  unless logs.empty?
    puts
    puts "==== logs ===="
    puts logs
  end
  puts "================"
end

benchmark_implementation("UDP sync", "STATSD_BUFFER_CAPACITY" => "0")
benchmark_implementation("UDP batched")
benchmark_implementation("UDP batched with aggregation", "STATSD_ENABLE_AGGREGATION" => "true")