#! /usr/bin/env ruby

# First require ruby-prof
require 'ruby-prof'

# Now setup option parser
require 'ostruct'
require 'optparse'

module RubyProf
  # == Synopsis
  #
  # Profiles a Ruby program.
  #
  # == Usage
  # ruby-prof [options] <script.rb> [--] [profiled-script-command-line-options]
  #
  # Options:
  #        --allow_exceptions           Raise exceptions encountered during profiling (true) or suppress them (false)
  #    -E, --eval-noprof=code           execute the ruby statements (not profiled)
  #        --exclude=methods    A comma separated list of methods to exclude.
  #                                       Specify instance methods via # (Integer#times)
  #                                       Specify class methods via . (Integer.superclass)
  #        --exclude-common             Remove common methods from the profile
  #    -f, --file=path                  Output results to a file instead of standard out.
  #    -m, --min_percent=min_percent    The minimum percent a method must take before
  #                                       being included in output reports.
  #                                       This option is not supported for call tree.
  #        --mode=measure_mode          Select what ruby-prof should measure:
  #                                       wall - Wall time (default).
  #                                       process - Process time.
  #                                       allocations - Object allocations (requires patched Ruby interpreter).
  #                                       memory - Allocated memory in KB (requires patched Ruby interpreter).
  #    -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
  #                                       call_stack - prints a HTML visualization of the call tree
  #                                       dot - Prints a graph profile as a dot file
  #                                       multi - Creates several reports in output directory
  #    -R, --require-noprof=lib         require a specific library (not profiled)
  #    -s, --sort=sort_mode             Select how ruby-prof results should be sorted:
  #                                       total - Total time
  #                                       self - Self time
  #                                       wait - Wait time
  #                                       child - Child time
  #        --track_allocations          Track allocations while profiling
  #    -v, --version version            Show version (1.1.0)
  #    -h, --help                       Show help message

  class Cmd
    # :enddoc:
    attr_accessor :options
    attr_reader :profile

    def initialize
      setup_options
      parse_args

      load_pre_libs
      load_pre_execs
    end

    def setup_options
      @options = OpenStruct.new
      options.allow_exceptions = false
      options.exclude = Array.new
      options.exclude_common = false
      options.file = nil
      options.measure_mode = RubyProf::WALL_TIME
      options.min_percent = 0
      options.pre_libs = Array.new
      options.pre_execs = Array.new
      options.printer = RubyProf::FlatPrinter
      options.track_allocations = false
    end

    # This is copied from ActiveSupport:
    def constantize(camel_cased_word)
      if !camel_cased_word.include?("::")
        Object.const_get(camel_cased_word)
      else
        names = camel_cased_word.split("::")

        # Trigger a built-in NameError exception including the ill-formed constant in the message.
        Object.const_get(camel_cased_word) if names.empty?

        # Remove the first blank element in case of '::ClassName' notation.
        names.shift if names.size > 1 && names.first.empty?

        names.inject(Object) do |constant, name|
          if constant == Object
            constant.const_get(name)
          else
            candidate = constant.const_get(name)
            next candidate if constant.const_defined?(name, false)
            next candidate unless Object.const_defined?(name)

            # Go down the ancestors to check if it is owned directly. The check
            # stops when we reach Object or the end of ancestors tree.
            constant = constant.ancestors.inject(constant) do |const, ancestor|
              break const    if ancestor == Object
              break ancestor if ancestor.const_defined?(name, false)
              const
            end

            # owner is in Object, so raise
            constant.const_get(name, false)
          end
        end
      end
    end

    def option_parser
      OptionParser.new do |opts|
        opts.banner = "ruby_prof #{RubyProf::VERSION}\n" +
            "Usage: ruby-prof [options] <script.rb> [--] [profiled-script-command-line-options]"

        opts.separator ""
        opts.separator "Options:"

        opts.on('--allow_exceptions', 'Raise exceptions encountered during profiling (true) or suppress them (false)') do
          options.allow_exceptions = true
        end

        opts.on('-E code', '--eval-noprof=code', 'execute the ruby statements (not profiled)') do |code|
          options.pre_execs << code
        end

        opts.on('--exclude=methods', String,
                'A comma separated list of methods to exclude.',
                '  Specify instance methods via # (Integer#times)',
                '  Specify class methods via . (Integer.superclass)') do |exclude_string|
          exclude_string.split(',').each do |string|
            match = string.strip.match(/(.*)(#|\.)(.*)/)
            klass = constantize(match[1])
            if match[2] == '.'
              klass = klass.singleton_class
            end
            method = match[3].to_sym
            options.exclude << [klass, method]
          end
        end

        opts.on('--exclude-common', 'Remove common methods from the profile') do
          options.exclude_common = true
        end

        opts.on('-f path', '--file=path',
                'Output results to a file instead of standard out.') do |file|
          options.file = file
          options.old_wd = Dir.pwd
        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('--mode=measure_mode',
                [:process, :wall, :allocations, :memory],
                'Select what ruby-prof should measure:',
                '  wall - Wall time (default).',
                '  process - Process time.',
                '  allocations - Object allocations (requires patched Ruby interpreter).',
                '  memory - Allocated memory in KB (requires patched Ruby interpreter).') do |measure_mode|

          case measure_mode
          when :wall
            options.measure_mode = RubyProf::WALL_TIME
          when :process
            options.measure_mode = RubyProf::PROCESS_TIME
          when :allocations
            options.measure_mode = RubyProf::ALLOCATIONS
          when :memory
            options.measure_mode = RubyProf::MEMORY
          end
        end

        opts.on('-p printer', '--printer=printer', [:flat, :flat_with_line_numbers, :graph, :graph_html, :call_tree, :call_stack, :dot, :multi],
                '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',
                '  call_stack - prints a HTML visualization of the call tree',
                '  dot - Prints a graph profile as a dot file',
                '  multi - Creates several reports in output directory'
        ) do |printer|

          case printer
            when :flat
              options.printer = RubyProf::FlatPrinter
            when :graph
              options.printer = RubyProf::GraphPrinter
            when :graph_html
              options.printer = RubyProf::GraphHtmlPrinter
            when :call_tree
              options.printer = RubyProf::CallTreePrinter
            when :call_stack
              options.printer = RubyProf::CallStackPrinter
            when :dot
              options.printer = RubyProf::DotPrinter
            when :multi
              options.printer = RubyProf::MultiPrinter
          end
        end

        opts.on('-R lib', '--require-noprof=lib', 'require a specific library (not profiled)') do |lib|
          options.pre_libs << lib
        end

        opts.on('-s sort_mode', '--sort=sort_mode', [:total, :self, :wait, :child],
                'Select how ruby-prof results should be sorted:',
                '  total - Total time',
                '  self - Self time',
                '  wait - Wait time',
                '  child - Child time') do |sort_mode|

          options.sort_method = case sort_mode
                                when :total
                                  :total_time
                                when :self
                                  :self_time
                                when :wait
                                  :wait_time
                                when :child
                                  :children_time
                                end
        end

        opts.on('--track_allocations', 'Track allocations while profiling') do
          options.track_allocations = true
        end

        opts.on_tail("-v version", "--version", "Show version (#{RubyProf::VERSION})") do
          puts "ruby_prof " + RubyProf::VERSION
          exit
        end

        opts.on_tail("-h", "--help", "Show help message") do
          puts opts
          exit
        end
      end
    end

    def parse_args
      # Make sure the user specified at least one file
      if ARGV.length < 1 and not options.exec
        puts self.option_parser
        puts ""
        puts "Must specify a script to run"
        exit(-1)
      end

      self.option_parser.parse! ARGV

      if options.printer.needs_dir?
        options.file ||= "."
        options.old_wd ||= Dir.pwd
        if !File.directory?(options.file)
          puts "'#{options.file}' is not a directory"
          puts "#{options.printer} needs an existing directory path to put profiles under."
          exit(-1)
        end
      end
    rescue OptionParser::InvalidOption, OptionParser::InvalidArgument, OptionParser::MissingArgument => e
      puts self.option_parser
      puts e.message
      exit(-1)
    end

    def load_pre_libs
      options.pre_libs.each do |lib|
        require lib
      end
    end

    def load_pre_execs
      options.pre_execs.each do |exec|
        eval(exec)
      end
    end

    def run
      profile_options = {:allow_exceptions => options.allow_exceptions,
                         :exclude_common => options.exclude_common,
                         :measure_mode => options.measure_mode,
                         :track_allocations => options.track_allocations}

      @profile = Profile.new(**profile_options)

      options.exclude.each do |klass, method|
        @profile.exclude_method!(klass, method)
      end

      script = ARGV.shift
      profile.profile do
       load script
      end
    end
  end
end

# Parse command line options
cmd = RubyProf::Cmd.new

# 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 {
  # Create a printer
  printer = cmd.options.printer.new(cmd.profile)
  printer_options = {:min_percent => cmd.options.min_percent,
                     :sort_method => cmd.options.sort_method}

  # Get output
  if cmd.options.file
    # write it relative to the dir they *started* in, as it's a bit surprising to write it in the dir they end up in.
    Dir.chdir(cmd.options.old_wd) do
      if printer.class.needs_dir?
        printer.print(printer_options.merge(:path => cmd.options.file))
      else
        File.open(cmd.options.file, 'w') do |file|
          printer.print(file, printer_options)
        end
      end
    end
  else
    # Print out results
    printer.print(STDOUT, printer_options)
  end
}

# Now profile some code
cmd.run