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

require 'erb'
require 'optparse'
require 'pathname'

require_relative 'lib/tracepoint_defs'

EXTRA_GROUPS = RubyTimelineTool::USDT_DEFS.keys - RubyTimelineTool::DEFAULT_GROUPS

options = {}

OptionParser.new do |parser|
  parser.accept(Pathname) do |value|
    Pathname.new(value)
  end
  parser.on('-r', '--ruby-bin RUBY', Pathname, 'Path to the ruby executable')
  parser.on('-m', '--gc-module-so GC_MODULE', Pathname, 'Path to the GC module .so')
  parser.on('-e', '--erb-debug', TrueClass, 'Print the Ruby script generated by ERB.  For debugging.')
  parser.on('-p', '--print-script', TrueClass, 'Print the content of the bpftrace script')
  parser.on('-d', '--dry-run', TrueClass, 'Dry run.  Print the bpftrace command to be executed, but do not actually execute it.')
  parser.on('-g', '--extra-groups=', "Enable extra USDT groups, separated by comma.  Available groups: #{EXTRA_GROUPS.join(', ')}")
end.parse!(into: options)

[:'ruby-bin'].each do |k|
  if !options.include?(k)
    raise "Option --#{k} is required"
  end
end

ruby_bin = options[:'ruby-bin']
if !ruby_bin.file?
  raise "Ruby executable '#{ruby_bin}' does not exist"
end

gc_module_so = options[:'gc-module-so']

gc_module_kind = case
when gc_module_so.nil?
  nil
when m = /librubygc\.(\w+)/.match(gc_module_so.basename.to_s)
  m[1]
else
  raise "Unexpected gc module so name: #{gc_module_so}"
end

script_path = Pathname(__dir__) / 'capture.bt'
script = File.read(script_path)
template = ERB.new(script, trim_mode: '%<>')
template.location = [script_path.to_s, 1]

if options[:'erb-debug']
  $stderr.puts template.src
end

extra_groups_opt = options[:'extra-groups']

selected_groups = RubyTimelineTool::DEFAULT_GROUPS.dup

if !extra_groups_opt.nil?
  extra_groups_opt.split(',').each do |group|
    if !EXTRA_GROUPS.include?(group)
      raise "Unknown USDT group #{group}"
    end

    selected_groups << group
  end
end

selected_groups.uniq!

$stderr.puts "Selected groups: #{selected_groups.join(',')}"

usdts = RubyTimelineTool::USDT_DEFS.values_at(*selected_groups).flatten

ScriptUsdtEntry = Struct.new('ScriptUsdtEntry', :file, :probe_name, :nargs)

usdt_entries = []

usdts.each do |t|
  file = case [gc_module_kind, t.where]
    in nil,   'ruby'                  then ruby_bin
    in nil,   'default'               then ruby_bin
    in x,     'ruby'                  then ruby_bin
    in x,     y           if x == y   then gc_module_so
    else                                   nil
  end

  if file.nil?
    @stderr.puts "Skipped probe '#{probe_name}' which is for '#{location}'.  Current GC module: '#{gc_module_kind}'."
    next
  end

  nargs = t.args.length

  usdt_entries << ScriptUsdtEntry.new(file, t.probe_name, nargs)
end

content = template.result_with_hash({
  usdt_entries:,
})

if options[:'print-script']
  $stderr.puts content
end

IO.write('capture.out.bt', content)

command_line = ['sudo', 'bpftrace', '-v', '--unsafe', 'capture.out.bt']
$stderr.puts 'Command to execute:'
$stderr.puts command_line.map { |s| "'#{s}'" }.join(' ')

if options[:'dry-run']
  $stderr.puts 'Dry run.  Exit...'
  exit 0
end

exec(*command_line)