#!/usr/bin/env ruby
require 'rubygems'
require 'json'
require 'socket'
require 'optparse'
require 'open3'
require 'rake'

options = OpenStruct.new
OptionParser.new do |opts|
  
  opts.banner = <<EOS
Usage: estool <command> [options..]

Commands include: 
    status                           Returns the status of INDEX
    list                             Returns a list of all indices    
    health                           Returns the health of the shards
    flush                            Performs a full flush of the INDEX
    create                           Create the specified INDEX
    delete                           Delete the specified INDEX. Requires confirmation.
    refresh                          Refresh the specified INDEX
    optimize                         Optimizes the specified INDEX to (-s) number of segments
    snapshot                         Snapshots the specified INDEX to the gateway
    segments                         Returns the segment information. Requires ElasticSearch v
    aliases                          Returns a list of Index/Alias pairs
    ialiases                         Returns a list of Alias/Index pairs
    count                            The number of documents in an index
    mapping
    set_replication
    search
    obj_types

Options include:
EOS

  options.host     = Socket.gethostname
  options.port     = 9200
  options.index    = "_all"
  options.segments = 3
  options.query    = "foo"
  options.raw      = false
  options.usage    = opts

  opts.on('-c', '--host HOSTNAME', 'Connect to ElasticSearch on HOSTNAME', 'Defaults to localhost') do |host|
    options.host     = host
  end

  opts.on('-p', '--port PORT', 'Connect to ElasticSearch using PORT', 'Defaults to 9200') do |port|
    options.port     = port
  end

  opts.on('-i','--index NAME','Name of index to query against', 'Defaults to _all') do |index|
    options.index    = index
  end

  opts.on('-s', '--segments INT', 'Number of segments to optimize to', 'Defaults to 3. Use with <optimize>') do |num|
    options.segments = num
  end

  opts.on('-r','--raw', 'Return raw JSON for parsing by another program') do 
    options.raw      = true
  end

  opts.on('-q', '--query STRING', 'Query INDEX with STRING.', 'Defaults to foo. Use with <search>') do |str|
    options.query    = str
  end

  opts.on('-h', '--help', 'Display this screen and exit'){ puts opts ; exit }
end.parse!

class ESTool
  
  attr_reader :options

  def initialize(options)
    @options = options
  end
  
  def connection() "http://#{options.host}:#{options.port}" ; end

  def shell_response(cmd, req="-XGET")
    url = File.join(connection, cmd)
    Open3.popen3('curl','-s',req, url){ |stdin, stdout, stderr, thread| JSON.parse(stdout.read, :max_nesting => 100) }
  end

  def display cmd 
    result  = self.send(cmd.to_sym)
    display = options.raw ? result.to_json : JSON.pretty_generate(result, :max_nesting => 100)
    puts display
  end
  
  def status()     shell_response(File.join(options.index, "_status?")) ; end

  def list()       status["indices"].keys.sort ; end

  def health()     shell_response("_cluster/health?") ; end

  def flush()      shell_response(File.join(options.index, "_flush?full=true")) ; end

  def create()     shell_response(options.index, "-XPUT") ; end

  def delete()    
    require_confirmation!("delete", options.index)
    shell_response(options.index, "-XDELETE") 
  end
  
  def refresh()    shell_response(File.join(options.index, "_refresh"), "-XPOST") ; end
  
  def optimize()   shell_response(File.join(options.index, "_optimize?max_num_segements=#{options.segments}"), "-XPOST") ; end

  def snapshot()   shell_response(File.join(options.index, "_gateway/snapshot"), "-XPOST") ; end
  
  def segments()   shell_response(File.join(options.index, "_segments")) ; end

  def mapping()    shell_response(File.join(options.index, "_mapping")) ; end
  
  def aliases()    shell_response('_aliases?').sort.inject({}){ |hsh, (index, info)| hsh[index] = info['aliases'].keys ; hsh } ; end
  
  def ialiases() 
    inverse = Hash.new{ |hsh, key| hsh[key] = [] }
    aliases.each{ |idx, als| als.each{ |a| inverse[a] << idx } }
    inverse
  end

  # estool status -r | ruby -rjson -e 'puts JSON.parse($stdin.read)["indices"]["item_count_legacy"]["docs"]'
  # {"num_docs"=>187749, "max_doc"=>187749, "deleted_docs"=>0}
  def count()
    if options.index == '_all'
      status['indices'].inject({}){ |hsh, (index, info)| hsh[index] = info['docs']['num_docs'] ; hsh }
    else
      { options.index => shell_response(File.join(options.index, '_count'))['count'] }
    end
  end

  # curl -s -XPUT http://host:port/index/_settings -d '{"index":{"number_of_replicas":num}}'
  def set_replication() { "error" => "method not yet implemented" }; end

  def search() shell_response(File.join(options.index, "_search?q=#{options.query}")) ; end

  def obj_types() mapping[options.index].keys ; end

  def require_confirmation!(meth, *args)
    print "#{meth.capitalize} method with args #{args} requires confirmation! [yN]?"
    response = STDIN.gets.chomp
    if response =~ /y/i
      print "#{meth.capitalize} method with args #{args} confirmed!"
    else
      print "#{meth.capitalize} method with args #{args} cancelled!"
      exit
    end
  end

  def method_missing meth, *args
    puts "invalid command: #{meth}", options.usage
    exit
  end

end

command = ARGV.first
ESTool.new(options).display(command)