#!/usr/bin/env ruby

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

class Riemann::Tools::Elasticsearch
  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 :es_host, 'Elasticsearch host', default: "localhost"
  opt :es_port, 'Elasticsearch port', type: :int, default: 9200


  # Handles HTTP connections and GET requests safely
  def safe_get(uri)
      # 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 => uri.host,
          :service => "elasticsearch health",
          :state => "critical",
          :description => "HTTP connection error: #{e.class} - #{e.message}"
        )
      end
      response
  end

  def health_url
    "http://#{options[:es_host]}:#{options[:es_port]}/_cluster/health"
  end

  def tick
    uri = URI(health_url)
    response = safe_get(uri)

    return if response.nil?

    if response.status != 200
        report(:host => uri.host,
          :service => "elasticsearch health",
          :state => "critical",
          :description => "HTTP connection error: #{response.status} - #{response.body}"
        )
    else
      # Assuming that a 200 will give json
      json = JSON.parse(response.body)
      cluster_name = json.delete("cluster_name")
      cluster_status = json.delete("status")
      state = case cluster_status
      when "green"
        "ok"
      when "yellow"
        "warning"
      when "red"
        "critical"
      end

      report(:host => uri.host,
             :service => "elasticsearch health",
             :state => state,
             :description => "Elasticsearch cluster: #{cluster_name} - #{cluster_status}")

      json.each_pair do |k,v|
        report(:host => uri.host,
               :service => "elasticsearch #{k}",
               :metric => v,
               :description => "Elasticsearch cluster #{k}"
        )

      end
    end
  end



end
Riemann::Tools::Elasticsearch.run