#!/usr/bin/env ruby
require 'rubygems'
require 'trollop'
require 'sup'
require 'sup/client'
require 'pp'
require 'yaml'
include Redwood

SUB_COMMANDS = %w(query count label add)
global_opts = Trollop::options do
  #version = "sup-cmd (sup #{Redwood::VERSION})"
  banner <<EOS
Connect to a running sup-server.

Usage:
  sup-cmd [global options] command [options]

  Valid commands: #{SUB_COMMANDS * ', '}

  Global options:
EOS

  opt :host, "server address", :type => :string, :default => 'localhost', :short => 'o'
  opt :port, "server port", :type => :int, :default => 4300
  opt :socket, "unix domain socket path", :type => :string, :default => nil
  opt :verbose

  conflicts :host, :socket
  conflicts :port, :socket

  stop_on SUB_COMMANDS
end

cmd = ARGV.shift
cmd_opts = case cmd
when "query"
  Trollop.options do
    opt :offset, "Offset", :default => 0, :type => :int
    opt :limit, "Limit", :type => :int
    opt :raw, "Retrieve raw message text", :default => false
  end
when "count"
  Trollop.options do
  end
when "label"
  Trollop.options do
    opt :add_labels, "Labels to add", :default => ""
    opt :remove_labels, "Labels to remove", :default => ""
  end
when "add"
  Trollop.options do
    opt :labels, "Labels separated by commas", :default => ""
    opt :mbox, "Treat input files as mboxes", :default => false
  end
else
  Trollop::die "unrecognized command #{cmd.inspect}"
end

class SupCmd < Redwood::Client
  def initialize cmd, args, opts
    @cmd = cmd
    @opts = opts
    @args = args
    super()
  end

  def get_query
    @args.first or fail "query argument required"
  end

  def connection_established
    case @cmd
    when "query"
      query get_query, @opts[:offset], @opts[:limit], @opts[:raw] do |result|
        if result
          puts YAML.dump(result['summary'])
          puts YAML.dump(result['raw']) if @opts[:raw]
        else
          close_connection
        end
      end
    when "count"
      count(get_query) do |x|
        puts x
        close_connection
      end
    when "label"
      label get_query, @opts[:remove_labels].split(','), @opts[:add_labels].split(',') do
        close_connection
      end
    when "add"
      ARGF.binmode
      labels = @opts[:labels].split(',')
      get_message = lambda do
        return ARGF.gets(nil) unless @opts[:mbox]
        str = ""
        l = ARGF.gets
        str << l until ARGF.closed? || ARGF.eof? || MBox::is_break_line?(l = ARGF.gets)
        str.empty? ? nil : str
      end
      i_s = i = 0
      t = Time.now
      while raw = get_message[]
        i += 1
        t_d = Time.now - t
        if t_d >= 5
          i_d = i - i_s
          puts "indexed #{i} messages (#{i_d/t_d} m/s)" if global_opts[:verbose]
          t = Time.now
          i_s = i
        end
        add raw, labels do
          close_connection
        end
      end
    else
      fail "#{@cmd} command unimplemented"
      close_connection
    end
  end

  def unbind
    EM.stop
  end
end


EM.run do
  if global_opts[:socket]
    EM.connect global_opts[:socket], SupCmd, cmd, ARGV, cmd_opts.merge(global_opts)
  else
    EM.connect global_opts[:host], global_opts[:port], SupCmd, cmd, ARGV, cmd_opts.merge(global_opts)
  end
end

exit 0

