benchmark.rb

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

require 'fileutils'
require 'optparse'

class BenchmarkRunner
  BASE_APP_PATH = File.expand_path('base_app', __dir__)
  METHODS_PATH = File.expand_path('methods', __dir__)
  attr_reader :verbose

  def initialize(method_name, verbose: false)
    @method_name = method_name
    @method_path = File.join(METHODS_PATH, method_name)
    @backup_gemfile = nil
    @backup_gemfile_lock = nil
    @backup_seeds = nil
    @verbose = verbose

    validate_method!
  end

  def log(msg)
    puts msg if @verbose
  end

  def run
    log '=' * 60
    log "Benchmarking: #{@method_name}"
    log '=' * 60
    setup_method
    install_gems
    reset_and_migrate_database
    load_operations
    seed_data
    run_benchmarks
    restore_project_files

    log "\n✅ Benchmark complete!"
  rescue StandardError => e
    puts "\n❌ Error: #{e.message}"
    puts e.backtrace.first(5)
    restore_project_files
    exit 1
  end

  def self.run_all(verbose: false)
    methods = Dir.glob(File.join(METHODS_PATH, '*'))
                 .select { |d| File.directory?(d) }
                 .map { |d| File.basename(d) }
                 .sort

    puts "Running all #{methods.length} benchmarks...\n\n" if verbose

    script_path = File.expand_path(__FILE__)
    results = {}
    
    methods.each do |method|
      puts "\n" + '=' * 60
      puts "Starting #{methods.index(method) + 1}/#{methods.length}: #{method}"
      puts '=' * 60 + "\n"
      
      # Run each benchmark in a separate Ruby process for isolation
      cmd = verbose ? "ruby #{script_path} -v #{method}" : "ruby #{script_path} #{method}"
      success = system(cmd)
      
      results[method] = success ? '✅' : '❌'
      
      puts "\n" + '=' * 60
      puts "Completed #{methods.index(method) + 1}/#{methods.length}: #{method} #{results[method]}"
      puts '=' * 60 + "\n"
    end

    puts "\n" + '=' * 60
    puts 'ALL BENCHMARKS COMPLETE'
    puts '=' * 60
    methods.each { |m| puts "#{results[m]} #{m}" }
  end

  private

  def validate_method!
    unless File.directory?(@method_path)
      puts "❌ Method '#{@method_name}' not found in methods/"
      puts 'Available methods:'
      Dir.glob(File.join(METHODS_PATH, '*')).each do |dir|
        puts "  - #{File.basename(dir)}" if File.directory?(dir)
      end
      exit 1
    end

    return if File.exist?(File.join(@method_path, 'Gemfile'))

    puts "❌ Gemfile not found in #{@method_path}"
    exit 1
  end

  def setup_method
    log "\n📦 Setting up #{@method_name}..."

    # Backup original Gemfile and Gemfile.lock
    gemfile_path = File.join(BASE_APP_PATH, 'Gemfile')
    gemfile_lock_path = File.join(BASE_APP_PATH, 'Gemfile.lock')
    @backup_gemfile = File.read(gemfile_path)
    @backup_gemfile_lock = File.read(gemfile_lock_path) if File.exist?(gemfile_lock_path)

    # Add method gems to Gemfile
    method_gems = File.read(File.join(@method_path, 'Gemfile'))
    File.open(gemfile_path, 'a') do |f|
      f.puts "\n# #{@method_name} method gems"
      f.puts method_gems
    end
    log '  ✓ Added gems to Gemfile'

    # Copy migrations
    migrations_src = File.join(@method_path, 'migrations')
    migrations_dst = File.join(BASE_APP_PATH, 'db', 'migrate')
    FileUtils.mkdir_p(migrations_dst)

    Dir.glob(File.join(migrations_src, '*.rb')).each do |migration|
      dst_file = File.join(migrations_dst, File.basename(migration))
      FileUtils.cp(migration, dst_file)
      log "  ✓ Copied migration: #{File.basename(migration)}"
    end

    # Copy models
    models_src = File.join(@method_path, 'models')
    models_dst = File.join(BASE_APP_PATH, 'app', 'models')
    FileUtils.mkdir_p(models_dst)

    Dir.glob(File.join(models_src, '*.rb')).each do |model|
      dst_file = File.join(models_dst, File.basename(model))
      FileUtils.cp(model, dst_file)
      log "  ✓ Copied model: #{File.basename(model)}"
    end

    # Copy operations
    copy_operations

    # Copy seeds.rb if method defines custom seeds
    copy_seeds
  end

  def copy_operations
    operations_file = File.join(@method_path, 'operations.rb')
    return unless File.exist?(operations_file)

    operations_dst = File.join(BASE_APP_PATH, 'lib', 'operations')
    FileUtils.mkdir_p(operations_dst)

    dst_file = File.join(operations_dst, "#{@method_name}.rb")
    FileUtils.cp(operations_file, dst_file)
    log "  ✓ Copied operations: #{@method_name}.rb"
  end

  def copy_seeds
    seeds_file = File.join(@method_path, 'seeds.rb')
    return unless File.exist?(seeds_file)

    seeds_dst = File.join(BASE_APP_PATH, 'db', 'seeds.rb')
    @backup_seeds = File.exist?(seeds_dst) ? File.read(seeds_dst) : nil
    FileUtils.cp(seeds_file, seeds_dst)
    log '  ✓ Copied custom seeds.rb'
  end

  def install_gems
    log "\n💎 Installing gems..."
    run_command("cd #{BASE_APP_PATH} && bundle install --quiet", silent: true)
    log '  ✓ Gems installed'
  end

  def reset_and_migrate_database
    log "\n🗄️  Resetting database..."
    db_name = 'array_enum_development'
    
    # Disconnect any active ActiveRecord connections before dropping
    if defined?(ActiveRecord::Base)
      ActiveRecord::Base.connection_pool.disconnect!
    end
    
    run_command("cd #{BASE_APP_PATH} && bundle exec rake db:drop db:create", silent: true)
    log '  ✓ Database reset'

    log "\n📋 Running migrations..."
    run_command("cd #{BASE_APP_PATH} && bundle exec rake db:migrate", silent: true)
    log '  ✓ Migrations complete'
  end

  def load_operations
    operations_file = File.join(BASE_APP_PATH, 'lib', 'operations', "#{@method_name}.rb")
    require operations_file if File.exist?(operations_file)
  end

  def seed_data
    log "\n🌱 Seeding data..."
    require_relative 'base_app/config/environment'

    # Load custom seeds if method defines them
    custom_seeds_path = File.join(@method_path, 'seeds.rb')
    if File.exist?(custom_seeds_path)
      log '  Loading custom seeds...'
      load custom_seeds_path
    end

    # Clear existing users (but preserve colors for junction table)
    User.delete_all if defined?(User)

    # Create sample users with favorite colors using Operations module
    # Use subsets of COLORS constant
    users_data = [
      COLORS[0..2],      # red, green, blue
      COLORS[3..6],      # yellow, purple, orange, pink
      COLORS[7..11],     # cyan, magenta, lime, teal, indigo
      [COLORS[0], COLORS[3]], # red, yellow
      COLORS[2..4] + [COLORS[7]] # blue, purple, pink, cyan
    ]

    users_data.each do |colors|
      Operations.create_user(colors)
    end

    log "  ✓ Created #{User.count} users"
  end

  def run_benchmarks
    puts "\n⚡ Running benchmarks..."
    require 'benchmark'

    n = 1000
    test_colors = [COLORS[0], COLORS[2]] # red, blue
    test_color = COLORS[0] # red
    update_colors = [COLORS[1], COLORS[3]] # green, yellow

    Benchmark.bm(20) do |x|
      x.report('Create user:') do
        n.times do
          Operations.create_user(test_colors)
        end
      end

      x.report('Find by color:') do
        n.times do
          users = Operations.find_by_color(test_color)
          raise "incorrect result #{users.length}" if users.length != n + 2
        end
      end

      x.report('Update colors:') do
        User.limit(100).each do |user|
          Operations.update_user_colors(user, update_colors)
        end
      end
    end

    raise "Unexpected number of end user count: #{User.count}" if User.count != n + 5
    if Operations.count_by_color(test_color) != (n - 100) + 5
      raise "Unexpected number of users liking '#{test_color}': #{Operations.count_by_color(test_color)}"
    end

    puts "\n📊 Database stats:"
    puts "  Total users: #{User.count}"
    puts "  Users liking '#{test_color}': #{Operations.count_by_color(test_color)}"
  end

  def restore_project_files
    log "\n🧹 Restoring project files..."

    # Restore original Gemfile
    if @backup_gemfile
      gemfile_path = File.join(BASE_APP_PATH, 'Gemfile')
      File.write(gemfile_path, @backup_gemfile)
      log '  ✓ Restored Gemfile'
    end

    # Restore original Gemfile.lock
    if @backup_gemfile_lock
      gemfile_lock_path = File.join(BASE_APP_PATH, 'Gemfile.lock')
      File.write(gemfile_lock_path, @backup_gemfile_lock)
      log '  ✓ Restored Gemfile.lock'
    end

    # Remove schema.rb (generated during migrations)
    schema_path = File.join(BASE_APP_PATH, 'db', 'schema.rb')
    if File.exist?(schema_path)
      FileUtils.rm_f(schema_path)
      log '  ✓ Removed schema.rb'
    end

    # Remove copied migrations
    migrations_src = File.join(@method_path, 'migrations')
    migrations_dst = File.join(BASE_APP_PATH, 'db', 'migrate')

    Dir.glob(File.join(migrations_src, '*.rb')).each do |migration|
      dst_file = File.join(migrations_dst, File.basename(migration))
      FileUtils.rm_f(dst_file)
      log "  ✓ Removed migration: #{File.basename(migration)}"
    end

    # Remove copied models
    models_src = File.join(@method_path, 'models')
    models_dst = File.join(BASE_APP_PATH, 'app', 'models')

    Dir.glob(File.join(models_src, '*.rb')).each do |model|
      dst_file = File.join(models_dst, File.basename(model))
      FileUtils.rm_f(dst_file)
      log "  ✓ Removed model: #{File.basename(model)}"
    end

    # Remove copied operations
    remove_operations

    # Restore original seeds.rb
    restore_seeds

    # NOTE: Database cleanup happens at the start of the next run
  end

  def remove_operations
    operations_dst = File.join(BASE_APP_PATH, 'lib', 'operations', "#{@method_name}.rb")
    return unless File.exist?(operations_dst)

    FileUtils.rm_f(operations_dst)
    log "  ✓ Removed operations: #{@method_name}.rb"
  end

  def restore_seeds
    seeds_dst = File.join(BASE_APP_PATH, 'db', 'seeds.rb')
    if @backup_seeds
      File.write(seeds_dst, @backup_seeds)
      log '  ✓ Restored original seeds.rb'
    elsif File.exist?(seeds_dst)
      FileUtils.rm_f(seeds_dst)
      log '  ✓ Removed custom seeds.rb'
    end
  end

  def run_command(cmd, silent: false)
    output = silent ? ' > /dev/null 2>&1' : ''
    system("#{cmd}#{output}") || raise("Command failed: #{cmd}")
  end
end

# Parse command line arguments
options = {}
OptionParser.new do |opts|
  opts.banner = 'Usage: ruby benchmark.rb [options] METHOD'
  opts.on('-a', '--all', 'Run all benchmarks') do
    options[:all] = true
  end
  opts.on('-v', '--verbose', 'Show detailed logs') do
    options[:verbose] = true
  end
  opts.on('-h', '--help', 'Show this help message') do
    puts opts
    exit
  end
end.parse!

verbose = options[:verbose] || false

if options[:all]
  BenchmarkRunner.run_all(verbose: verbose)
else
  method_name = ARGV[0]

  if method_name.nil?
    puts '❌ Please provide a method name'
    puts 'Usage: ruby benchmark.rb METHOD'
    puts "\nExample:"
    puts '  ruby benchmark.rb array_enum'
    puts "\nOr run all benchmarks:"
    puts '  ruby benchmark.rb --all'
    exit 1
  end

  runner = BenchmarkRunner.new(method_name, verbose: verbose)
  runner.run
end