-- Log statistics to Prometheus

--

-- Copyright (C) 2014 Daurnimator

-- Copyright (C) 2018 Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>

-- Copyright (C) 2021 Jonas Schäfer <jonas@zombofant.net>

--

-- This module is MIT/X11 licensed.

module:set_global();

local tostring = tostring;

local t_insert = table.insert;

local t_concat = table.concat;

local socket = require "socket";

local statsman = require "core.statsmanager";

local get_stats = statsman.get_stats;

local get_metric_registry = statsman.get_metric_registry;

local collect = statsman.collect;

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 get_timestamp()

	-- Using LuaSocket for that because os.time() only has second precision.

	return math.floor(socket.gettime() * 1000);

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

local get_metrics;

if statsman.get_metric_registry then

	module:log("debug", "detected OpenMetrics statsmanager")

	-- Prosody 0.12+ with OpenMetrics

	function get_metrics(event)

		local response = event.response;

		response.headers.content_type = "application/openmetrics-text; version=0.0.4";

		if collect then

			-- Ensure to get up-to-date samples when running in manual mode

			collect()

		end

		local registry = get_metric_registry()

		if registry == nil then

			response.headers.content_type = "text/plain; charset=utf-8"

			response.status_code = 404

			return "No statistics provider configured\n"

		end

		local answer = {};

		for metric_family_name, metric_family in pairs(registry: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

else

	module:log("debug", "detected pre-OpenMetrics statsmanager")

	-- Pre-OpenMetrics

	local allowed_extras = { min = true, max = true, average = true };

	local function insert_extras(data, key, name, timestamp, extra)

		if not extra then

			return false;

		end

		local has_extra = false;

		for extra_name in pairs(allowed_extras) do

			if extra[extra_name] then

				local field = {

					value = extra[extra_name],

					labels = {

						["type"] = name,

						field = extra_name,

					},

					typ = "gauge";

					timestamp = timestamp,

				};

				t_insert(data[key], field);

				has_extra = true;

			end

		end

		return has_extra;

	end

	local function parse_stats()

		local timestamp = tostring(get_timestamp());

		local data = {};

		local stats, changed_only, extras = get_stats();

		for stat, value in pairs(stats) do

			-- module:log("debug", "changed_stats[%q] = %s", stat, tostring(value));

			local extra = extras[stat];

			local host, sect, name, typ = stat:match("^/([^/]+)/([^/]+)/(.+):(%a+)$");

			if host == nil then

				sect, name, typ = stat:match("^([^.]+)%.(.+):(%a+)$");

			elseif host == "*" then

				host = nil;

			end

			if sect:find("^mod_measure_.") then

				sect = sect:sub(13);

			elseif sect:find("^mod_statistics_.") then

				sect = sect:sub(16);

			end

			local key = escape_name("prosody_"..sect);

			local field = {

				value = value,

				labels = { ["type"] = name},

				-- TODO: Use the other types where it makes sense.

				typ = (typ == "rate" and "counter" or "gauge"),

				timestamp = timestamp,

			};

			if host then

				field.labels.host = host;

			end

			if data[key] == nil then

				data[key] = {};

			end

			if not insert_extras(data, key, name, timestamp, extra) then

				t_insert(data[key], field);

			end

		end

		return data;

	end

	function get_metrics(event)

		local response = event.response;

		response.headers.content_type = "text/plain; version=0.0.4";

		if statsman.collect then

			statsman.collect()

		end

		local answer = {};

		for key, fields in pairs(parse_stats()) do

			t_insert(answer, repr_help(key, ""));

			t_insert(answer, repr_type(key, fields[1].typ));

			for _, field in pairs(fields) do

				t_insert(answer, repr_sample(key, nil, nil, field.labels, field.value, field.timestamp));

			end

		end

		return t_concat(answer, "");

	end

end

function module.add_host(module)

	module:depends "http";

	module:provides("http", {

		default_path = "metrics";

		route = {

			GET = get_metrics;

		};

	});

end