#! /usr/bin/env ruby

# == Synopsis
#
# Profiles a Ruby program.
#
# == Usage
#
# ruby_prof [options] <script.rb> [options-for-script]
#
# Options:
#     -p, --printer=printer            Select a printer:
#                                        flat - Prints a flat profile as text (default).
#                                        graph - Prints a graph profile as text.
#                                        graph_html - Prints a graph profile as html.
#                                        call_tree - format for KCacheGrind
#     -f, --file=path                  Output results to a file instead of standard out.
#     -m, --measure-mode=measure_mode  Select a measurement mode:
#                                        process - Use process time (default).
#                                        wall - Use wall time.
#                                        cpu - Use the CPU clock counter
#                                              (only supported on Pentium and PowerPCs).
#                                        allocations - Tracks object allocations
#                                              (requires a patched Ruby interpreter).
#         --replace-progname           Replace $0 when loading the .rb files.
#     -h, --help                       Show help message
#         --version                    Show version
# 
#
# See also: {flat profiles}[link:files/examples/flat_txt.html], {graph profiles}[link:files/examples/graph_txt.html], {html graph profiles}[link:files/examples/graph_html.html]
#

require 'ostruct'
require 'optparse'
require 'ruby-prof'

options = OpenStruct.new
options.measure_mode = RubyProf::PROCESS_TIME
options.printer = RubyProf::FlatPrinter
options.min_percent = 0
options.file = nil
options.replace_prog_name = false

opts = OptionParser.new do |opts|
  opts.banner = "ruby_prof #{RubyProf::VERSION}\n" +
                "Usage: ruby_prof [options] <script.rb> [--extra-options-for-script]"
 
  opts.separator ""
  opts.separator "Options:"

    
  opts.on('-p printer', '--printer=printer', [:flat, :graph, :graph_html, :call_tree],
          'Select a printer:',
          '  flat - Prints a flat profile as text (default).',
          '  graph - Prints a graph profile as text.',
          '  graph_html - Prints a graph profile as html.',
          '  call_tree - format for KCacheGrind' ) do |printer|

          
    case printer
      when :flat
        options.printer = RubyProf::FlatPrinter
      when :graph
        options.printer = RubyProf::GraphPrinter
      when :graph_html
        options.printer = RubyProf::GraphHtmlPrinter
    end
  end
    
  opts.on('-m min_percent', '--min_percent=min_percent', Float,
          'The minimum percent a method must take before ',
          '  being included in output reports.',
                                        '  this option is not supported for call tree.') do |min_percent|
    options.min_percent = min_percent
  end

  opts.on('-f path', '--file=path',
        'Output results to a file instead of standard out.') do |file|
    options.file = file
  end
    
  opts.on('-m measure_mode', '--measure-mode=measure_mode',
      [:process, :wall, :cpu, :allocations],
      'Select what ruby-prof should measure:',
      '  process - Process time (default).',
      '  wall - Wall time.',
      '  cpu - CPU time',
      '        (only supported on Pentium and PowerPCs).',
      '  allocations - O3bject allocations (required patched Ruby interpreter).') do |measure_mode|
      
      case measure_mode
      when :process
        options.measure_mode = RubyProf::PROCESS_TIME     
      when :wall
        options.measure_mode = RubyProf::WALL_TIME      
      when :cpu
        options.measure_mode = RubyProf::CPU_TIME
      when :allocations
        options.measure_mode = RubyProf::ALLOCATIONS
      end
  end
        
        opts.on("--replace-progname", "Replace $0 when loading the .rb files.") do
                options.replace_prog_name = true
        end
    
  opts.on_tail("-h", "--help", "Show help message") do
      puts opts
      exit
  end
  opts.on_tail("-v", "--version", "Show version") do
      puts "ruby_prof " + RubyProf::VERSION
      exit
  end
end

begin
  opts.parse! ARGV
rescue OptionParser::InvalidOption, OptionParser::InvalidArgument,
       OptionParser::MissingArgument => e
  puts opts
  puts
  puts e.message
  exit(-1)
end

# Make sure the user specified at least one file
if ARGV.length < 1
  puts opts
  puts ""
  puts "Must specify a script to run"
  exit(-1)
end


# Install at_exit handler.  It is important that we do this 
# before loading the scripts so our at_exit handler run
# *after* any other one that will be installed. 

at_exit {
  # Stop profiling
  result = RubyProf.stop

  # Create a printer
  printer = options.printer.new(result)

  # Get output
  if options.file
    File.open(options.file, 'w') do |file|
      printer.print(file, options.min_percent)
    end
  else
    # Print out results 
    printer.print(STDOUT, options.min_percent)
  end
}

# Now set measure mode
RubyProf.measure_mode = options.measure_mode

# Get the script we will execute
script = ARGV.shift
if options.replace_prog_name
        $0 = File.basename(File.expand_path(script))
end

# Start profiling
RubyProf.start 

# Load the script
load script