#!/usr/bin/env ruby

# tail like CLI for mongo capped collection

require 'optparse'
require 'json'

require 'mongo'

TailConfig = {
  d: 'fluent',
  c: 'out_mongo_backup',
  h: 'localhost',
  p: 27017,
  n: 10
}

OptionParser.new { |opt|
  opt.on('-d VAL', 'The name of database') { |v| TailConfig[:d] = v }
  opt.on('-c VAL', 'The name of collection') { |v| TailConfig[:c] = v }
  opt.on('-h VAL', 'The host of mongodb server') { |v| TailConfig[:h] = v }
  opt.on('-p VAL', 'The port of mongodb server') { |v| TailConfig[:p] = Integer(v) }
  opt.on('-n VAL', 'The number of documents') { |v| TailConfig[:n] = Integer(v) }
  opt.on('-f', 'This option causes tail to not stop when end of collection is reached, but rather to wait for additional data to be appended to the collection') { |v|
    TailConfig[:f] = true
  }

  opt.parse!(ARGV)
}

def get_client_options(conf)
  client_options = {}
  client_options[:database] = conf[:d]
  client_options
end

def get_collection_options
  collection_options = {}
  collection_options[:capped] = true
  collection_options
end

def get_capped_collection(conf)
  client_options = get_client_options(conf)
  collection_options = get_collection_options
  client = Mongo::Client.new(["#{conf[:h]}:#{conf[:p]}"], client_options)
  collection = client["#{conf[:c]}", collection_options]
  if collection.capped?
    collection
  else
    puts "#{conf[:c]} is not capped. mongo-tail can not tail normal collection."
  end
end

def create_skip_number(collection, conf)
  skip = collection.count - conf[:n]
  skip
end

def tail_n(collection, conf)
  collection.find().skip(create_skip_number(collection, conf)).each {|doc|
    puts doc.to_json
  }
end

def tail_forever(collection, conf)
  loop {
    option['_id'] = {'$gt' => BSON::ObjectId(@last_id)} if @last_id

    documents = collection.find(option)
    if documents.count >= 1
      documents.each {|doc|
        STDOUT.puts doc.to_json
        STDOUT.flush
        if id = doc.delete('_id')
          @last_id = id.to_s
        end
      }
    else
      sleep 1
    end
  }
end

exit unless collection = get_capped_collection(TailConfig)

if TailConfig[:f]
  tail_forever(collection, TailConfig)
else
  tail_n(collection, TailConfig)
end
