#!/usr/bin/env ruby
#

require File.expand_path('../../lib/riemann/tools', __FILE__)

class Riemann::Tools::Rabbitmq
  include Riemann::Tools

  require 'faraday'
  require 'json'
  require 'uri'


  opt :read_timeout, 'Faraday read timeout', type: :int, default: 2
  opt :open_timeout, 'Faraday open timeout', type: :int, default: 1

  opt :monitor_user, 'RabbitMQ monitoring user', type: :string
  opt :monitor_pass, 'RabbitMQ monitoring user password', type: :string
  opt :monitor_port, 'RabbitMQ monitoring port', default: 15672
  opt :monitor_host, 'RabbitMQ monitoring host', default: "localhost"

  def monitor_url
    "http://#{options[:monitor_user]}:#{options[:monitor_pass]}@#{options[:monitor_host]}:#{options[:monitor_port]}/api/overview"
  end

  def event_host
    if options[:event_host]
      return options[:event_host]
    else
      return options[:monitor_host]
    end
  end

  def safe_get(uri, event_host)
      # Handle connection timeouts
      response = nil
      begin
        connection = Faraday.new(uri)
        response = connection.get do |req|
          req.options[:timeout] = options[:read_timeout]
          req.options[:open_timeout] = options[:open_timeout]
        end
      rescue => e
        report(:host => event_host,
          :service => "rabbitmq monitoring",
          :state => "critical",
          :description => "HTTP connection error: #{e.class} - #{e.message}"
        )
      end
      response
  end

  def tick
    uri = URI(monitor_url)
    response = safe_get(uri, event_host)

    return if response.nil?

    json = JSON.parse(response.body)

    if response.status != 200
        report(:host => event_host,
          :service => "rabbitmq",
          :state => "critical",
          :description => "HTTP connection error: #{response.status} - #{response.body}"
        )
    else
      report(:host => event_host,
        :service => "rabbitmq monitoring",
        :state => "ok",
        :description => "HTTP connection ok"
      )

      %w( message_stats queue_totals object_totals ).each do |stat|
        # NOTE / BUG ?
        # Brand new servers can have blank message stats. Is this ok?
        # I can't decide.
        next if json[stat].empty?
        json[stat].each_pair do |k,v|
          service = "rabbitmq.#{stat}.#{k}"
          if k =~ /details$/
            metric = v['rate']
          else
            metric = v
          end

          # TODO: Set state via thresholds which can be configured

          report(:host => event_host,
                 :service => service,
                 :metric => metric,
                 :description => "RabbitMQ monitor"
                )
        end
      end
    end
  end
end
Riemann::Tools::Rabbitmq.run