#!/usr/bin/env ruby
#
# kumofs
#
# Copyright (C) 2009 FURUHASHI Sadayuki
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
#

load File.dirname(__FILE__) + "/kumoctl"


class KumoServer < KumoRPC
	def initialize(host, port)
		super(host, port)
	end

	def GetStatus(key)
		send_request_sync_ex(Protocol::GetStatus, [key])
	end

	def SetConfig(key, arg)
		send_request_sync_ex(Protocol::SetConfig, [key, arg])
	end

	STAT_PID         = 0
	STAT_UPTIME      = 1
	STAT_TIME        = 2
	STAT_VERSION     = 3
	STAT_CMD_GET     = 4
	STAT_CMD_SET     = 5
	STAT_CMD_DELETE  = 6
	STAT_DB_ITEMS    = 7
	STAT_CLOCKTIME   = 8
	STAT_RHS         = 9
	STAT_WHS         = 10
	STAT_REPLACE     = 11

	CONF_TCP_NODELAY = 0

	STAT_NAME_OF = {
		STAT_PID        => "pid",
		STAT_UPTIME     => "uptime",
		STAT_TIME       => "time",
		STAT_VERSION    => "version",
		STAT_CMD_GET    => "cmd_get",
		STAT_CMD_SET    => "cmd_set",
		STAT_CMD_DELETE => "cmd_delete",
		STAT_DB_ITEMS   => "curr_items",
	}

	def self.replace_stat_str(flags)
		if flags & 0x01 != 0
			"inactive"
		elsif flags & 0x08 != 0
			"Deleting"
		elsif flags & 0x04 != 0
			"Copying"
		elsif flags & 0x02 != 0
			"WAIT"
		else
			"ready"
		end
	end
end

if $0 == __FILE__


def usage
	puts "Usage: #{File.basename($0)} server-address[:port=#{KumoRPC::SERVER_DEFAULT_PORT}] command"
	puts "       #{File.basename($0)} -m manager-address[:port=#{KumoRPC::MANAGER_DEFAULT_PORT}] command"
	puts "command:"
	puts "   pid                        get pid of server process"
	puts "   uptime                     get uptime"
	puts "   time                       get UNIX time"
	puts "   version                    get version"
	puts "   cmd_get                    get number of get requests"
	puts "   cmd_set                    get number of set requests"
	puts "   cmd_delete                 get number of delete requests"
	puts "   items                      get number of stored items"
	puts "   stats                      get statistics like memcached's 'stats' command"
	puts "   rhs                        get rhs (routing table for Get)"
	puts "   whs                        get whs (routing table for Set/Delete)"
	puts "   hscheck                    check if rhs == whs"
	puts "   replstat                   get replace status"
	puts "   set_delay                  maximize throughput at the expense of latency"
	puts "   unset_delay                minimize latency at the expense of throughput"
	exit 1
end

@commands = {
	"pid"         => [KumoServer::STAT_PID],
	"uptime"      => [KumoServer::STAT_UPTIME],
	"time"        => [KumoServer::STAT_TIME],
	"version"     => [KumoServer::STAT_VERSION],
	"cmd_get"     => [KumoServer::STAT_CMD_GET],
	"cmd_set"     => [KumoServer::STAT_CMD_SET],
	"cmd_delete"  => [KumoServer::STAT_CMD_DELETE],
	"items"       => [KumoServer::STAT_DB_ITEMS],
	"rhs"         => Proc.new{|s| KumoRPC::HSSeed.parse(s.GetStatus(KumoServer::STAT_RHS)).inspect },
	"whs"         => Proc.new{|s| KumoRPC::HSSeed.parse(s.GetStatus(KumoServer::STAT_WHS)).inspect },
	"hscheck"     => Proc.new{|s| s.GetStatus(KumoServer::STAT_RHS) == s.GetStatus(KumoServer::STAT_WHS) },
	"replstat"    => Proc.new{|s| KumoServer.replace_stat_str s.GetStatus(KumoServer::STAT_REPLACE) },
	"set_delay"   => Proc.new{|s| s.SetConfig(KumoServer::CONF_TCP_NODELAY, false) },
	"unset_delay" => Proc.new{|s| s.SetConfig(KumoServer::CONF_TCP_NODELAY, true) },
	"stats"       => [
		KumoServer::STAT_PID,
		KumoServer::STAT_UPTIME,
		KumoServer::STAT_TIME,
		KumoServer::STAT_CMD_GET,
		KumoServer::STAT_CMD_SET,
		KumoServer::STAT_CMD_DELETE,
		KumoServer::STAT_DB_ITEMS,
	],
}


def run_command(cmd, host, port)
	s = KumoServer.new(host, port)
	begin
		result = case cmd
						 when Array
							if cmd.length > 1
								res = ""
								cmd.each {|c|
									res += sprintf "STAT %s %s\n", KumoServer::STAT_NAME_OF[c], s.GetStatus(*[c])
								}
								res+"END"
							else
								s.GetStatus(*cmd)
							end

				 else Proc
					 cmd.call(s)

				 end

	rescue
		puts $!.to_s
	ensure
		s.close
	end
	result.to_s+"\n"
end

if ARGV[0] == "-m"
	manager = true
	ARGV.shift
end

if ARGV.length != 2
	usage
end

addr = ARGV.shift
cmd  = ARGV.shift
cmd = @commands[cmd]

unless cmd
	usage
end


if manager
	host, port = addr.split(':', 2)
	port ||= KumoRPC::MANAGER_DEFAULT_PORT

	mgr = KumoManager.new(host, port)
	attached, not_attached, date, clock = mgr.GetStatus
	mgr.close

	attached.each {|host, port, active|
		if active
			print "#{host}:#{port}:  "
			$stdout.flush
			print run_command(cmd, host, port)
		end
	}

else
	host, port = addr.split(':', 2)
	port ||= KumoRPC::SERVER_DEFAULT_PORT

	print run_command(cmd, host, port)

end


end   # if $0 == __FILE__
