#!/usr/bin/env ruby

# Connects to a server (first arg) and populates it with a constant stream of
# events for testing.

require 'riemann/client'
require 'pp'

class Riemann::Bench
  attr_accessor :client, :hosts, :services, :states
  def initialize
    @hosts = [nil] + (0...10).map { |i| "host#{i}" }
    @hosts = ['test']
    @services = %w(per)
    @states = {}
    @client = Riemann::Client.new(host: (ARGV.first || 'localhost'))
  end

  def evolve(state)
    m = rand
    s = case m
    when 0...0.75
      'ok'
    when 0.75...0.9
      'warning'
    when 0.9..1.0
      'critical'
    end

    {
      metric_f: m,
      state: s,
      host: state[:host],
      service: state[:service],
      description: "at #{Time.now}"
    }
  end
  
  def tick
#    pp @states
    hosts.product(services).each do |id|
      client << (states[id] = evolve(states[id]))
    end
  end

  def run
    start
    loop do
#      sleep 0.01
      tick
    end
  end

  def start
    hosts.product(services).each do |host, service|
      states[[host, service]] = {
        metric_f: 0.5,
        state: 'ok',
        description: "Starting up",
        host: host,
        service: service
      }
    end
  end
end

Riemann::Bench.new.run