#!/usr/bin/env ruby

$: << File.join(File.dirname(__FILE__), '..', 'lib')
require 'benchmark'
require 'moneta'
require 'fileutils'

class String
  def random(n)
    (1..n).map { self[rand(size),1] }.join
  end
end

class Array
  def sum
    inject(0, &:+)
  end

  def mean
    sum / size
  end

  def stddev
    m = mean
    Math.sqrt(map {|s| (s - m) ** 2 }.mean)
  end
end

class MonetaBenchmarks
  DIR = __FILE__ + '.tmp'
  FileUtils.mkpath(DIR)

  STORES = {
    # SDBM is unstable
    # :SDBM => { :file => "#{DIR}/sdbm" },
    # YAML is so fucking slow
    # :YAML => { :file => "#{DIR}/yaml" },
    :ActiveRecord => { :connection => { :adapter  => 'sqlite3', :database => ':memory:' } },
    :Cassandra => {},
    :Client => {},
    :Couch => {},
    :DBM => { :file => "#{DIR}/dbm" },
    :DataMapper => { :setup => "sqlite3:#{DIR}/datamapper" },
    :Daybreak => { :file => "#{DIR}/daybreak" },
    :File => { :dir => "#{DIR}/file" },
    :GDBM => { :file => "#{DIR}/gdbm" },
    :HBase => {},
    :HashFile => { :dir => "#{DIR}/hashfile" },
    :KyotoCabinet => { :file => 'kyotocabinet.kch' },
    :LRUHash => {},
    :LevelDB => { :dir => "#{DIR}/leveldb" },
    :LocalMemCache => { :file => "#{DIR}/lmc" },
    :MemcachedDalli => {},
    :MemcachedNative => {},
    :Memory => {},
    :Mongo => {},
    :PStore => { :file => "#{DIR}/pstore" },
    :Redis => {},
    :RestClient => { :url => 'http://localhost:8808/' },
    :Riak => {},
    :Sequel => { :db => 'sqlite:/' },
    :Sqlite => { :file => ':memory:' },
    :TDB => { :file => "#{DIR}/tdb" },
    :TokyoCabinet => { :file => 'tokyocabinet' },
  }

  CONFIGS = {
    :uniform_small => {
      :runs => 3,
      :keys => 1000,
      :min_key_len => 1,
      :max_key_len => 32,
      :key_dist => :uniform,
      :min_val_len => 0,
      :max_val_len => 256,
      :val_dist => :uniform
    },
    :uniform_medium => {
      :runs => 3,
      :keys => 1000,
      :min_key_len => 3,
      :max_key_len => 200,
      :key_dist => :uniform,
      :min_val_len => 0,
      :max_val_len => 1024,
      :val_dist => :uniform
    },
    :uniform_large => {
      :runs => 3,
      :keys => 100,
      :min_key_len => 3,
      :max_key_len => 200,
      :key_dist => :uniform,
      :min_val_len => 0,
      :max_val_len => 10240,
      :val_dist => :uniform
    },
    :normal_small => {
      :runs => 3,
      :keys => 1000,
      :min_key_len => 1,
      :max_key_len => 32,
      :key_dist => :normal,
      :min_val_len => 0,
      :max_val_len => 256,
      :val_dist => :normal
    },
    :normal_medium => {
      :runs => 3,
      :keys => 1000,
      :min_key_len => 3,
      :max_key_len => 200,
      :key_dist => :normal,
      :min_val_len => 0,
      :max_val_len => 1024,
      :val_dist => :normal
    },
    :normal_large => {
      :runs => 3,
      :keys => 100,
      :min_key_len => 3,
      :max_key_len => 200,
      :key_dist => :normal,
      :min_val_len => 0,
      :max_val_len => 10240,
      :val_dist => :normal
    },
  }

  DICT = 'ABCDEFGHIJKLNOPQRSTUVWXYZabcdefghijklnopqrstuvwxyz123456789'.freeze
  HEADER = "\n                         Minimum  Maximum    Total     Mean   Stddev    Ops/s"
  SEPARATOR = '=' * 77

  module Rand
    extend self

    def normal_rand(mean, stddev)
      # Box-Muller transform
      theta = 2 * Math::PI * (rand(1e10) / 1e10)
      scale = stddev * Math.sqrt(-2 * Math.log(1 - (rand(1e10) / 1e10)))
      [mean + scale * Math.cos(theta),
       mean + scale * Math.sin(theta)]
    end

    def uniform(min, max)
      rand(max - min) + min
    end

    def normal(min, max)
      mean = (min + max) / 2
      stddev = (max - min) / 4
      loop do
        val = normal_rand(mean, stddev)
        return val.first if val.first >= min && val.first <= max
        return val.last if val.last >= min && val.last <= max
      end
    end
  end

  def parallel(&block)
    if defined?(JRUBY_VERSION)
      Thread.new(&block)
    else
      Process.fork(&block)
    end
  end

  def write_histogram(file, sizes)
    min = sizes.min
    delta = sizes.max - min
    histogram = []
    sizes.each do |s|
      s = 10 * (s - min) / delta
      histogram[s] ||= 0
      histogram[s] += 1
    end
    File.open(file, 'w') do |f|
      histogram.each_with_index { |n,i| f.puts "#{i*delta/10+min} #{n}" }
    end
  end

  def start_servers
    parallel do
      begin
        Moneta::Server.new(Moneta.new(:Memory)).run
      rescue Exception => ex
        puts "\e[31mFailed to start Moneta server - #{ex.message}\e[0m"
      end
    end

    parallel do
      begin
        require 'rack'
        require 'webrick'
        require 'rack/moneta_rest'

        # Keep webrick quiet
        ::WEBrick::HTTPServer.class_eval do
          def access_log(config, req, res); end
        end
        ::WEBrick::BasicLog.class_eval do
          def log(level, data); end
        end

        Rack::Server.start(:app => Rack::Builder.app do
                             use Rack::Lint
                             run Rack::MonetaRest.new(:store => :Memory)
                           end,
                           :environment => :none,
                           :server => :webrick,
                           :Port => 8808)
      rescue Exception => ex
        puts "\e[31mFailed to start Rack server - #{ex.message}\e[0m"
      end
    end

    sleep 1 # Wait for servers
  end

  def test_stores
    STORES.each do |name, options|
      begin
        if name == :DataMapper
          begin
            require 'dm-core'
            DataMapper.setup(:default, :adapter => :in_memory)
          rescue LoadError => ex
            puts "\e[31mFailed to load DataMapper - #{ex.message}\e[0m"
          end
        elsif name == :Riak
          require 'riak'
          Riak.disable_list_keys_warnings = true
        end

        cache = Moneta.new(name, options.dup)
        cache['test'] = 'test'
      rescue Exception => ex
        puts "\e[31m#{name} not benchmarked - #{ex.message}\e[0m"
        STORES.delete(name)
      ensure
        (cache.close rescue nil) if cache
      end
    end
  end

  def generate_data
    until @data.size == @config[:keys]
      key = DICT.random(Rand.send(@config[:key_dist], @config[:min_key_len], @config[:max_key_len]))
      @data[key] = DICT.random(Rand.send(@config[:val_dist], @config[:min_val_len], @config[:max_val_len]))
    end

    key_lens, val_lens = @data.keys.map(&:size), @data.values.map(&:size)
    @data = @data.to_a

    write_histogram("#{DIR}/key.histogram", key_lens)
    write_histogram("#{DIR}/value.histogram", val_lens)

    puts "\n\e[1m\e[34m#{SEPARATOR}\n\e[34mComputing keys and values...\n\e[34m#{SEPARATOR}\e[0m"
    puts %{                         Minimum  Maximum    Total     Mean   Stddev}
    puts 'Key Length              % 8d % 8d % 8d % 8d % 8d' % [key_lens.min, key_lens.max, key_lens.sum, key_lens.mean, key_lens.stddev]
    puts 'Value Length            % 8d % 8d % 8d % 8d % 8d' % [val_lens.min, val_lens.max, val_lens.sum, val_lens.mean, val_lens.stddev]
  end

  def print_config
    puts "\e[1m\e[36m#{SEPARATOR}\n\e[36mConfig #{@config_name}\n\e[36m#{SEPARATOR}\e[0m"
    @config.each do |k,v|
      puts '%-16s = %-10s' % [k,v]
    end
  end

  def print_store_stats(name)
    puts HEADER
    [:write, :read, :sum].each do |i|
      ops = (1000 * @config[:runs] * @data.size) / @stats[name][i].sum
      line = '%-17.17s %-5s % 8d % 8d % 8d % 8d % 8d % 8d' %
        [name, i, @stats[name][i].min, @stats[name][i].max, @stats[name][i].sum,
         @stats[name][i].mean, @stats[name][i].stddev, ops]
      @summary << [-ops, line << "\n"] if i == :sum
      puts line
    end

    errors = @stats[name][:error].sum
    if errors > 0
      puts "\e[31m%-23.23s % 8d % 8d % 8d % 8d\e[0m" %
        ['Read errors', @stats[name][:error].min, @stats[name][:error].max, errors, errors / @config[:runs]]
    else
      puts "\e[32mNo read errors"
    end
  end

  def benchmark_store(name, options)
    puts "\n\e[1m\e[34m#{SEPARATOR}\n\e[34m#{name}\n\e[34m#{SEPARATOR}\e[0m"

    store = Moneta.new(name, options.dup)

    @stats[name] = {
      :write => [],
      :read => [],
      :sum => [],
      :error => []
    }

    %w(Rehearse Measure).each do |type|
      state = ''
      print "%s [%#{2 * @config[:runs]}s] " % [type, state]

      @config[:runs].times do |run|
        store.clear
        print "%s[%-#{2 * @config[:runs]}s] " % ["\b" * (2 * @config[:runs] + 3), state << 'W']

        @data.shuffle!
        m1 = Benchmark.measure do
          @data.each {|k,v| store[k] = v }
        end

        print "%s[%-#{2 * @config[:runs]}s] " % ["\b" * (2 * @config[:runs] + 3), state << 'R']

        @data.shuffle!
        error = 0
        m2 = Benchmark.measure do
          @data.each do |k, v|
            error += 1 if v != store[k]
          end
        end

        if type == 'Measure'
          @stats[name][:write] << m1.real * 1000
          @stats[name][:error] << error
          @stats[name][:read] << m2.real * 1000
          @stats[name][:sum] << (m1.real + m2.real) * 1000
        end
      end
    end

    print_store_stats(name)
  rescue StandardError => ex
    puts "\n\e[31mFailed to benchmark #{name} - #{ex.message}\e[0m\n"
  ensure
    store.close if store
  end

  def run_benchmarks
    STORES.each {|name, options| benchmark_store(name, options) }
  end

  def print_summary
    puts "\n\e[1m\e[36m#{SEPARATOR}\n\e[36mSummary #{@config_name}: #{@config[:runs]} runs, #{@data.size} keys\n\e[36m#{SEPARATOR}\e[0m#{HEADER}\n"
    @summary.sort_by(&:first).each do |entry|
      puts entry.last
    end
  end

  def initialize(args)
    @config_name = args.size == 1 ? args.first.to_sym : :uniform_medium
    unless @config = CONFIGS[@config_name]
      puts "Configuration #{@config_name} not found"
      exit
    end

    # Disable jruby stdout pollution by memcached
    if defined?(JRUBY_VERSION)
      require 'java'
      properties = java.lang.System.getProperties();
      properties.put('net.spy.log.LoggerImpl', 'net.spy.memcached.compat.log.SunLogger');
      java.lang.System.setProperties(properties);
      java.util.logging.Logger.getLogger('').setLevel(java.util.logging.Level::OFF)
    end

    @stats, @data, @summary = {}, {}, []
  end

  def run
    start_servers
    test_stores
    print_config
    generate_data
    run_benchmarks
    print_summary
  end
end

MonetaBenchmarks.new(ARGV).run