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

require "bundler/setup"
require "benchmark/ips"
require "benchmark/memory"
require "memory_profiler"
require "json"
require "fileutils"

def run_benchmark(commit_hash, with_yjit)
  results = {}
  results["commit"] = commit_hash
  results["yjit"] = with_yjit
  results["timestamp"] = Time.now.to_i
  results["ruby_version"] = RUBY_VERSION
  results["benchmark_results"] = {}

  test_name = "my first test"
  results["benchmark_results"][test_name] = {}

  default_benchmark = Benchmark.ips do |x|
    # x.report ...

    x.compare!
  end

  # Extract IPS results
  default_benchmark.entries.each do |entry|
    results["benchmark_results"][test_name][entry.label] = {
      "ips" => entry.ips,
      "microseconds_per_iter" => (1_000_000 / entry.ips)
    }
  end

  # Run memory benchmarks for default encoder
  memory_benchmark = Benchmark.memory do |x|
    # x.report ...

    x.compare!
  end

  # Extract memory results
  memory_benchmark.entries.each do |entry|
    results["benchmark_results"][test_name][entry.label]["memory"] = {
      "memsize" => entry.measurement.memory.allocated,
      "objects" => entry.measurement.objects.allocated,
      "strings" => entry.measurement.strings.allocated
    }
  end

  # Run more benchmarks

  # Memory profiler for more detailed memory analysis
  report = MemoryProfiler.report do
    # do stuff...
  end

  results["memory_profiler"] = {
    "total_allocated" => report.total_allocated_memsize,
    "total_allocated_objects" => report.total_allocated,
    "total_retained" => report.total_retained_memsize,
    "total_retained_objects" => report.total_retained
  }

  # Add top allocation locations
  results["memory_profiler"]["top_allocations"] = report.allocated_memory_by_location.sort_by { |l| -l[:count] }.first(5).map do |loc|
    file, line = loc[:data].split(":")
    {
      "file" => file,
      "line" => line,
      "bytes" => loc[:count]
    }
  end

  results
end

def save_results(results, commit_hash, with_yjit)
  yjit_suffix = with_yjit ? "yjit" : "no_yjit"
  filename = File.join(RESULTS_DIR, "#{commit_hash}_#{yjit_suffix}.json")
  File.write(filename, JSON.pretty_generate(results))
  puts "Saved results to #{filename}"
end

def summarize_results
  puts "\n\n==== BENCHMARK SUMMARY ====\n\n"

  # Load all result files
  result_files = Dir.glob(File.join(RESULTS_DIR, "*.json"))
  all_results = result_files.map { |file| JSON.parse(File.read(file)) }

  # Group by YJIT status
  yjit_results = all_results.select { |r| r["yjit"] }
  no_yjit_results = all_results.select { |r| !r["yjit"] }

  # Sort results by commit timestamp
  yjit_results.sort_by! { |r| r["timestamp"] }
  no_yjit_results.sort_by! { |r| r["timestamp"] }

  # Print summary tables
  puts "=== Performance without YJIT ==="
  
  # Each test...
  test_name = "..."
  entry_label = "..."
  print_performance_table(no_yjit_results, test_name, entry_label)
  # ...
  
  puts "\n=== Performance with YJIT ==="
  # Each test...
  test_name = "..."
  entry_label = "..."
  print_performance_table(yjit_results, test_name, entry_label)

  puts "\n=== Memory Usage ==="

  test_name = "..."
  entry_label = "..."
  print_memory_table(no_yjit_results, test_name, entry_label)
  # ...

  puts "\n=== Top Allocation Sites ==="
  print_allocation_sites(no_yjit_results)
  
  puts "\n\n==== HIGHLIGHTS ====\n"
  print_highlights_table(no_yjit_results, yjit_results)
end

def print_performance_table(results, benchmark_group, test_case)
  puts "\n#{benchmark_group} - #{test_case} (iterations per second):"
  puts "-" * 80
  puts "%-10s %-20s %-15s %-15s" % ["Commit", "Description", "IPS", "μs/iteration"]
  puts "-" * 80

  baseline_ips = nil

  results.each do |result|
    commit = result["commit"][0..7]
    commit_msg = `git log --format=%s -n 1 #{result["commit"]}`.strip[0..17] + "..."

    if result["benchmark_results"][benchmark_group] && result["benchmark_results"][benchmark_group][test_case]
      ips = result["benchmark_results"][benchmark_group][test_case]["ips"]
      us_per_iter = result["benchmark_results"][benchmark_group][test_case]["microseconds_per_iter"]

      baseline_ips ||= ips
      comparison = baseline_ips ? (ips / baseline_ips).round(2) : 1.0
      comparison_str = if comparison == 1.0
        "baseline"
      else
        ((comparison > 1.0) ? "#{comparison}x faster" : "#{(1.0 / comparison).round(2)}x slower")
      end

      puts "%-10s %-20s %-15.2f %-15.2f %s" % [commit, commit_msg, ips, us_per_iter, comparison_str]
    else
      puts "%-10s %-20s %-15s %-15s" % [commit, commit_msg, "N/A", "N/A"]
    end
  end
end

def print_memory_table(results, benchmark_group, test_case)
  puts "\n#{benchmark_group} - #{test_case} (memory usage):"
  puts "-" * 80
  puts "%-10s %-20s %-15s %-15s %-15s" % ["Commit", "Description", "Bytes", "Objects", "Strings"]
  puts "-" * 80

  baseline_bytes = nil

  results.each do |result|
    commit = result["commit"][0..7]
    commit_msg = `git log --format=%s -n 1 #{result["commit"]}`.strip[0..17] + "..."

    if result["benchmark_results"][benchmark_group] &&
        result["benchmark_results"][benchmark_group][test_case] &&
        result["benchmark_results"][benchmark_group][test_case]["memory"]

      bytes = result["benchmark_results"][benchmark_group][test_case]["memory"]["memsize"]
      objects = result["benchmark_results"][benchmark_group][test_case]["memory"]["objects"]
      strings = result["benchmark_results"][benchmark_group][test_case]["memory"]["strings"]

      baseline_bytes ||= bytes
      comparison = baseline_bytes ? (bytes.to_f / baseline_bytes).round(2) : 1.0
      comparison_str = if comparison == 1.0
        "baseline"
      else
        ((comparison < 1.0) ? "#{(1.0 / comparison).round(2)}x better" : "#{comparison}x worse")
      end

      puts "%-10s %-20s %-15d %-15d %-15d %s" % [commit, commit_msg, bytes, objects, strings, comparison_str]
    else
      puts "%-10s %-20s %-15s %-15s %-15s" % [commit, commit_msg, "N/A", "N/A", "N/A"]
    end
  end
end

def print_allocation_sites(results)
  results.each do |result|
    commit = result["commit"][0..7]
    commit_msg = `git log --format=%s -n 1 #{result["commit"]}`.strip

    puts "\n#{commit} - #{commit_msg}"
    puts "-" * 80

    if result["memory_profiler"] && result["memory_profiler"]["top_allocations"]
      puts "%-40s %-10s %-15s" % ["File:Line", "Bytes", "Class"]
      puts "-" * 80

      result["memory_profiler"]["top_allocations"].each do |alloc|
        file_line = "#{File.basename(alloc["file"])}:#{alloc["line"]}"
        puts "%-40s %-10d %-15s" % [file_line, alloc["bytes"], alloc["class"]]
      end
    else
      puts "No memory profiler data available"
    end
  end
end

def print_highlights_table(no_yjit_results, yjit_results = [])
  puts "Performance and Memory Highlights"
  puts "-" * 145
  header = "%-10s %-25s %-15s %-15s %-15s %-15s %-15s"
  puts header % [
    "Commit", 
    "Description", 
    "IPS Change", 
    "YJIT IPS Change", 
    "YJIT vs No YJIT", 
    "Memory Change", 
    "Objects Change"
  ]
  puts "-" * 145
  
  # Sort results by timestamp
  no_yjit_sorted = no_yjit_results.sort_by { |r| r["timestamp"] }
  
  # Create a hash of commit to yjit result for easy lookup
  yjit_lookup = {}
  yjit_results.each do |res|
    yjit_lookup[res["commit"]] = res
  end
  
  # Use the first commit as baseline
  baseline = no_yjit_sorted.first
  
  # Reference values for comparison
  test_name = "..." 
  entry_label = "..."
  baseline_ips = baseline.dig("benchmark_results", test_name, entry_label, "ips")
  baseline_memory = baseline.dig("benchmark_results", test_name, entry_label, "memory", "memsize")
  baseline_objects = baseline.dig("benchmark_results", test_name, entry_label, "memory", "objects")
  
  # Get YJIT baseline if available
  yjit_baseline = yjit_lookup[baseline["commit"]]
  baseline_yjit_ips = yjit_baseline&.dig("benchmark_results", test_name, entry_label, "ips")
  
  # Show baseline
  commit = baseline["commit"][0..7]
  commit_msg = `git log --format=%s -n 1 #{baseline["commit"]}`.strip[0..22] + "..."
  puts "%-10s %-25s %-15s %-15s %-15s %-15s %-15s" % [
    commit, 
    commit_msg, 
    "baseline", 
    "baseline", 
    "#{(baseline_yjit_ips && baseline_ips) ? (baseline_yjit_ips / baseline_ips).round(2) : 'N/A'}x",
    "baseline", 
    "baseline"
  ]
  
  # Skip the first one (baseline)
  no_yjit_sorted[1..].each do |result|
    commit = result["commit"][0..7]
    commit_msg = `git log --format=%s -n 1 #{result["commit"]}`.strip[0..22] + "..."
    
    # Get performance metrics for this commit
    current_ips = result.dig("benchmark_results", test_name, entry_label, "ips")
    current_memory = result.dig("benchmark_results", test_name, entry_label, "memory", "memsize") 
    current_objects = result.dig("benchmark_results", test_name, entry_label, "memory", "objects")
    
    # Get corresponding YJIT results if available
    yjit_result = yjit_lookup[result["commit"]]
    current_yjit_ips = yjit_result&.dig("benchmark_results", test_name, entry_label, "ips")
    
    # Calculate changes
    if current_ips && baseline_ips
      ips_comparison = (current_ips / baseline_ips).round(2)
      ips_change = if ips_comparison > 1.0
        "+#{((ips_comparison - 1) * 100).round(1)}% (#{ips_comparison}x)"
      elsif ips_comparison < 1.0
        "-#{((1 - ips_comparison) * 100).round(1)}% (#{ips_comparison}x)"
      else
        "No change"
      end
    else
      ips_change = "N/A"
    end
    
    # Calculate YJIT changes
    if current_yjit_ips && baseline_yjit_ips
      yjit_ips_comparison = (current_yjit_ips / baseline_yjit_ips).round(2)
      yjit_ips_change = if yjit_ips_comparison > 1.0
        "+#{((yjit_ips_comparison - 1) * 100).round(1)}% (#{yjit_ips_comparison}x)"
      elsif yjit_ips_comparison < 1.0
        "-#{((1 - yjit_ips_comparison) * 100).round(1)}% (#{yjit_ips_comparison}x)"
      else
        "No change"
      end
    else
      yjit_ips_change = "N/A"
    end
    
    # Compare YJIT vs no YJIT for this commit
    yjit_vs_no_yjit = if current_yjit_ips && current_ips
      "#{(current_yjit_ips / current_ips).round(2)}x"
    else
      "N/A"
    end
    
    if current_memory && baseline_memory
      memory_comparison = (current_memory.to_f / baseline_memory).round(2)
      memory_change = if memory_comparison < 1.0
        "-#{((1 - memory_comparison) * 100).round(1)}% (#{memory_comparison}x)"
      elsif memory_comparison > 1.0
        "+#{((memory_comparison - 1) * 100).round(1)}% (#{memory_comparison}x)"
      else
        "No change"
      end
    else
      memory_change = "N/A"
    end
    
    if current_objects && baseline_objects
      objects_comparison = (current_objects.to_f / baseline_objects).round(2)
      objects_change = if objects_comparison < 1.0
        "-#{((1 - objects_comparison) * 100).round(1)}% (#{objects_comparison}x)"
      elsif objects_comparison > 1.0
        "+#{((objects_comparison - 1) * 100).round(1)}% (#{objects_comparison}x)"
      else
        "No change"
      end
    else
      objects_change = "N/A"
    end
    
    puts "%-10s %-25s %-15s %-15s %-15s %-15s %-15s" % [
      commit, 
      commit_msg, 
      ips_change, 
      yjit_ips_change, 
      yjit_vs_no_yjit,
      memory_change, 
      objects_change
    ]
  end
end

RESULTS_DIR = File.join(Dir.pwd, "benchmark_results")
FileUtils.mkdir_p(RESULTS_DIR)

# Check if we're being called to run benchmarks for a specific commit
if ARGV.include?("--bench")
  # Get the commit hash
  commit_index = ARGV.index("--bench") + 1
  if commit_index >= ARGV.length
    puts "Error: No commit hash provided after --bench"
    exit 1
  end
  commit = ARGV[commit_index]
  with_yjit = ARGV.include?("--yjit")

  # Any setup work...

  # Run benchmark and save results
  results = run_benchmark(commit, with_yjit)
  save_results(results, commit, with_yjit)
  exit 0
elsif ARGV.include?("--results-only")
  # Just display results without running benchmarks
  if Dir.exist?(RESULTS_DIR) && !Dir.empty?(RESULTS_DIR)
    summarize_results
    exit 0
  else
    puts "Error: No benchmark results found in #{RESULTS_DIR}"
    puts "Run benchmarks first or specify commits to benchmark."
    exit 1
  end
elsif ARGV.length != 2
  puts "Usage: #{$PROGRAM_NAME} <start_commit> <end_commit>"
  puts "       #{$PROGRAM_NAME} --results-only (to view results without running benchmarks)"
  exit 1
end

START_COMMIT = ARGV[0]
END_COMMIT = ARGV[1]

# Get list of commits between start and end (inclusive)
commits = `git log --format=%H #{START_COMMIT}..#{END_COMMIT}`.split("\n")
# Add the start commit to the beginning
commits.unshift(`git rev-parse #{START_COMMIT}`.strip)
# Add the end commit if it's not already in the list
end_commit_hash = `git rev-parse #{END_COMMIT}`.strip
commits << end_commit_hash unless commits.include?(end_commit_hash)

puts "Found #{commits.length} commits to benchmark"

# Save current branch
current_branch = `git branch --show-current`.strip
puts "Current branch: #{current_branch}"

# For each commit, checkout and run benchmarks in separate processes
commits.each_with_index do |commit, index|
  puts "\n\n==== Benchmarking commit #{index + 1}/#{commits.length}: #{commit} ===="
  commit_msg = `git log --format=%s -n 1 #{commit}`.strip
  puts "#{commit}: #{commit_msg}"

  # Checkout the commit
  `git checkout #{commit}`

  # Wait for filesystem to settle
  sleep 1

  # Run benchmarks without YJIT in a separate process
  puts "\nRunning benchmarks without YJIT"
  system("ruby #{$PROGRAM_NAME} --bench #{commit}")

  # Run benchmarks with YJIT in a separate process
  puts "\nRunning benchmarks with YJIT"
  system("ruby --yjit #{$PROGRAM_NAME} --bench #{commit} --yjit")
end

# Return to original branch
`git checkout #{current_branch}`

# Summarize results
summarize_results