File

local config = require "core.configmanager";
local log = require "util.logger".init("stats");
local timer = require "util.timer";
local fire_event = prosody.events.fire_event;
local array = require "util.array";
local timed = require "util.openmetrics".timed;

local stats_interval_config = config.get("*", "statistics_interval");
local stats_interval = tonumber(stats_interval_config);
if stats_interval_config and not stats_interval and stats_interval_config ~= "manual" then
	log("error", "Invalid 'statistics_interval' setting, statistics will be disabled");
end

local stats_provider_name;
local stats_provider_config = config.get("*", "statistics");
local stats_provider = stats_provider_config;

if not stats_provider and stats_interval then
	stats_provider = "internal";
elseif stats_provider and not stats_interval then
	stats_interval = 60;
end
if stats_interval_config == "manual" then
	stats_interval = nil;
end

local builtin_providers = {
	internal = "util.statistics";
	statsd = "util.statsd";
};


local stats, stats_err = false, nil;

if stats_provider then
	if stats_provider:sub(1,1) == ":" then
		stats_provider = stats_provider:sub(2);
		stats_provider_name = "external "..stats_provider;
	elseif stats_provider then
		stats_provider_name = "built-in "..stats_provider;
		stats_provider = builtin_providers[stats_provider];
		if not stats_provider then
			log("error", "Unrecognized statistics provider '%s', statistics will be disabled", stats_provider_config);
		end
	end

	local have_stats_provider, stats_lib = pcall(require, stats_provider);
	if not have_stats_provider then
		stats, stats_err = nil, stats_lib;
	else
		local stats_config = config.get("*", "statistics_config");
		stats, stats_err = stats_lib.new(stats_config);
		stats_provider_name = stats_lib._NAME or stats_provider_name;
	end
end

if stats == nil then
	log("error", "Error loading statistics provider '%s': %s", stats_provider, stats_err);
end

local measure, collect, metric, cork, uncork;

if stats then
	function metric(type_, name, unit, description, labels, extra)
		local registry = stats.metric_registry
		local f = assert(registry[type_], "unknown metric family type: "..type_);
		return f(registry, name, unit or "", description or "", labels, extra);
	end

	local function new_legacy_metric(stat_type, name, unit, description, fixed_label_key, fixed_label_value, extra)
		local label_keys = array()
		local conf = extra or {}
		if fixed_label_key then
			label_keys:push(fixed_label_key)
		end
		unit = unit or ""
		local mf = metric(stat_type, "prosody_" .. name, unit, description, label_keys, conf);
		if fixed_label_key then
			mf = mf:with_partial_label(fixed_label_value)
		end
		return mf:with_labels()
	end

	local function unwrap_legacy_extra(extra, type_, name, unit)
		local description = extra and extra.description or "Legacy "..type_.." metric "..name
		unit = extra and extra.unit or unit
		return description, unit
	end

	-- These wrappers provide the pre-OpenMetrics interface of statsmanager
	-- and moduleapi (module:measure).
	local legacy_metric_wrappers = {
		amount = function(name, fixed_label_key, fixed_label_value, extra)
			local initial = 0
			if type(extra) == "number" then
				initial = extra
			else
				initial = extra and extra.initial or initial
			end
			local description, unit = unwrap_legacy_extra(extra, "amount", name)

			local m = new_legacy_metric("gauge", name, unit, description, fixed_label_key, fixed_label_value)
			m:set(initial or 0)
			return function(v)
				m:set(v)
			end
		end;

		counter = function(name, fixed_label_key, fixed_label_value, extra)
			if type(extra) == "number" then
				-- previous versions of the API allowed passing an initial
				-- value here; we do not allow that anymore, it is not a thing
				-- which makes sense with counters
				extra = nil
			end

			local description, unit = unwrap_legacy_extra(extra, "counter", name)

			local m = new_legacy_metric("counter", name, unit, description, fixed_label_key, fixed_label_value)
			m:set(0)
			return function(v)
				m:add(v)
			end
		end;

		rate = function(name, fixed_label_key, fixed_label_value, extra)
			if type(extra) == "number" then
				-- previous versions of the API allowed passing an initial
				-- value here; we do not allow that anymore, it is not a thing
				-- which makes sense with counters
				extra = nil
			end

			local description, unit = unwrap_legacy_extra(extra, "counter", name)

			local m = new_legacy_metric("counter", name, unit, description, fixed_label_key, fixed_label_value)
			m:set(0)
			return function()
				m:add(1)
			end
		end;

		times = function(name, fixed_label_key, fixed_label_value, extra)
			local conf = {}
			if extra and extra.buckets then
				conf.buckets = extra.buckets
			else
				conf.buckets = { 0.001, 0.01, 0.1, 1.0, 10.0, 100.0 }
			end
			local description, _ = unwrap_legacy_extra(extra, "times", name)

			local m = new_legacy_metric("histogram", name, "seconds", description, fixed_label_key, fixed_label_value, conf)
			return function()
				return timed(m)
			end
		end;

		sizes = function(name, fixed_label_key, fixed_label_value, extra)
			local conf = {}
			if extra and extra.buckets then
				conf.buckets = extra.buckets
			else
				conf.buckets = { 1024, 4096, 32768, 131072, 1048576, 4194304, 33554432, 134217728, 1073741824 }
			end
			local description, _ = unwrap_legacy_extra(extra, "sizes", name)

			local m = new_legacy_metric("histogram", name, "bytes", description, fixed_label_key, fixed_label_value, conf)
			return function(v)
				m:sample(v)
			end
		end;

		distribution = function(name, fixed_label_key, fixed_label_value, extra)
			if type(extra) == "string" then
				-- compat with previous API
				extra = { unit = extra }
			end
			local description, unit = unwrap_legacy_extra(extra, "distribution", name, "")
			local m = new_legacy_metric("summary", name, unit, description, fixed_label_key, fixed_label_value)
			return function(v)
				m:sample(v)
			end
		end;
	};

	-- Argument order switched here to support the legacy statsmanager.measure
	-- interface.
	function measure(stat_type, name, extra, fixed_label_key, fixed_label_value)
		local wrapper = assert(legacy_metric_wrappers[stat_type], "unknown legacy metric type "..stat_type)
		return wrapper(name, fixed_label_key, fixed_label_value, extra)
	end

	if stats.cork then
		function cork()
			return stats:cork()
		end

		function uncork()
			return stats:uncork()
		end
	else
		function cork() end
		function uncork() end
	end

	if stats_interval or stats_interval_config == "manual" then

		local mark_collection_start = measure("times", "stats.collection");
		local mark_processing_start = measure("times", "stats.processing");

		function collect()
			local mark_collection_done = mark_collection_start();
			fire_event("stats-update");
			-- ensure that the backend is uncorked, in case it got stuck at
			-- some point, to avoid infinite resource use
			uncork()
			mark_collection_done();
			local manual_result = nil

			if stats.metric_registry then
				-- only if supported by the backend, we fire the event which
				-- provides the current metric values
				local mark_processing_done = mark_processing_start();
				local metric_registry = stats.metric_registry;
				fire_event("openmetrics-updated", { metric_registry = metric_registry })
				mark_processing_done();
				manual_result = metric_registry;
			end

			return stats_interval, manual_result;
		end
		if stats_interval then
			log("debug", "Statistics enabled using %s provider, collecting every %d seconds", stats_provider_name, stats_interval);
			timer.add_task(stats_interval, collect);
			prosody.events.add_handler("server-started", function () collect() end, -1);
			prosody.events.add_handler("server-stopped", function () collect() end, -1);
		else
			log("debug", "Statistics enabled using %s provider, no scheduled collection", stats_provider_name);
		end
	else
		log("debug", "Statistics enabled using %s provider, collection is disabled", stats_provider_name);
	end
else
	log("debug", "Statistics disabled");
	function measure() return measure; end

	local dummy_mt = {}
	function dummy_mt.__newindex()
	end
	function dummy_mt:__index()
		return self
	end
	function dummy_mt:__call()
		return self
	end
	local dummy = {}
	setmetatable(dummy, dummy_mt)

	function metric() return dummy; end
	function cork() end
	function uncork() end
end

local exported_collect = nil;
if stats_interval_config == "manual" then
	exported_collect = collect;
end

return {
	collect = exported_collect;
	measure = measure;
	cork = cork;
	uncork = uncork;
	metric = metric;
	get_metric_registry = function ()
		return stats and stats.metric_registry or nil
	end;
};