#!/usr/bin/env ruby

require 'rake'
require 'rake/testtask'
require 'fileutils'

module RubyProf

  # Define a task library for profiling unit tests with ruby-prof.
  #
  # All of the options provided by
  # the Rake:TestTask are supported except the loader
  # which is set to ruby-prof.  For detailed information
  # please refer to the Rake:TestTask documentation.
  #
  # ruby-prof specific options include:
  #
  #   output_dir - For each file specified an output 
  #                file with profile information will be
  #                written to the output directory.
  #                By default, the output directory is
  #                called "profile" and is created underneath
  #                the current working directory.
  #
  #   printer - Specifies the output printer.  Valid values include
  #             :flat, :graph, :graph_html and :call_tree.
  #
  #   min_percent - Methods that take less than the specified percent
  #                 will not be written out.
  #
  # Example:
  #   
  #   require 'ruby-prof/task'
  #   
  #   RubyProf::ProfileTask.new do |t|
  #     t.test_files = FileList['test/test*.rb']
  #     t.output_dir = "c:/temp"
  #     t.printer = :graph
  #     t.min_percent = 10
  #   end
  #
  # If rake is invoked with a "TEST=filename" command line option,
  # then the list of test files will be overridden to include only the
  # filename specified on the command line.  This provides an easy way
  # to run just one test.
  #
  # If rake is invoked with a "TESTOPTS=options" command line option,
  # then the given options are passed to the test process after a
  # '--'.  This allows Test::Unit options to be passed to the test
  # suite.
  #
  # Examples:
  #
  #   rake profile                           # run tests normally
  #   rake profile TEST=just_one_file.rb     # run just one test file.
  #   rake profile TESTOPTS="-v"             # run in verbose mode
  #   rake profile TESTOPTS="--runner=fox"   # use the fox test runner
  
  class ProfileTask < Rake::TestTask
    attr_accessor :output_dir 
    attr_accessor :min_percent 
    attr_accessor :printer
    
    def initialize(name = :profile)
      super(name)
    end
    
    # Create the tasks defined by this task lib.
    def define
      lib_path = @libs.join(File::PATH_SEPARATOR)
      desc "Profile" + (@name==:profile ? "" : " for #{@name}")
      
      task @name do
        create_output_directory
        
        @ruby_opts.unshift( "-I#{lib_path}" )
        @ruby_opts.unshift( "-w" ) if @warning
        @ruby_opts.push("-S ruby-prof")
        @ruby_opts.push("--printer #{@printer}")
        @ruby_opts.push("--min_percent #{@min_percent}")

        file_list.each do |file_path|  
          run_script(file_path)
        end
      end
      self
    end
    
    # Run script
    def run_script(script_path)
      run_code = ''
      RakeFileUtils.verbose(@verbose) do
        file_name = File.basename(script_path, File.extname(script_path))
        case @printer
          when :flat, :graph, :call_tree
            file_name += ".txt"
          when :graph_html
            file_name += ".html"
          else
            file_name += ".txt"
        end
          
        output_file_path = File.join(output_directory, file_name)
          
        command_line = @ruby_opts.join(" ") + 
                      " --file=" + output_file_path +
                      " " + script_path

        puts "ruby " + command_line 
        # We have to catch the exeption to continue on.  However,
        # the error message will have been output to STDERR
        # already by the time we get here so we don't have to
        # do that again
        begin
          ruby command_line
        rescue => e
          STDOUT << e << "\n"
          STDOUT.flush
        end
        puts ""
        puts ""
      end
    end

    def output_directory
      File.expand_path(@output_dir)
    end
    
    def create_output_directory
      if not File.exist?(output_directory)
        Dir.mkdir(output_directory)
      end
    end

    def clean_output_directory
      if File.exist?(output_directory)
        files = Dir.glob(output_directory + '/*')
        FileUtils.rm(files)
      end
    end
    
    def option_list # :nodoc:
      ENV['OPTIONS'] || @options.join(" ") || ""
    end
  end
end