require 'rack'
require 'perftools'
require 'rbconfig'
require 'pstore'

# REQUIREMENTS

# You'll need graphviz to generate call graphs using dot (for the GIF printer):
#
#    sudo port install graphviz     # osx
#    sudo apt-get install graphviz  # debian/ubuntu

# You'll need ps2pdf to generate PDFs 
# On OS X, ps2pdf comes is installed as part of Ghostscript
# 
#    sudo port install ghostscript # osx
#    brew install ghostscript      # homebrew
#    sudo apt-get install ps2pdf   # debian/ubuntu

# USAGE

# To configure your Rack app to use PerftoolsProfiler, call the 'use' method
#
#    use Rack::PerftoolsProfiler, options
# 
# For example:
#    use Rack::PerftoolsProfiler, :default_printer => 'gif'

# OPTIONS

# :default_printer - can be set to 'text', 'gif', or 'pdf'. Default is :text
# :mode            - can be set to 'cputime' or 'walltime'. Default is :cputime
# :frequency       - in :cputime mode, the number of times per second the app will be sampled.
#                    Default is 100 (times/sec)

# MODES 

# There are two modes for the profiler
# 
# First, you can run in 'simple' mode. Just visit the url you want to profile, but
# add the 'profile' and (optionally) the 'times' GET params
# 
# example: 
# curl http://localhost:8080/foobar?profile=true&times=3
#
# Note that this will change the status, body, and headers of the response (you'll get
# back the profiling data, NOT the original response.
#
#
# The other mode is start/stop mode.
# 
# example:
# curl http://localhost:8080/__start__
# curl http://localhost:8080/foobar
# curl http://localhost:8080/foobaz
# curl http://localhost:8080/__stop__
# curl http://localhost:8080/__data__
#
# In this mode, all responses are normal. You must visit __stop__ to complete profiling and
# then you can view the profiling data by visiting __data__

# PROFILING DATA OPTIONS
#
# In both simple and start/stop modes, you can add additional params to change how the data
# is displayed. In simple mode, these params are just added to the URL being profiled. In
# start/stop mode, they are added to the __data__ URL

# printer - overrides the default_printer option (see above)
# ignore  - a regular expression of the area of code to ignore 
# focus   - a regular expression of the area of code to solely focus on.

# (for ignore and focus, please see http://google-perftools.googlecode.com/svn/trunk/doc/cpuprofile.html#pprof
# for more details)

module Rack

  class PerftoolsProfiler
    include Rack::Utils

    PRINTER_CONTENT_TYPE = {
      :text => 'text/plain',
      :gif => 'image/gif',
      :pdf => 'application/pdf'
    }
    
    def self.clear_data
      Profiler.clear_data
    end

    def self.with_profiling_off(app, options = {})
      instance = self.new(app, options)
      instance.force_stop
      instance
    end

    def initialize(app, options = {})
      @app = app
      @profiler = Profiler.new(@app, options.clone)
    end
    
    def call(env)
      action = Action.for_env(env.clone, @profiler, self)
      action.act
      action.response
    end

    def call_app(env)
      @app.call(env)
    end

    def force_stop
      @profiler.stop
    end

    def profiler_data_response(profiling_data)
      format, body = profiling_data
      if format==:none
        [404, {'Content-Type' => 'text/plain'}, ['No profiling data available.']] 
      else
        [200, headers(format, body), Array(body)]
       end
    end

    private

    def headers(printer, body)
      headers = { 
        'Content-Type' => PRINTER_CONTENT_TYPE[printer],
        'Content-Length' => content_length(body)
      }
      if printer==:pdf
        filetype = printer
        filename='profile_data'
        headers['Content-Disposition'] = %(attachment; filename="#{filename}.#{filetype}")
      end
      headers
    end

    def content_length(body)
      body.inject(0) { |len, part| len + bytesize(part) }.to_s
    end


  end

  class Profiler

    def self.tmpdir
      dir = nil
      Dir.chdir Dir.tmpdir do dir = Dir.pwd end # HACK FOR OSX
      dir
    end

    PROFILING_DATA_FILE = ::File.join(self.tmpdir, 'rack_perftools_profiler.prof')
    PROFILING_SETTINGS_FILE = ::File.join(self.tmpdir, 'rack_perftools_profiler.config')
    DEFAULT_PRINTER = :text
    DEFAULT_MODE = :cputime
    UNSET_FREQUENCY = -1

    def initialize(app, options)
      @printer   = (options.delete(:default_printer) { DEFAULT_PRINTER }).to_sym
      @frequency = (options.delete(:frequency) { UNSET_FREQUENCY }).to_s
      @mode      = (options.delete(:mode) { DEFAULT_MODE }).to_sym
      raise ArgumentError, "Invalid option(s): #{options.keys.join(' ')}" unless options.empty?
    end
    
    def profile
      start
      yield
    ensure
      stop
    end

    def self.clear_data
      ::File.delete(PROFILING_DATA_FILE) if ::File.exists?(PROFILING_DATA_FILE)
    end

    def start
      set_env_vars
      PerfTools::CpuProfiler.start(PROFILING_DATA_FILE)
      self.profiling = true
    end

    def stop
      PerfTools::CpuProfiler.stop
      self.profiling = false
      unset_env_vars
    end

    def profiling?
      pstore_transaction(true) do |store|
        store[:profiling?]
      end
    end

    def data(options = {})
      printer = (options.fetch('printer') {@printer}).to_sym
      ignore = options.fetch('ignore') { nil }
      focus = options.fetch('focus') { nil }
      if ::File.exists?(PROFILING_DATA_FILE)
        args = "--#{printer}"
        args += " --ignore=#{ignore}" if ignore
        args += " --focus=#{focus}" if focus
        cmd = "pprof.rb #{args} #{PROFILING_DATA_FILE}"
        [printer, `#{cmd}`]
      else
        [:none, nil]
      end
    end

    private 

    def set_env_vars
      ENV['CPUPROFILE_REALTIME'] = '1' if @mode == :walltime
      ENV['CPUPROFILE_FREQUENCY'] = @frequency if @frequency != UNSET_FREQUENCY
    end

    def unset_env_vars
      ENV.delete('CPUPROFILE_REALTIME')
      ENV.delete('CPUPROFILE_FREQUENCY')
    end
    
    def profiling=(value)
      pstore_transaction(false) do |store|
        store[:profiling?] = value
      end
    end

    def pstore_transaction(read_only)
      pstore = PStore.new(PROFILING_SETTINGS_FILE)
      pstore.transaction(read_only) do
        yield pstore if block_given?
      end
    end

  end

  class Action

    def initialize(env, profiler, middleware)
      @env = env
      @request = Request.new(env)
      @data_params = @request.params.clone
      @profiler = profiler
      @middleware = middleware
    end
    
    def act
      # do nothing
    end

    def self.for_env(env, profiler, middleware)
      request = Request.new(env)
      klass = 
        case request.path
        when '/__start__'
          StartProfiling
        when '/__stop__'
          StopProfiling
        when '/__data__'
          ReturnData
        else
          if ProfileOnce.has_special_param?(request)
            ProfileOnce
          else
            CallAppDirectly
          end
        end
      klass.new(env, profiler, middleware)
    end

  end

  class StartProfiling < Action

    def act
      @profiler.start
    end

    def response
      [200, {'Content-Type' => 'text/plain'}, 
       [<<-EOS
Profiling is now enabled.
Visit the URLS that should be profiled.
When you are finished, visit /__stop__, then visit /__data__ to view the results.
EOS
       ]]
    end

  end

  class StopProfiling < Action
    
    def act
      @profiler.stop
    end

    def response
      [200, {'Content-Type' => 'text/plain'}, 
       [<<-EOS
Profiling is now disabled.
Visit /__data__ to view the results.
EOS
       ]]
    end

  end

  class ReturnData < Action

    def response
      if @profiler.profiling?
        [400, {'Content-Type' => 'text/plain'}, ['No profiling data available.']]
      else
        @middleware.profiler_data_response(@profiler.data(@data_params))
      end
    end

  end

  class ProfileOnce < Action
    include Rack::Utils

    def self.has_special_param?(request)
      request.params['profile'] != nil
    end

    def initialize(*args)
      super
      @times = (Request.new(@env).params.fetch('times') {1}).to_i
      @new_env = delete_custom_params(@env)
    end
    
    def act
      @profiler.profile do
        @times.times { @middleware.call_app(@new_env) }
      end
    end

    def response
      @middleware.profiler_data_response(@profiler.data(@data_params))
    end

    def delete_custom_params(env)
      new_env = env.clone
      
      params = Request.new(new_env).params
      params.delete('profile')
      params.delete('times')
      params.delete('printer')
      params.delete('ignore')
      params.delete('focus')

      new_env.delete('rack.request.query_string')
      new_env.delete('rack.request.query_hash')

      new_env['QUERY_STRING'] = build_query(params)
      new_env
    end

  end

  class CallAppDirectly < Action
    
    def act
      @result = @middleware.call_app(@env)
    end

    def response
      @result
    end

  end

end