module StatsD_Checker {

/* Verifies that  StatsD metrics in a test match the expected values
 * Uses StatsD_CodecPort to receive the statsd messages from the DUT
 * and a separate VTY connection to reset and trigger the stats.
 *
 * When using this you should configure your stats reporter to disable
 * interval-based reports and always send all metrics:
 * > stats interval 0
 * > stats reporter statsd
 * >  remote-ip a.b.c.d
 * >  remote-port 8125
 * >  level subscriber
 * >  flush-period 1
 * >  mtu 1024
 * >  enable
 *
 * (C) 2020 by sysmocom s.f.m.c. GmbH <info@sysmocom.de>
 * All rights reserved.
 *
 * Author: Daniel Willmann <dwillmann@sysmocom.de>
 *
 * Released under the terms of GNU General Public License, Version 2 or
 * (at your option) any later version.
 * SPDX-License-Identifier: GPL-2.0-or-later
 */

import from Misc_Helpers all;
import from Socket_API_Definitions all;

import from StatsD_Types all;
import from StatsD_CodecPort all;
import from StatsD_CodecPort_CtrlFunct all;

import from General_Types all;
import from Osmocom_Types all;
#ifdef STATSD_HAVE_VTY
import from Osmocom_VTY_Functions all;
import from TELNETasp_PortType all;
#endif

modulepar {
	/* Whether to test stats values */
	boolean mp_enable_stats := true;
}

type record StatsDMetricKey {
	MetricName	name,
	MetricType	mtype
};
type set of StatsDMetricKey StatsDMetricKeys;

template (value) StatsDMetricKey ts_StatsDMetricKey(template (value) MetricName	name, template (value) MetricType mtype) := {
	name := name,
	mtype := mtype
}

type record StatsDExpect {
	MetricName	name,
	MetricType	mtype,
	MetricValue	min,
	MetricValue	max
};
type set of StatsDExpect StatsDExpects;

type enumerated StatsDResultType {
	e_Matched,
	e_Mismatched,
	e_NotFound
}

type record StatsDExpectResult {
	StatsDResultType kind,
	integer idx
}

type component StatsD_Checker_CT {
#ifdef STATSD_HAVE_VTY
	port TELNETasp_PT STATSVTY;
#endif
	port STATSD_PROC_PT STATSD_PROC;
	port STATSD_CODEC_PT STATS;
	timer T_statsd := 5.0;
}

type component StatsD_ConnHdlr {
	port STATSD_PROC_PT STATSD_PROC;
}

signature STATSD_reset();
signature STATSD_snapshot(in StatsDMetricKeys keys, in boolean since_last_snapshot) return StatsDMetrics;
signature STATSD_expect(in StatsDExpects expects, in boolean wait_converge, in boolean use_snapshot, in StatsDMetrics snapshot) return boolean;

type port STATSD_PROC_PT procedure {
	inout STATSD_reset, STATSD_snapshot, STATSD_expect;
} with {extension "internal"};

/* Expect templates and functions */


/* StatsD checker component */
function main(charstring statsd_host, integer statsd_port) runs on StatsD_Checker_CT {
	var StatsD_ConnHdlr vc_conn;
	var StatsDMetricKeys keys;
	var boolean since_last_snapshot;
	var StatsDExpects expects;
	var boolean wait_converge;
	var boolean use_snapshot;
	var StatsDMetrics snapshot;
	var Result res;

	while (not mp_enable_stats) {
		log("StatsD checker disabled by modulepar");
		f_sleep(3600.0);
	}

	map(self:STATS, system:STATS);
	res := StatsD_CodecPort_CtrlFunct.f_IPL4_listen(STATS, statsd_host, statsd_port, { udp := {} }, {});
	if (not ispresent(res.connId)) {
		Misc_Helpers.f_shutdown(__BFILE__, __LINE__, fail,
					"Could not bind StatsD socket, check your configuration");
	}

#ifdef STATSD_HAVE_VTY
	/* Connect to VTY and reset stats */
	map(self:STATSVTY, system:STATSVTY);
	f_vty_set_prompts(STATSVTY);
	f_vty_transceive(STATSVTY, "enable");

	/* Reset the stats system at start */
	f_vty_transceive(STATSVTY, "stats reset");
#endif

	while (true) {
		alt {
		[] STATSD_PROC.getcall(STATSD_reset:{}) -> sender vc_conn {
#ifdef STATSD_HAVE_VTY
			f_vty_transceive(STATSVTY, "stats reset");
			STATSD_PROC.reply(STATSD_reset:{}) to vc_conn;
#else
			Misc_Helpers.f_shutdown(__BFILE__, __LINE__, fail,
						"STATSD_reset not supported, StatsD_Checker was built without VTY support");
#endif
			}
		[] STATSD_PROC.getcall(STATSD_snapshot:{?, ?}) -> param(keys, since_last_snapshot) sender vc_conn {
			snapshot := f_statsd_checker_snapshot(keys, since_last_snapshot);
			STATSD_PROC.reply(STATSD_snapshot:{keys, since_last_snapshot} value snapshot) to vc_conn;
			}
		[] STATSD_PROC.getcall(STATSD_expect:{?, ?, ?, ?}) -> param(expects, wait_converge, use_snapshot, snapshot) sender vc_conn {
			var boolean success := f_statsd_checker_expect(expects, wait_converge, use_snapshot, snapshot);
			STATSD_PROC.reply(STATSD_expect:{expects, wait_converge, use_snapshot, snapshot} value success) to vc_conn;
			}
		}
	}
}

/* Updates "metrics" & "seen" with content from "it". Returns true if the metric becomes known (for first time) as a result. */
private function f_statsd_metrics_update_value(inout StatsDMetrics metrics, inout Booleans seen, StatsDMetric it) return boolean
{
	for (var integer i := 0; i < lengthof(metrics); i := i + 1) {
		if (it.name != metrics[i].name or it.mtype != metrics[i].mtype) {
			continue;
		}
		metrics[i] := it;
		if (seen[i]) {
			return false;
		} else {
			seen[i] := true;
			return true;
		}
	}
	return false;
}


private function f_statsd_checker_snapshot(StatsDMetricKeys keys, boolean since_last_snapshot := true) runs on StatsD_Checker_CT return StatsDMetrics {
	var StatsDMessage msg;
	var StatsDMetrics metrics := {};
	var Booleans seen := {};
	var integer seen_remain := 0;

	for (var integer i := 0; i < lengthof(keys); i := i + 1) {
		metrics := metrics & {valueof(ts_StatsDMetric(keys[i].name, 0, keys[i].mtype))};
		seen := seen & {false};
		seen_remain := seen_remain + 1;
	}

	if (not since_last_snapshot) {
		STATS.clear;
	}

	T_statsd.start;
	while (seen_remain > 0) {
		var StatsD_RecvFrom rf;
		alt {
		[] STATS.receive(tr_StatsD_RecvFrom(?, ?)) -> value rf {
			msg := rf.msg;
			}
		[] T_statsd.timeout {
			for (var integer i := 0; i < lengthof(metrics); i := i + 1) {
				/* We're still missing some expects, keep looking */
				if (not seen[i]) {
					log("Timeout waiting for ", metrics[i].name);
				}
			}
			Misc_Helpers.f_shutdown(__BFILE__, __LINE__, fail,
						log2str("Timeout waiting for metrics: ", keys, seen));
		}
		}

		for (var integer i := 0; i < lengthof(msg); i := i + 1) {
			var StatsDMetric metric := msg[i];
			if (f_statsd_metrics_update_value(metrics, seen, metric)) {
				seen_remain := seen_remain - 1;
			}
		}
	}
	T_statsd.stop;

	return metrics;
}

private function get_val_from_snapshot(inout integer val, StatsDMetric metric, StatsDMetrics snapshot) return boolean
{
	for (var integer i := 0; i < lengthof(snapshot); i := i + 1) {
		if (metric.name != snapshot[i].name or metric.mtype != snapshot[i].mtype) {
			continue;
		}
		val := snapshot[i].val;
		return true;
	}
	return false;
}

/* Return false if the expectation doesn't match the metric, otherwise return true */
private function f_compare_expect(StatsDMetric metric,
				  StatsDExpect expect,
				  boolean use_snapshot := false,
				  StatsDMetrics snapshot := {}) return boolean {
	var integer val := 0;
	if ((metric.name != expect.name) or (metric.mtype != expect.mtype)) {
		return false;
	}
	if (use_snapshot) {
		var integer prev_val := 0;
		if (not get_val_from_snapshot(prev_val, metric, snapshot)) {
			Misc_Helpers.f_shutdown(__BFILE__, __LINE__, fail,
						log2str("Metric ", metric.name, " not found in snapshot ", snapshot));
		}
		val := metric.val - prev_val;
	} else {
		val := metric.val;
	}

	if ((val < expect.min) or (val > expect.max)) {
		return false;
	}
	return true;
}

private function f_statsd_checker_metric_expects(StatsDExpects expects,
						 StatsDMetric metric,
						 boolean use_snapshot := false,
						 StatsDMetrics snapshot := {})
return StatsDExpectResult {
	var StatsDExpectResult result := {
		kind := e_NotFound,
		idx := -1
	};

	for (var integer i := 0; i < lengthof(expects); i := i + 1) {
		var StatsDExpect exp := expects[i];
		if (exp.name != metric.name) {
			continue;
		}
		if (not f_compare_expect(metric, exp, use_snapshot, snapshot)) {
			log("EXP mismatch: ", metric, " vs exp ", exp, " | use_snapshot=", use_snapshot, ", snapshot=", snapshot);
			result := {
				kind := e_Mismatched,
				idx := i
			};
			break;
		} else {
			log("EXP match: ", metric, " vs exp ", exp);
			result := {
				kind := e_Matched,
				idx := i
			};
			break;
		}
	}
	return result;
}

private function f_statsd_checker_expect(StatsDExpects expects,
					 boolean wait_converge := false,
					 boolean use_snapshot := false,
					 StatsDMetrics snapshot := {}) runs on StatsD_Checker_CT return boolean {
	var default t;
	var StatsDMessage msg;
	var StatsDExpectResult res;
	var Booleans matched := {};
	var integer matched_remain := 0;

	for (var integer i := 0; i < lengthof(expects); i := i + 1) {
		matched := matched & {false};
		matched_remain := matched_remain + 1;
	}

	/* Dismiss any messages we might have skipped from the last report */
	STATS.clear;

	if (not use_snapshot) {
#ifdef STATSD_HAVE_VTY
		f_vty_transceive(STATSVTY, "stats report");
#else
		/* Assume caller knows previous state, eg. gauges may have been 0 due to IUT being reset */
#endif
	}

	T_statsd.start;
	while (matched_remain > 0) {
		var StatsD_RecvFrom rf;
		alt {
		[] STATS.receive(tr_StatsD_RecvFrom(?, ?)) -> value rf {
			msg := rf.msg;
			}
		[] T_statsd.timeout {
			for (var integer i := 0; i < lengthof(expects); i := i + 1) {
				/* We're still missing some expects, keep looking */
				if (not matched[i]) {
					log("Timeout waiting for ", expects[i].name,
					    " (min: ", expects[i].min, ", max: ", expects[i].max, ")");
				}
			}
			setverdict(fail, "Timeout waiting for metrics");
			return false;
		}
		}

		for (var integer i := 0; i < lengthof(msg); i := i + 1) {
			var StatsDMetric metric := msg[i];
			res := f_statsd_checker_metric_expects(expects, metric, use_snapshot, snapshot);
			if (res.kind == e_NotFound) {
				continue;
			}
			if (res.kind == e_Mismatched) {
				if (wait_converge and not matched[res.idx]) {
					log("Waiting convergence: Ignoring metric mismatch metric=", metric, " expect=", expects[res.idx])
					continue;
				}
				log("Metric: ", metric);
				log("Expect: ", expects[res.idx]);
				setverdict(fail, "Metric failed expectation ", metric, " vs ", expects[res.idx]);
				return false;
			}
			if (res.kind == e_Matched) {
				if (not matched[res.idx]) {
					matched[res.idx] := true;
					matched_remain := matched_remain - 1;
				}
				continue;
			}
		}
	}
	T_statsd.stop;
	return true;
}

function f_init_statsd(charstring id, inout StatsD_Checker_CT vc_STATSD, charstring local_addr, integer local_port) {
	id := id & "-STATS";

	vc_STATSD := StatsD_Checker_CT.create(id);
	vc_STATSD.start(StatsD_Checker.main(local_addr, local_port));
}


/* StatsD connhdlr */
function f_statsd_reset() runs on StatsD_ConnHdlr {
	if (not mp_enable_stats) {
		return;
	}

	STATSD_PROC.call(STATSD_reset:{}) {
		[] STATSD_PROC.getreply(STATSD_reset:{}) {}
	}
}

/* Useful to automatically generate param for f_statsd_snapshot() from StatsDExpects used in f_statsd_expect_from_snapshot() */
function f_statsd_keys_from_expect(StatsDExpects expects) return StatsDMetricKeys
{
	var StatsDMetricKeys keys := {}
	for (var integer i := 0; i < lengthof(expects); i := i + 1) {
		keys := keys & {valueof(ts_StatsDMetricKey(expects[i].name, expects[i].mtype))}
	}
	return keys;
}

/* Retrieve current values obtained at statsd server.
* If since_last_snapshot is false, then clear the received packets in port. */
function f_statsd_snapshot(StatsDMetricKeys keys, boolean since_last_snapshot := true) runs on StatsD_ConnHdlr return StatsDMetrics {
	var StatsDMetrics snapshot;
	if (not mp_enable_stats) {
		return {};
	}

	STATSD_PROC.call(STATSD_snapshot:{keys, since_last_snapshot}) {
		[] STATSD_PROC.getreply(STATSD_snapshot:{keys, since_last_snapshot}) -> value snapshot;
	}
	return snapshot;
}

function f_statsd_expect(StatsDExpects expects, boolean wait_converge := false) runs on StatsD_ConnHdlr return boolean {
	var boolean res;

	if (not mp_enable_stats) {
		return true;
	}

	STATSD_PROC.call(STATSD_expect:{expects, wait_converge, false, {}}) {
		[] STATSD_PROC.getreply(STATSD_expect:{expects, wait_converge, false, {}}) -> value res;
	}
	return res;
}

function f_statsd_expect_from_snapshot(StatsDExpects expects, boolean wait_converge := false,
				       StatsDMetrics snapshot := {}) runs on StatsD_ConnHdlr return boolean {
	var boolean res;

	if (not mp_enable_stats) {
		return true;
	}

	STATSD_PROC.call(STATSD_expect:{expects, wait_converge, true, snapshot}) {
		[] STATSD_PROC.getreply(STATSD_expect:{expects, wait_converge, true, snapshot}) -> value res;
	}
	return res;
}

}