#!/usr/bin/env ruby
require "rubygems"
require "bundler/setup"

$:.unshift(File.expand_path(File.dirname(__FILE__) + '/../lib'))

require 'pp'
require 'bunny'
require 'yajl'
require 'uuidtools'
require 'word_salad'

require 'opscode/expander/configuration'

Opscode::Expander.init_config(ARGV)

MESSAGES_TO_SEND = 10_000

NUM_RAND_KEY_PAIRS = 50
NUM_RAND_VALUE_PAIRS = 50

PERSISTENT_MESSAGES = true

KEYS = NUM_RAND_VALUE_PAIRS.words

SAMPLE_NODES = []
Dir.glob(File.expand_path(File.dirname(__FILE__)) + '/../data/*_node.json') do |node_file|
  SAMPLE_NODES << Yajl::Parser.parse(IO.read(node_file))
end

NUM_NODES = SAMPLE_NODES.size
puts "Read #{NUM_NODES} sample nodes"

puts "Using rabbitmq config #{Opscode::Expander.config.amqp_config.inspect}"

puts "connecting to rabbitmq"
amqp_client = Bunny.new(Opscode::Expander.config.amqp_config)
amqp_client.start

puts 'declaring queues'
queues = {}
0.upto(1023) do |vnode|
  queues[vnode] = amqp_client.queue("vnode-#{vnode}", :durable => true)
end
  
def add_rand_keys(node)
  rand_key_vals = Hash[*((2 * NUM_RAND_KEY_PAIRS).words)]
  rand_vals = Hash[*(KEYS.zip(NUM_RAND_VALUE_PAIRS.words)).flatten]
  node.merge(rand_key_vals.merge(rand_vals))
end

puts "sending #{MESSAGES_TO_SEND} messages"
start_time = Time.now
sent_messages = 0
1.upto(MESSAGES_TO_SEND) do
  node = SAMPLE_NODES[rand(NUM_NODES)]
  node = add_rand_keys(node)
  index_data = {:action => :add}
  index_data[:payload] = {:item => node}
  index_data[:payload][:type] = :node
  index_data[:payload][:database] = :testdb
  index_data[:payload][:enqueued_at] = Time.now.utc.to_i

  id = node["name"]
  vnode = rand(1024)
  index_data[:payload][:id] = id

  puts "queue: vnode-#{vnode} (#{sent_messages} / #{MESSAGES_TO_SEND})"
  amqp_client.tx_select if PERSISTENT_MESSAGES
  queues[vnode].publish(Yajl::Encoder.encode(index_data), :persistent => PERSISTENT_MESSAGES)
  amqp_client.tx_commit if PERSISTENT_MESSAGES
  sent_messages += 1
end
end_time = Time.now

total_time = end_time - start_time
rate = MESSAGES_TO_SEND.to_f / total_time
puts "done (#{total_time}s, #{rate} msg/s)"
