#!/usr/bin/env lua

local print = print;

local pairs = pairs;

local socket = require"socket";

local stats = {};

stats.c2s = {

	graph_title = "Prosody C2S Connections";

	graph_vlabel = "users";

	graph_category = "Prosody";

	all_client_connections = {

		label = "client connections";

		_key = "total_c2s";

	}

}

stats.s2s = {

	graph_title = "Prosody S2S Connections";

	graph_vlabel = "servers";

	graph_category = "Prosody";

	outgoing_connections = {

		label = "outgoing connections";

		_key = "total_s2sout";

	};

	incoming_connections = {

		label = "incoming connections";

		_key = "total_s2sin";

	}

}

stats.mem = {

	graph_title = "Prosody Memory Usage";

	graph_vlabel = "Bytes";

	graph_args = "--base 1024 -l 0";

	graph_category = "Prosody"; --memory_unused

	graph_order = "memory_total memory_rss memory_allocated memory_used memory_lua memory_returnable";

	memory_allocated = { label = "Allocated", draw = "AREA"  };

	memory_lua = { label = "Lua", draw = "AREA" };

	memory_rss = { label = "RSS", draw = "AREA" };

	memory_total = { label = "Total", draw = "AREA" };

	-- memory_unused = { label = "Unused", draw = "AREA" };

	memory_used = { label = "Used", draw = "AREA" };

	memory_returnable = { label = "Returnable", draw = "AREA" };

}

stats.cpu = {

	graph_title = "Prosody CPU Usage";

	graph_category = "Prosody";

	graph_args = "-l 0";

	graph_vlabel = "CPU time used in milliseconds";

	cpu_total = { label = "CPU"; type = "DERIVE"; min = 0; };

}

stats.auth = {

	graph_title = "Prosody Authentications";

	graph_category = "Prosody";

	graph_args = "--base 1000";

	c2s_auth = {

		label = "Logins";

		type = "DERIVE";

		min = 0;

	};

	c2s_authfail = {

		label = "Failed logins";

		type = "DERIVE";

		min = 0;

	};

}

local function onerror(msg, err, exit)

	io.stderr:write(msg, '\n');

	if err then

		io.stderr:write(err, '\n');

	end

	os.exit(exit or 1);

end

local function connect()

	local conn, err = socket.connect(os.getenv"host" or "localhost", os.getenv"port" or 5782);

	if not conn then onerror("Could not connect to prosody", err); end

	conn:settimeout(1);

	return conn;

end

local function get_config(item)

	for k,v in pairs(item) do

		if type(v) == "string" then

			print(k .. " " .. v);

		elseif type(v) == "table" then

			for sk,v in pairs(v) do

				if not sk:match("^_") then

					print(k.."."..sk.." "..v);

				end

			end

		end

	end

end

local function get_stats(item)

	local labels = {};

	for key, val in pairs(item) do

		if type(val) == "table" and val.label then

			labels[val._key or key] = key;

		end

	end

	local conn = connect();

	local line, err = conn:receive("*l");

	local stat, value, label;

	while line and line ~= "" and next(labels) ~= nil do

		stat, value = line:match('^STAT%s+"([^"]*)"%s*(%b())');

		label = stat and labels[stat];

		if label then

			print(label..".value "..tonumber(value:sub(2,-2)));

			labels[stat] = nil;

		end

		line, err = conn:receive("*l");

	end

	if err then onerror(err); end

end

local function main(stat, mode)

	if mode == "suggest" then

		for available_stat in pairs(stats) do

			print(available_stat);

		end

	elseif mode == "config" then

		return get_config(stats[stat]);

	elseif stats[stat] then

		return get_stats(stats[stat]);

	end

end

if arg then return main(arg[0]:match("prosody_(%w*)"), ...); end

return {

	stats = stats,

	get_stats = get_stats,

	get_config = get_config,

}