// SPDX-License-Identifier: BSD-3-Clause
// Copyright (c) 2024 Robin Jarry

#include "metrics.h"
#include "module.h"
#include "port.h"
#include "rxtx.h"
#include "vec.h"
#include "worker.h"

#include <gr_infra.h>

#include <rte_common.h>
#include <rte_ethdev.h>

#include <fnmatch.h>

static bool skip_stat(const struct gr_stat *s, gr_stats_flags_t flags) {
	if (flags & GR_STATS_F_ZERO)
		return false;

	if (s->packets != 0)
		return false;

	if (strcmp(s->name, "idle") == 0 || strcmp(s->name, "overhead") == 0) {
		if (s->batches != 0)
			return false;
	}

	return true;
}

static struct api_out stats_get(const void *request, struct api_ctx *) {
	uint64_t rx_burst_histogram[RTE_GRAPH_BURST_SIZE + 1];
	const struct gr_stats_get_req *req = request;
	struct gr_stats_get_resp *resp = NULL;
	vec struct gr_stat *stats = NULL;
	size_t len, n_stats;
	struct gr_stat *s;
	int ret;

	if (req->flags & GR_STATS_F_SW) {
		stats = worker_dump_stats(req->cpu_id);
		if (stats == NULL && req->cpu_id != UINT16_MAX)
			return api_out(ENODEV, 0, NULL);
	}

	if (req->flags & GR_STATS_F_HW) {
		struct rte_eth_xstat_name *names = NULL;
		struct rte_eth_xstat *xstats = NULL;
		struct iface *iface = NULL;
		unsigned num;

		while ((iface = iface_next(GR_IFACE_TYPE_PORT, iface)) != NULL) {
			struct iface_info_port *port = iface_info_port(iface);

			// call first with NULL/0 to get the exact count
			if ((ret = rte_eth_xstats_get(port->port_id, NULL, 0)) < 0)
				goto err;
			if (ret == 0)
				continue;
			num = ret;

			if ((xstats = calloc(num, sizeof(*xstats))) == NULL) {
				ret = -ENOMEM;
				goto free_xstat;
			}
			if ((ret = rte_eth_xstats_get(port->port_id, xstats, num)) < 0)
				goto free_xstat;

			if ((names = calloc(num, sizeof(*names))) == NULL) {
				ret = -ENOMEM;
				goto free_xstat;
			}
			if ((ret = rte_eth_xstats_get_names(port->port_id, names, num)) < 0)
				goto free_xstat;

			// xstats and names are matched by array index
			for (unsigned i = 0; i < num; i++) {
				struct gr_stat stat = {
					.packets = xstats[i].value,
					.batches = 0,
					.cycles = 0,
				};
				// prefix each xstat name with interface name
				snprintf(
					stat.name,
					sizeof(stat.name),
					"%s.%s",
					iface->name,
					names[i].name
				);
				vec_add(stats, stat);
			}
free_xstat:
			free(xstats);
			xstats = NULL;
			free(names);
			names = NULL;
			if (ret < 0)
				goto err;

			rx_burst_histogram_get(
				port->port_id, rx_burst_histogram, ARRAY_DIM(rx_burst_histogram)
			);
			for (unsigned s = 0; s < ARRAY_DIM(rx_burst_histogram); s++) {
				struct gr_stat stat = {0};
				snprintf(
					stat.name,
					sizeof(stat.name),
					"%s.rx_burst_%u_pkts",
					iface->name,
					s
				);
				stat.batches = rx_burst_histogram[s];
				stat.packets = stat.batches * s;
				vec_add(stats, stat);
			}
		}
	}

	// iterate once to determine the number of stats matching pattern
	n_stats = 0;
	vec_foreach_ref (s, stats) {
		if (skip_stat(s, req->flags))
			continue;
		switch (fnmatch(req->pattern, s->name, 0)) {
		case 0:
			n_stats++;
		case FNM_NOMATCH:
			continue;
		default:
			ret = -errno;
			goto err;
		}
	}

	// allocate correct response size
	len = sizeof(*resp) + n_stats * sizeof(struct gr_stat);
	if ((resp = calloc(1, len)) == NULL) {
		ret = -ENOMEM;
		goto err;
	}

	// fill in response
	vec_foreach_ref (s, stats) {
		if (skip_stat(s, req->flags))
			continue;
		switch (fnmatch(req->pattern, s->name, 0)) {
		case 0:
			resp->stats[resp->n_stats++] = *s;
		case FNM_NOMATCH:
			continue;
		default:
			ret = -errno;
			goto err;
		}
	}

	vec_free(stats);
	return api_out(0, len, resp);
err:
	vec_free(stats);
	free(resp);
	return api_out(-ret, 0, NULL);
}

static struct api_out stats_reset(const void * /*request*/, struct api_ctx *) {
	struct worker *worker;
	struct iface *iface;
	int ret;

	STAILQ_FOREACH (worker, &workers, next) {
		atomic_store(&worker->stats_reset, true);
		worker_wakeup(worker);
	}

	iface = NULL;

	// Reset software stats for all interfaces.
	memset(iface_stats, 0, sizeof(iface_stats));

	while ((iface = iface_next(GR_IFACE_TYPE_PORT, iface)) != NULL) {
		struct iface_info_port *port = iface_info_port(iface);
		if ((ret = rte_eth_stats_reset(port->port_id)) < 0)
			return api_out(-ret, 0, NULL);
		if ((ret = rte_eth_xstats_reset(port->port_id)) < 0)
			return api_out(-ret, 0, NULL);
	}

	rx_burst_histogram_reset();

	return api_out(0, 0, NULL);
}

static struct api_out iface_stats_get(const void * /*request*/, struct api_ctx *) {
	struct gr_iface_stats_get_resp *resp = NULL;
	vec struct gr_iface_stats *stats_vec = NULL;
	struct iface *iface = NULL;
	int ret = 0;

	while ((iface = iface_next(GR_IFACE_TYPE_UNDEF, iface)) != NULL) {
		// Create a single stats object per interface
		struct gr_iface_stats s;
		s.iface_id = iface->id;
		s.rx_packets = 0;
		s.rx_bytes = 0;
		s.rx_drops = 0;
		s.tx_packets = 0;
		s.tx_bytes = 0;
		s.tx_errors = 0;
		s.cp_rx_packets = 0;
		s.cp_rx_bytes = 0;
		s.cp_tx_packets = 0;
		s.cp_tx_bytes = 0;

		// Aggregate per-core stats
		for (int i = 0; i < RTE_MAX_LCORE; i++) {
			struct iface_stats *sw_stats = iface_get_stats(i, iface->id);
			s.rx_packets += sw_stats->rx_packets;
			s.rx_bytes += sw_stats->rx_bytes;
			s.tx_packets += sw_stats->tx_packets;
			s.tx_bytes += sw_stats->tx_bytes;
			s.cp_rx_packets += sw_stats->cp_rx_packets;
			s.cp_rx_bytes += sw_stats->cp_rx_bytes;
			s.cp_tx_packets += sw_stats->cp_tx_packets;
			s.cp_tx_bytes += sw_stats->cp_tx_bytes;
		}

		if (iface->type == GR_IFACE_TYPE_PORT) {
			// If possible, use the hardware statistics from the driver
			struct iface_info_port *port = iface_info_port(iface);
			struct rte_eth_stats stats = {0};
			if (rte_eth_stats_get(port->port_id, &stats) == 0) {
				s.rx_drops = stats.imissed;
				s.tx_errors = stats.oerrors;
			}
		}

		vec_add(stats_vec, s);
	}

	size_t n_stats = vec_len(stats_vec);
	size_t len = sizeof(*resp) + n_stats * sizeof(struct gr_iface_stats);
	if ((resp = calloc(1, len)) == NULL) {
		ret = -ENOMEM;
		goto err;
	}

	resp->n_stats = n_stats;
	if (vec_len(stats_vec) > 0)
		memcpy(resp->stats, stats_vec, n_stats * sizeof(struct gr_iface_stats));

	vec_free(stats_vec);
	return api_out(0, len, resp);
err:
	vec_free(stats_vec);
	free(resp);
	return api_out(-ret, 0, NULL);
}

METRIC_COUNTER(m_packets, "node_packets", "Number of packets processed by a node.");
METRIC_COUNTER(m_batches, "node_batches", "Number of times a node was visited.");
METRIC_COUNTER(m_cycles, "node_cycles", "Number of cycles spent per node.");

static void graph_metrics_collect(struct metrics_writer *w) {
	vec struct gr_stat *stats = worker_dump_stats(UINT16_MAX);
	struct metrics_ctx ctx;

	vec_foreach_ref (const struct gr_stat *s, stats) {
		metrics_ctx_init(&ctx, w, "name", s->name, NULL);
		metric_emit(&ctx, &m_packets, s->packets);
		metric_emit(&ctx, &m_batches, s->batches);
		metric_emit(&ctx, &m_cycles, s->cycles);
	}

	vec_free(stats);
}

static struct metrics_collector graph_collector = {
	.name = "graph",
	.collect = graph_metrics_collect,
};

METRIC_COUNTER(m_idle_cycles, "idle_cycles", "Number of idle CPU cycles.");
METRIC_COUNTER(m_busy_cycles, "busy_cycles", "Number of busy CPU cycles.");

static void cpu_collect_metrics(struct metrics_writer *w) {
	struct metrics_ctx ctx;
	char cpu[16], numa[16];
	struct worker *worker;

	STAILQ_FOREACH (worker, &workers, next) {
		const struct worker_stats *stats = atomic_load(&worker->stats);
		if (stats == NULL)
			continue;

		snprintf(cpu, sizeof(cpu), "%u", worker->cpu_id);
		snprintf(numa, sizeof(numa), "%u", rte_lcore_to_socket_id(worker->lcore_id));

		metrics_ctx_init(&ctx, w, "cpu", cpu, "numa", numa, NULL);
		metric_emit(&ctx, &m_idle_cycles, stats->total_cycles - stats->busy_cycles);
		metric_emit(&ctx, &m_busy_cycles, stats->busy_cycles);
	}
}

static struct metrics_collector cpu_collector = {
	.name = "cpu",
	.collect = cpu_collect_metrics,
};

METRIC_HISTOGRAM(m_rx_burst, "rx_burst_size", "Distribution of RX burst sizes in packets.");

static const unsigned rx_burst_buckets[] = {0, 1, 2, 4, 8, 16, 32, 64, 128};

static void rx_burst_metrics_collect(struct metrics_writer *w) {
	uint64_t histogram[RTE_GRAPH_BURST_SIZE + 1];
	struct metrics_ctx ctx;
	struct iface *iface = NULL;

	while ((iface = iface_next(GR_IFACE_TYPE_PORT, iface)) != NULL) {
		struct iface_info_port *port = iface_info_port(iface);

		rx_burst_histogram_get(port->port_id, histogram, ARRAY_DIM(histogram));
		metrics_ctx_init(&ctx, w, "iface", iface->name, NULL);
		metric_emit_histogram(
			&ctx,
			&m_rx_burst,
			histogram,
			ARRAY_DIM(histogram),
			rx_burst_buckets,
			ARRAY_DIM(rx_burst_buckets)
		);
	}
}

static struct metrics_collector rx_burst_collector = {
	.name = "rx_burst",
	.collect = rx_burst_metrics_collect,
};

RTE_INIT(infra_stats_init) {
	api_handler(GR_STATS_GET, stats_get);
	api_handler(GR_STATS_RESET, stats_reset);
	api_handler(GR_IFACE_STATS_GET, iface_stats_get);
	metrics_register(&graph_collector);
	metrics_register(&cpu_collector);
	metrics_register(&rx_burst_collector);
}
