--[[

This module implements a subset of the OpenMetrics Internet Draft version 00.

URL: https://datatracker.ietf.org/doc/html/draft-richih-opsawg-openmetrics-00

The following metric types are supported:

- Counter

- Gauge

- Histogram

- Summary

It is used by util.statsd and util.statistics to provide the OpenMetrics API.

To understand what this module is about, it is useful to familiarize oneself

with the terms MetricFamily, Metric, LabelSet, Label and MetricPoint as

defined in the I-D linked above.

--]]

-- metric constructor interface:

-- metric_ctor(..., family_name, labels, extra)

local time = require "prosody.util.time".now;

local select = select;

local array = require "prosody.util.array";

local log = require "prosody.util.logger".init("util.openmetrics");

local new_multitable = require "prosody.util.multitable".new;

local iter_multitable = require "prosody.util.multitable".iter;

local t_concat, t_insert = table.concat, table.insert;

local t_pack, t_unpack = table.pack, table.unpack;

-- BEGIN of Utility: "metric proxy"

-- This allows to wrap a MetricFamily in a proxy which only provides the

-- `with_labels` and `with_partial_label` methods. This allows to pre-set one

-- or more labels on a metric family. This is used in particular via

-- `with_partial_label` by the moduleapi in order to pre-set the `host` label

-- on metrics created in non-global modules.

local metric_proxy_mt = {}

metric_proxy_mt.__name = "metric_proxy"

metric_proxy_mt.__index = metric_proxy_mt

local function new_metric_proxy(metric_family, with_labels_proxy_fun)

	return setmetatable({

		_family = metric_family,

		with_labels = function(self, ...)

			return with_labels_proxy_fun(self._family, ...)

		end;

		with_partial_label = function(self, label)

			return new_metric_proxy(self._family, function(family, ...)

				return family:with_labels(label, ...)

			end)

		end

	}, metric_proxy_mt);

end

-- END of Utility: "metric proxy"

-- BEGIN Rendering helper functions (internal)

local function escape(text)

	return text:gsub("\\", "\\\\"):gsub("\"", "\\\""):gsub("\n", "\\n");

end

local function escape_name(name)

	return name:gsub("/", "__"):gsub("[^A-Za-z0-9_]", "_"):gsub("^[^A-Za-z_]", "_%1");

end

local function repr_help(metric, docstring)

	docstring = docstring:gsub("\\", "\\\\"):gsub("\n", "\\n");

	return "# HELP "..escape_name(metric).." "..docstring.."\n";

end

local function repr_unit(metric, unit)

	if not unit then

		unit = ""

	else

		unit = unit:gsub("\\", "\\\\"):gsub("\n", "\\n");

	end

	return "# UNIT "..escape_name(metric).." "..unit.."\n";

end

-- local allowed_types = { counter = true, gauge = true, histogram = true, summary = true, untyped = true };

-- local allowed_types = { "counter", "gauge", "histogram", "summary", "untyped" };

local function repr_type(metric, type_)

	-- if not allowed_types:contains(type_) then

	-- 	return;

	-- end

	return "# TYPE "..escape_name(metric).." "..type_.."\n";

end

local function repr_label(key, value)

	return key.."=\""..escape(value).."\"";

end

local function repr_labels(labelkeys, labelvalues, extra_labels)

	local values = {}

	if labelkeys then

		for i, key in ipairs(labelkeys) do

			local value = labelvalues[i]

			t_insert(values, repr_label(escape_name(key), escape(value)));

		end

	end

	if extra_labels then

		for key, value in pairs(extra_labels) do

			t_insert(values, repr_label(escape_name(key), escape(value)));

		end

	end

	if #values == 0 then

		return "";

	end

	return "{"..t_concat(values, ",").."}";

end

local function repr_sample(metric, labelkeys, labelvalues, extra_labels, value)

	return escape_name(metric)..repr_labels(labelkeys, labelvalues, extra_labels).." "..string.format("%.17g", value).."\n";

end

-- END Rendering helper functions (internal)

local function render_histogram_le(v)

	if v == 1/0 then

		-- I-D-00: 4.1.2.2.1:

		--    Exposers MUST produce output for positive infinity as +Inf.

		return "+Inf"

	end

	return string.format("%.14g", v)

end

-- BEGIN of generic MetricFamily implementation

local metric_family_mt = {}

metric_family_mt.__name = "metric_family"

metric_family_mt.__index = metric_family_mt

local function histogram_metric_ctor(orig_ctor, buckets)

	return function(family_name, labels, extra)

		return orig_ctor(buckets, family_name, labels, extra)

	end

end

local function new_metric_family(backend, type_, family_name, unit, description, label_keys, extra)

	local metric_ctor = assert(backend[type_], "statistics backend does not support "..type_.." metrics families")

	local labels = label_keys or {}

	local user_labels = #labels

	if type_ == "histogram" then

		local buckets = extra and extra.buckets

		if not buckets then

			error("no buckets given for histogram metric")

		end

		buckets = array(buckets)

		buckets:push(1/0)  -- must have +inf bucket

		metric_ctor = histogram_metric_ctor(metric_ctor, buckets)

	end

	local data

	if #labels == 0 then

		data = metric_ctor(family_name, nil, extra)

	else

		data = new_multitable()

	end

	local mf = {

		family_name = family_name,

		data = data,

		type_ = type_,

		unit = unit,

		description = description,

		user_labels = user_labels,

		label_keys = labels,

		extra = extra,

		_metric_ctor = metric_ctor,

	}

	setmetatable(mf, metric_family_mt);

	return mf

end

function metric_family_mt:new_metric(labels)

	return self._metric_ctor(self.family_name, labels, self.extra)

end

function metric_family_mt:clear()

	for _, metric in self:iter_metrics() do

		metric:reset()

	end

end

function metric_family_mt:with_labels(...)

	local count = select('#', ...)

	if count ~= self.user_labels then

		error("number of labels passed to with_labels does not match number of label keys")

	end

	if count == 0 then

		return self.data

	end

	local metric = self.data:get(...)

	if not metric then

		local values = t_pack(...)

		metric = self:new_metric(values)

		values[values.n+1] = metric

		self.data:set(t_unpack(values, 1, values.n+1))

	end

	return metric

end

function metric_family_mt:with_partial_label(label)

	return new_metric_proxy(self, function (family, ...)

		return family:with_labels(label, ...)

	end)

end

function metric_family_mt:iter_metrics()

	if #self.label_keys == 0 then

		local done = false

		return function()

			if done then

				return nil

			end

			done = true

			return {}, self.data

		end

	end

	local searchkeys = {};

	local nlabels = #self.label_keys

	for i=1,nlabels do

		searchkeys[i] = nil;

	end

	local it, state = iter_multitable(self.data, t_unpack(searchkeys, 1, nlabels))

	return function(_s)

		local label_values = t_pack(it(_s))

		if label_values.n == 0 then

			return nil, nil

		end

		local metric = label_values[label_values.n]

		label_values[label_values.n] = nil

		label_values.n = label_values.n - 1

		return label_values, metric

	end, state

end

-- END of generic MetricFamily implementation

-- BEGIN of MetricRegistry implementation

-- Helper to test whether two metrics are "equal".

local function equal_metric_family(mf1, mf2)

	if mf1.type_ ~= mf2.type_ then

		return false

	end

	if #mf1.label_keys ~= #mf2.label_keys then

		return false

	end

	-- Ignoring unit here because in general it'll be part of the name anyway

	-- So either the unit was moved into/out of the name (which is a valid)

	-- thing to do on an upgrade or we would expect not to see any conflicts

	-- anyway.

	--[[

	if mf1.unit ~= mf2.unit then

		return false

	end

	]]

	for i, key in ipairs(mf1.label_keys) do

		if key ~= mf2.label_keys[i] then

			return false

		end

	end

	return true

end

-- If the unit is not empty, add it to the full name as per the I-D spec.

local function compose_name(name, unit)

	local full_name = name

	if unit and unit ~= "" then

		full_name = full_name .. "_" .. unit

	end

	-- TODO: prohibit certain suffixes used by metrics if where they may cause

	-- conflicts

	return full_name

end

local metric_registry_mt = {}

metric_registry_mt.__name = "metric_registry"

metric_registry_mt.__index = metric_registry_mt

local function new_metric_registry(backend)

	local reg = {

		families = {},

		backend = backend,

	}

	setmetatable(reg, metric_registry_mt)

	return reg

end

function metric_registry_mt:register_metric_family(name, metric_family)

	local existing = self.families[name];

	if existing then

		if not equal_metric_family(metric_family, existing) then

			-- We could either be strict about this, or replace the

			-- existing metric family with the new one.

			-- Being strict is nice to avoid programming errors /

			-- conflicts, but causes issues when a new version of a module

			-- is loaded.

			--

			-- We will thus assume that the new metric is the correct one;

			-- That is probably OK because unless you're reaching down into

			-- the util.openmetrics or core.statsmanager API, your metric

			-- name is going to be scoped to `prosody_mod_$modulename`

			-- anyway and the damage is thus controlled.

			--

			-- To make debugging such issues easier, we still log.

			log("debug", "replacing incompatible existing metric family %s", name)

			-- Below is the code to be strict.

			--error("conflicting declarations for metric family "..name)

		else

			return existing

		end

	end

	self.families[name] = metric_family

	return metric_family

end

function metric_registry_mt:gauge(name, unit, description, labels, extra)

	name = compose_name(name, unit)

	local mf = new_metric_family(self.backend, "gauge", name, unit, description, labels, extra)

	mf = self:register_metric_family(name, mf)

	return mf

end

function metric_registry_mt:counter(name, unit, description, labels, extra)

	name = compose_name(name, unit)

	local mf = new_metric_family(self.backend, "counter", name, unit, description, labels, extra)

	mf = self:register_metric_family(name, mf)

	return mf

end

function metric_registry_mt:histogram(name, unit, description, labels, extra)

	name = compose_name(name, unit)

	local mf = new_metric_family(self.backend, "histogram", name, unit, description, labels, extra)

	mf = self:register_metric_family(name, mf)

	return mf

end

function metric_registry_mt:summary(name, unit, description, labels, extra)

	name = compose_name(name, unit)

	local mf = new_metric_family(self.backend, "summary", name, unit, description, labels, extra)

	mf = self:register_metric_family(name, mf)

	return mf

end

function metric_registry_mt:get_metric_families()

	return self.families

end

function metric_registry_mt:render()

	local answer = {};

	for metric_family_name, metric_family in pairs(self:get_metric_families()) do

		t_insert(answer, repr_help(metric_family_name, metric_family.description))

		t_insert(answer, repr_unit(metric_family_name, metric_family.unit))

		t_insert(answer, repr_type(metric_family_name, metric_family.type_))

		for labelset, metric in metric_family:iter_metrics() do

			for suffix, extra_labels, value in metric:iter_samples() do

				t_insert(answer, repr_sample(metric_family_name..suffix, metric_family.label_keys, labelset, extra_labels, value))

			end

		end

	end

	t_insert(answer, "# EOF\n")

	return t_concat(answer, "");

end

-- END of MetricRegistry implementation

-- BEGIN of general helpers for implementing high-level APIs on top of OpenMetrics

local function timed(metric)

	local t0 = time()

	local submitter = assert(metric.sample or metric.set, "metric type cannot be used with timed()")

	return function()

		local t1 = time()

		submitter(metric, t1-t0)

	end

end

-- END of general helpers

return {