require 'benchmark'

# DumbPro is just a stupid manual profiler that only profiles the methods you tell it to.
# Example:
#     require 'dumbpro'
#     Array.instrument :<<
#     arr = Array.new
#     arr << 1
#     # => 0ms >  Array#<< < 0ms
module DumbPro
  class Timer
    def self.current
      Thread.current['__dumbpro']
    end

    def self.current=(timer)
      Thread.current['__dumbpro'] = timer
    end

    attr_reader :depth

    def initialize(obj, meth, options, *args)
      @depth = 1
      @elapsed = 0
      @options = options
      if @options[:as]
        @label = @options[:as]
      else
        klass = Class === obj ? "#{obj.name}." : "#{obj.class.name}#"
        @label = "#{klass}#{meth}"
      end
      @args = args
      @parent = Timer.current
      if @parent
        @parent.add(self)
        @depth = @parent.depth + 1
      end
    end

    def invoke(&block)
      Timer.current = self
      result = nil
      @elapsed = Benchmark.realtime do
        result = block.call
      end
      result
    ensure
      Timer.current = @parent
      unless @parent
        report
      end
    end

    def add(child)
      children << child
    end

    def children
      @children ||= []
    end

    def elapsed
      @elapsed ||= 0
    end

    def report
      $stdout.puts "#{format_str % self_time}ms >#{'  ' * @depth}#{@label}#{args} < #{total_time}ms"
      children.each {|c| c.report}
    end

    def ms(t)
      (t * 1000).to_i
    end

    def format_str
      @format_str ||= @parent.format_str if @parent
      @format_str ||= "%#{total_time.to_s.length}s"
    end

    def self_time
      @self_time ||= ms(elapsed - children.inject(0) {|sum,el| sum += el.elapsed}).to_s
    end

    def total_time
      ms(elapsed)
    end

    def args
      if @options[:with]
        "(#{@args[@options[:with]].inspect})"
      end
    end
  end
end

class Module
  # Instrument the given methods on instances of this class.
  # Optional: pass a :as => "label" value as the last argument to the method to give
  # the instrumented method a custom label.
  def instrument(*meths)
    options = Hash === meths.last ? meths.pop : {}
    meths.each do |meth|
      unless instrumented?(meth)
        define_method("#{meth}__instrumented") { options }
        alias_method "#{meth}__orig_instrument", meth
        alias_method "#{meth}__without_instrument", meth
        module_eval(<<-METH, __FILE__, __LINE__)
          def #{meth}(*args, &block)
            DumbPro::Timer.new(self, "#{meth}", #{meth}__instrumented, *args).invoke do
              #{meth}__without_instrument(*args, &block)
            end
          end
        METH
      end
    end
  end

  # Is this method currently instrumented?
  def instrumented?(meth)
    instance_methods.include?("#{meth}__instrumented")
  end

  # Uninstrument the given methods
  def uninstrument(*meths)
    meths.each do |meth|
      if instrumented?(meth)
        remove_method "#{meth}__instrumented"
        alias_method meth, "#{meth}__orig_instrument"
      end
    end
  end

  # Uninstrument (turn off) all instrumented methods in this class or module
  def uninstrument_all
    uninstrument(*(instance_methods.select {|m| m =~ /__instrumented$/}.map {|m| m[0...-14]}))
  end
end

if __FILE__ == $0
  require 'test/unit'
  require 'stringio'

  class Foo
    def slow_method
      sleep 0.5
    end
    def hello
    end
  end
  class Bar
    def slow_method
      call_foo
    end

    def call_foo
      Foo.new.slow_method
    end

    def several_times
      10.times do
        slow_method
      end
    end

    def call(*args)
    end
  end

  class DumbProTest < Test::Unit::TestCase
    def capture
      @save = $stdout
      @stdout = $stdout = StringIO.new
    end
    
    def teardown
      Foo.uninstrument_all
      Bar.uninstrument_all
      $stdout = @save if @save
      puts @out if @out
    end

    def test_formatting
      Bar.instrument :slow_method, :several_times
      Foo.instrument :slow_method
      Bar.new.several_times
    end

    def test_output_profile
      Bar.instrument :slow_method
      Foo.instrument :slow_method
      capture
      Bar.new.slow_method
      @out = @stdout.string
      assert @out =~ /Bar#slow_method/
      assert @out =~ /Foo#slow_method/
    end

    def test_custom_label
      capture
      Foo.instrument :hello, :as => "Hello World!"
      Foo.new.hello
      assert @stdout.string =~ /Hello World!/
    end

    def test_instrument_uninstrument
      capture
      Foo.new.hello
      assert @stdout.string.empty?
      Foo.instrument :hello
      Foo.new.hello
      assert @stdout.string =~ /Foo#hello/
      @stdout.string = ''
      Foo.uninstrument :hello
      Foo.new.hello
      assert @stdout.string.empty?
    end

    def test_instrument_with_first_arg
      capture
      Bar.instrument :call, :with => 0
      Bar.new.call("abc")
      Bar.new.call("def")
      assert @stdout.string =~ /call.*"abc"/
      assert @stdout.string =~ /call.*"def"/
    end
  end
end