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

#include "graph.h"
#include "log.h"
#include "module.h"
#include "port.h"
#include "rxtx.h"
#include "vec.h"
#include "worker.h"

#include <gr_infra.h>
#include <gr_string.h>

#include <rte_build_config.h>
#include <rte_errno.h>
#include <rte_ethdev.h>
#include <rte_malloc.h>

#include <stdatomic.h>
#include <sys/queue.h>
#include <unistd.h>

LOG_TYPE("graph");

struct node_infos node_infos = STAILQ_HEAD_INITIALIZER(node_infos);
static vec const char **base_node_names;
static vec char **rx_node_names;
static vec char **tx_node_names;
static rte_node_t port_rx_node;
static rte_node_t port_tx_node;
static rte_node_t port_output_node;
static rte_edge_t port_output_invalid;

rte_edge_t gr_node_attach_parent(const char *parent, const char *node) {
	rte_node_t parent_id;
	rte_edge_t edge;
	char **names;

	if ((parent_id = rte_node_from_name(parent)) == RTE_NODE_ID_INVALID)
		ABORT("'%s' parent node not found", parent);

	edge = rte_node_edge_update(parent_id, RTE_EDGE_ID_INVALID, &node, 1);
	if (edge == RTE_EDGE_ID_INVALID)
		ABORT("rte_node_edge_update: %s", rte_strerror(rte_errno));

	names = calloc(rte_node_edge_count(parent_id), sizeof(char *));
	if (names == NULL)
		ABORT("calloc(rte_node_edge_count('%s')) failed", parent);
	if (rte_node_edge_get(parent_id, names) == RTE_EDGE_ID_INVALID)
		ABORT("rte_node_edge_get('%s')) failed", parent);
	for (edge = 0; edge < rte_node_edge_count(parent_id); edge++) {
		if (strcmp(names[edge], node) == 0)
			break;
	}
	free(names);
	if (edge == rte_node_edge_count(parent_id))
		ABORT("cannot find added edge");

	LOG(DEBUG, "attach %s -> %s (edge=%u)", parent, node, edge);

	return edge;
}

void worker_graph_free(struct worker *worker) {
	int ret;
	for (int i = 0; i < 2; i++) {
		if (worker->graph[i] != NULL) {
			if ((ret = rte_graph_destroy(worker->graph[i]->id)) < 0)
				LOG(ERR, "rte_graph_destroy: %s", rte_strerror(-ret));
			worker->graph[i] = NULL;
		}
	}
}

static struct iface_info_port *find_port(vec struct iface_info_port **ports, uint16_t port_id) {
	struct iface_info_port *port = NULL;
	vec_foreach (struct iface_info_port *p, ports) {
		if (p->port_id == port_id) {
			port = p;
			break;
		}
	}
	assert(port != NULL);
	return port;
}

static struct gr_graph_conf graph_conf = {
	.rx_burst_max = 64,
	.vector_max = 64,
};

static int
worker_graph_new(struct worker *worker, uint8_t index, vec struct iface_info_port **ports) {
	vec const char **graph_nodes = NULL;
	char graph_name[RTE_GRAPH_NAMESIZE];
	char node_name[RTE_NODE_NAMESIZE];
	vec char **rx_nodes = NULL;
	struct iface_info_port *port;
	struct queue_map *qmap;
	struct rte_node *node;
	uint16_t burst_size;
	uint16_t graph_uid;
	unsigned n_rxqs;
	int ret = 0;

	n_rxqs = 0;
	vec_foreach_ref (qmap, worker->rxqs) {
		if (ports != NULL && qmap->enabled)
			n_rxqs++;
	}
	if (n_rxqs == 0) {
		worker->graph[index] = NULL;
		return 0;
	}

	// unique suffix for this graph
	graph_uid = (worker->cpu_id << 1) | (0x1 & index);
	snprintf(graph_name, sizeof(graph_name), "gr-%04x", graph_uid);

	// generate graph nodes list
	vec_foreach_ref (qmap, worker->rxqs) {
		if (!qmap->enabled)
			continue;
		char *name = astrcat(NULL, RX_NODE_FMT, qmap->port_id, qmap->queue_id);
		vec_add(rx_nodes, name);
		vec_add(graph_nodes, name);
	}

	vec_extend(graph_nodes, base_node_names);
	vec_extend(graph_nodes, tx_node_names);

	// graph init
	struct rte_graph_param params = {
		.socket_id = rte_lcore_to_socket_id(worker->lcore_id),
		.nb_node_patterns = vec_len(graph_nodes),
		.node_patterns = graph_nodes,
	};
	if (rte_graph_create(graph_name, &params) == RTE_GRAPH_ID_INVALID) {
		if (rte_errno == 0)
			rte_errno = EINVAL;
		ret = -rte_errno;
		goto out;
	}
	worker->graph[index] = rte_graph_lookup(graph_name);

	// set rx nodes context data
	vec_foreach_ref (qmap, worker->rxqs) {
		if (!qmap->enabled)
			continue;
		snprintf(node_name, sizeof(node_name), RX_NODE_FMT, qmap->port_id, qmap->queue_id);
		node = rte_graph_node_get_by_name(graph_name, node_name);
		struct rx_node_ctx *ctx = rx_node_ctx(node);
		port = find_port(ports, qmap->port_id);
		ctx->iface = RTE_PTR_SUB(port, offsetof(struct iface, info));
		ctx->rxq.port_id = qmap->port_id;
		ctx->rxq.queue_id = qmap->queue_id;
		burst_size = RTE_MAX(1u, graph_conf.vector_max / vec_len(rx_nodes));
		ctx->burst_size = RTE_MIN(graph_conf.rx_burst_max, burst_size);
		LOG(DEBUG,
		    "[CPU %d] <- port=%u rxq=%u burst_size=%u",
		    worker->cpu_id,
		    qmap->port_id,
		    qmap->queue_id,
		    ctx->burst_size);
		// select the appropriate node process callback
		if (ctx->iface->mode == GR_IFACE_MODE_BOND) {
			// virtio driver supports Rx vlan stripping but requires help for L4 csum
			if (port->virtio_offloads)
				node->process = rx_bond_virtio_process;
			else if (port->rx_offloads & RTE_ETH_RX_OFFLOAD_VLAN_STRIP)
				node->process = rx_bond_offload_process;
			else
				node->process = rx_bond_process;
		} else {
			// virtio driver supports Rx vlan stripping but requires help for L4 csum
			if (port->virtio_offloads)
				node->process = rx_virtio_process;
			else if (port->rx_offloads & RTE_ETH_RX_OFFLOAD_VLAN_STRIP)
				node->process = rx_offload_process;
			else
				node->process = rx_process;
		}
	}

	// initialize all tx nodes context to invalid ports and queues
	vec_foreach (const char *name, tx_node_names) {
		node = rte_graph_node_get_by_name(graph_name, name);
		struct tx_node_ctx *ctx = tx_node_ctx(node);
		ctx->txq.port_id = UINT16_MAX;
		ctx->txq.queue_id = UINT16_MAX;
		ctx->lock = NULL;
	}

	// initialize the port_output node context to point to the port_output_invalid edge
	struct port_output_edges *out = rte_malloc(__func__, sizeof(*out), RTE_CACHE_LINE_SIZE);
	if (out == NULL) {
		ret = -ENOMEM;
		goto out;
	}
	for (unsigned i = 0; i < ARRAY_DIM(out->edges); i++)
		out->edges[i] = port_output_invalid;

	vec_foreach_ref (qmap, worker->txqs) {
		if (!qmap->enabled)
			continue;
		LOG(DEBUG,
		    "[CPU %d] -> port %u txq %u",
		    worker->cpu_id,
		    qmap->port_id,
		    qmap->queue_id);
		// find the corresponding port_tx clone for this port-txq pair
		snprintf(node_name, sizeof(node_name), TX_NODE_FMT, qmap->port_id, qmap->queue_id);
		node = rte_graph_node_get_by_name(graph_name, node_name);
		// and update its context data to correct values
		struct tx_node_ctx *ctx = tx_node_ctx(node);
		ctx->txq.port_id = qmap->port_id;
		ctx->txq.queue_id = qmap->queue_id;
		port = find_port(ports, qmap->port_id);

		// set spinlock only if multiple workers share this TX queue
		unsigned txq_users = 0;
		struct worker *w = NULL;
		STAILQ_FOREACH (w, &workers, next) {
			vec_foreach_ref (struct queue_map *q, w->txqs) {
				if (q->port_id == qmap->port_id && q->queue_id == qmap->queue_id)
					txq_users++;
			}
		}
		if (txq_users > 1) {
			ctx->lock = &port->txq_locks[qmap->queue_id];
			node->process = tx_shared_process;
			LOG(WARNING,
			    "[CPU %d] port %s txq %u shared by %u workers",
			    worker->cpu_id,
			    port->devargs,
			    qmap->queue_id,
			    txq_users);
		} else {
			node->process = tx_process;
		}

		for (rte_edge_t edge = 0; edge < vec_len(tx_node_names); edge++) {
			if (strcmp(tx_node_names[edge], node_name) == 0) {
				// update the port_output context data to map this port to the
				// correct edge
				out->edges[ctx->txq.port_id] = edge;
				break;
			}
		}
	}

	// finally, set the port_output context data
	rte_graph_node_get_by_name(graph_name, "port_output")->ctx_ptr = out;

out:
	vec_free(graph_nodes);
	strvec_free(rx_nodes);

	return errno_set(-ret);
}

int worker_graph_reload(struct worker *worker, vec struct iface_info_port **ports) {
	unsigned next = !atomic_load(&worker->cur_config);
	int ret;

	if ((ret = worker_graph_new(worker, next, ports)) < 0)
		return errno_log(-ret, "worker_graph_new");

	// wait for datapath worker to pickup the config update
	atomic_store(&worker->next_config, next);
	worker_wakeup(worker);
	for (unsigned i = 0; atomic_load(&worker->cur_config) != next; i++) {
		if (i >= 10000) // 10000 * 500us -> 5s
			return errno_log(ETIMEDOUT, "worker not responding to graph reload");
		usleep(500);
	}

	// free old config
	next = !next;

	if (worker->graph[next] != NULL) {
		if ((ret = rte_graph_destroy(worker->graph[next]->id)) < 0)
			errno_log(-ret, "rte_graph_destroy");
		worker->graph[next] = NULL;
	}

	return 0;
}

static char *ensure_queue_node(
	rte_node_t base_node,
	const char *name_fmt,
	uint16_t port_id,
	uint16_t queue_id,
	vec char **old_list
) {
	char *name = astrcat(NULL, name_fmt, port_id, queue_id);
	assert(name != NULL);

	// remove node from the old list so that only unused nodes remain
	for (unsigned i = 0; i < vec_len(old_list); i++) {
		if (strcmp(old_list[i], name) == 0) {
			free(old_list[i]);
			vec_del(old_list, i);
			break;
		}
	}
	if (rte_node_from_name(name) == RTE_NODE_ID_INVALID) {
		// node does not exist yet, clone it from the base
		rte_node_t node = rte_node_clone(base_node, strstr(name, "-") + 1);
		assert(node != RTE_NODE_ID_INVALID);
	}

	return name;
}

static vec rte_node_t *worker_graph_nodes_add_missing(vec struct iface_info_port **ports) {
	vec char **old_rx = vec_clone(rx_node_names);
	vec char **old_tx = vec_clone(tx_node_names);
	rte_node_t *unused_nodes = NULL;
	rte_edge_t edge;
	char *name;

	vec_free(rx_node_names);
	vec_free(tx_node_names);

	// clone port_rx and port_tx to match all possible port-queue pairs
	vec_foreach (struct iface_info_port *port, ports) {
		for (uint16_t rxq = 0; rxq < port->n_rxq; rxq++) {
			name = ensure_queue_node(
				port_rx_node, RX_NODE_FMT, port->port_id, rxq, old_rx
			);
			vec_add(rx_node_names, name);
		}
		for (uint16_t txq = 0; txq < port->n_txq; txq++) {
			name = ensure_queue_node(
				port_tx_node, TX_NODE_FMT, port->port_id, txq, old_tx
			);
			vec_add(tx_node_names, name);
		}
	}

	// update the port_output edges to allow reaching all possible port_tx clones
	edge = rte_node_edge_update(
		port_output_node, 0, (const char **)tx_node_names, vec_len(tx_node_names)
	);
	assert(edge != RTE_EDGE_ID_INVALID);
	port_output_invalid = vec_len(tx_node_names);

	// this error edge must be last
	const char *invalid_port = "port_output_invalid";
	edge = rte_node_edge_update(port_output_node, vec_len(tx_node_names), &invalid_port, 1);
	assert(edge != RTE_EDGE_ID_INVALID);

	edge = rte_node_edge_shrink(port_output_node, vec_len(tx_node_names) + 1);
	assert(edge != RTE_EDGE_ID_INVALID);

	// store all unused node_ids in a list to be returned to the caller
	vec_foreach (name, old_rx)
		vec_add(unused_nodes, rte_node_from_name(name));
	vec_foreach (name, old_tx)
		vec_add(unused_nodes, rte_node_from_name(name));
	strvec_free(old_rx);
	strvec_free(old_tx);

	return unused_nodes;
}

int worker_graph_reload_all(vec struct iface_info_port **ports) {
	struct worker *worker;
	int ret;

	vec rte_node_t *unused_nodes = worker_graph_nodes_add_missing(ports);

	// stop all workers by switching them to NULL graphs
	STAILQ_FOREACH (worker, &workers, next) {
		if ((ret = worker_graph_reload(worker, NULL)) < 0)
			return ret;
	}

	// create new graphs and start all workers
	STAILQ_FOREACH (worker, &workers, next) {
		if ((ret = worker_graph_reload(worker, ports)) < 0)
			return ret;
	}

	// these port_rx and port_tx clones are now not referenced in any graph
	// we can safely delete them
	vec_foreach (rte_node_t node, unused_nodes)
		rte_node_free(node);
	vec_free(unused_nodes);

	return 0;
}

const struct gr_node_info *gr_node_info_get(rte_node_t node_id) {
	const struct gr_node_info *info;

	STAILQ_FOREACH (info, &node_infos, next)
		if (info->node->id == node_id)
			return info;

	return errno_set_null(ENOENT);
}

static struct api_out graph_dump(const void *request, struct api_ctx *) {
	static const struct {
		gr_node_type_t type;
		const char *name;
	} layers[] = {
		{GR_NODE_T_L1, "L1"},
		{GR_NODE_T_L2, "L2"},
		{GR_NODE_T_L3, "L3"},
		{GR_NODE_T_L4, "L4"},
	};
	const struct gr_graph_dump_req *req = request;
	vec const char **seen_edges = NULL;
	struct gr_node_info *info;
	char **edges = NULL;
	size_t buf_len = 0;
	char *buf = NULL;
	FILE *f = NULL;

	if ((f = open_memstream(&buf, &buf_len)) == NULL)
		return api_out(errno, 0, NULL);

	if (fprintf(f, "digraph {\n") < 0)
		goto err;
	if (!req->by_layer) {
		if (fprintf(f, "\trankdir=LR;\n") < 0)
			goto err;
	}
	if (!req->compact || req->by_layer) {
		if (fprintf(f, "\tnewrank=true;\n") < 0)
			goto err;
	}
	if (fprintf(f, "\tnode [margin=0.02 fontsize=11 fontname=sans];\n") < 0)
		goto err;
	if (fprintf(f, "\tgraph [fontsize=11 fontname=sans compound=true style=dotted];\n") < 0)
		goto err;

	for (unsigned i = 0; i < ARRAY_DIM(layers); i++) {
		if (req->by_layer) {
			if (fprintf(f, "\tsubgraph \"cluster_%s\" {\n", layers[i].name) < 0)
				goto err;
			if (fprintf(f, "\t\tlabel=\"%s\";\n", layers[i].name) < 0)
				goto err;
			if (fprintf(f, "\t\tl%u [style=invis];\n", i + 1) < 0)
				goto err;
		}

		STAILQ_FOREACH (info, &node_infos, next) {
			rte_node_t node_id = rte_node_from_name(info->node->name);
			unsigned nb_edges = rte_node_edge_count(node_id);
			const char *name = info->node->name;
			const char *attrs = "";

			if (!(info->type & layers[i].type))
				continue;

			if (node_id == port_output_node)
				nb_edges = info->node->nb_edges;

			if (!req->full) {
				if (nb_edges == 0)
					continue;
				if (strstr(name, "error"))
					continue;
				if (strcmp(name, "control_input") == 0)
					continue;
				if (strcmp(name, "control_output") == 0)
					continue;
			}
			if (fprintf(f, "\t\t\"%s\"", name) < 0)
				goto err;

			if (info->node->flags & RTE_NODE_SOURCE_F) {
				attrs = " [color=blue style=bold]";
				if (fprintf(f, "%s", attrs) < 0)
					goto err;
			} else if (info->type & GR_NODE_T_CONTROL) {
				attrs = " [color=dodgerblue shape=box margin=0.1]";
				if (fprintf(f, "%s", attrs) < 0)
					goto err;
			} else if (nb_edges == 0) {
				if (fprintf(f, " [fontcolor=darkorange shape=plain]") < 0)
					goto err;
			}

			if (fprintf(f, ";\n") < 0)
				goto err;
		}

		if (req->by_layer && fprintf(f, "\t}\n") < 0)
			goto err;
	}

	if (req->by_layer) {
		for (unsigned i = 1; i < ARRAY_DIM(layers); i++) {
			if (fprintf(f, "\tl%u -> l%u [style=invis weight=10];\n", i, i + 1) < 0)
				goto err;
		}
	}

	STAILQ_FOREACH (info, &node_infos, next) {
		rte_node_t node_id = rte_node_from_name(info->node->name);
		unsigned nb_edges = rte_node_edge_count(node_id);
		const char *name = info->node->name;
		const char *attrs = "";

		if (node_id == port_output_node)
			nb_edges = info->node->nb_edges;

		if (nb_edges == 0)
			continue;
		if (!req->full) {
			if (strstr(name, "error"))
				continue;
			if (strcmp(name, "control_input") == 0)
				continue;
			if (strcmp(name, "control_output") == 0)
				continue;
		}

		if (info->node->flags & RTE_NODE_SOURCE_F) {
			if (req->by_layer)
				attrs = " [color=blue style=bold constraint=false]";
			else
				attrs = " [color=blue style=bold]";
		} else if (info->type & GR_NODE_T_CONTROL) {
			if (req->by_layer)
				attrs = " [color=dodgerblue constraint=false]";
			else
				attrs = " [color=dodgerblue]";
		}

		if ((edges = calloc(nb_edges, sizeof(char *))) == NULL)
			goto err;
		if (node_id == port_output_node) {
			for (unsigned i = 0; i < nb_edges; i++)
				edges[i] = (char *)info->node->next_nodes[i];
		} else if (rte_node_edge_get(node_id, edges) == RTE_EDGE_ID_INVALID)
			goto err;

		for (unsigned i = 0; i < nb_edges; i++) {
			const char *edge = edges[i];
			const char *node_attrs = attrs;
			rte_node_t id = rte_node_from_name(edge);
			unsigned num = rte_node_edge_count(id);

			if (!req->full) {
				if (strstr(edge, "error"))
					continue;
				if (strcmp(edge, "control_input") == 0)
					continue;
				if (strcmp(edge, "control_output") == 0)
					continue;
			}
			if (id == port_output_node) {
				num = gr_node_info_get(id)->node->nb_edges;
			}
			if (num == 0) {
				if (!req->full)
					continue;
				if (req->by_layer)
					node_attrs = " [color=darkorange constraint=false]";
				else
					node_attrs = " [color=darkorange]";
			}
			vec_foreach (const char *e, seen_edges) {
				if (strcmp(e, edge) == 0)
					goto skip; // skip duplicate edges
			}
			vec_add(seen_edges, edge);

			if (fprintf(f, "\t\"%s\" -> \"%s\"%s;\n", name, edge, node_attrs) < 0)
				goto err;
skip:;
		}
		free(edges);
		edges = NULL;
		vec_free(seen_edges);
	}

	// terminate with nul character
	if (fprintf(f, "}\n%c", '\0') < 0)
		goto err;

	fflush(f);
	fclose(f);

	return api_out(0, buf_len, buf);

err:
	int errsave = errno;
	fclose(f);
	free(buf);
	free(edges);
	vec_free(seen_edges);
	return api_out(errsave, 0, NULL);
}

static struct api_out graph_conf_get(const void * /*request*/, struct api_ctx *) {
	struct gr_graph_conf *resp = malloc(sizeof(*resp));
	if (resp == NULL)
		return api_out(ENOMEM, 0, NULL);
	*resp = graph_conf;
	return api_out(0, sizeof(*resp), resp);
}

static struct api_out graph_conf_set(const void *request, struct api_ctx *) {
	const struct gr_graph_conf *req = request;
	struct gr_graph_conf prev = graph_conf;
	vec struct iface_info_port **ports = NULL;
	struct iface *iface = NULL;
	int ret;

	if (req->rx_burst_max > RTE_GRAPH_BURST_SIZE || req->vector_max > RTE_GRAPH_BURST_SIZE)
		return api_out(EOVERFLOW, 0, NULL);

	if (req->rx_burst_max > 0)
		graph_conf.rx_burst_max = req->rx_burst_max;
	if (req->vector_max > 0)
		graph_conf.vector_max = req->vector_max;

	if (graph_conf.rx_burst_max == prev.rx_burst_max
	    && graph_conf.vector_max == prev.vector_max)
		return api_out(0, 0, NULL); // no change

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

	LOG(NOTICE,
	    "vector_max=%u rx_burst_max=%u",
	    graph_conf.vector_max,
	    graph_conf.rx_burst_max);

	ret = worker_graph_reload_all(ports);
	if (ret < 0) {
		graph_conf = prev;
		worker_graph_reload_all(ports);
	}

	vec_free(ports);

	return api_out(-ret, 0, NULL);
}

static void graph_init(struct event_base *) {
	struct rte_node_register *reg;
	struct gr_node_info *info;

	// register nodes first
	STAILQ_FOREACH (info, &node_infos, next) {
		if (info->node == NULL)
			ABORT("info->node == NULL");
		reg = info->node;
		reg->parent_id = RTE_NODE_ID_INVALID;
		if (reg->xstats && reg->xstats->nb_xstats > GR_MAX_NODE_XSTATS)
			ABORT("node %s has %u xstats (max %u)",
			      reg->name,
			      reg->xstats->nb_xstats,
			      GR_MAX_NODE_XSTATS);
		reg->id = __rte_node_register(reg);
		if (reg->id == RTE_NODE_ID_INVALID)
			ABORT("__rte_node_register(%s): %s", reg->name, rte_strerror(rte_errno));

		if (strcmp(reg->name, RX_NODE_BASE) == 0)
			port_rx_node = reg->id;
		else if (strcmp(reg->name, TX_NODE_BASE) == 0)
			port_tx_node = reg->id;
		else
			vec_add(base_node_names, reg->name);

		if (strcmp(reg->name, "port_output") == 0)
			port_output_node = reg->id;
	}

	// then, invoke all registration callbacks where applicable
	STAILQ_FOREACH (info, &node_infos, next) {
		if (info->register_callback != NULL) {
			info->register_callback();
		}
	}
}

static void graph_fini(struct event_base *) {
	struct gr_node_info *info;

	STAILQ_FOREACH (info, &node_infos, next) {
		if (info->unregister_callback != NULL) {
			info->unregister_callback();
		}
	}

	vec_free(base_node_names);
	strvec_free(rx_node_names);
	strvec_free(tx_node_names);
}

static struct module graph_module = {
	.name = "graph",
	.depends_on = "mempool",
	.init = graph_init,
	.fini = graph_fini,
};

RTE_INIT(control_graph_init) {
	api_handler(GR_GRAPH_DUMP, graph_dump);
	api_handler(GR_GRAPH_CONF_GET, graph_conf_get);
	api_handler(GR_GRAPH_CONF_SET, graph_conf_set);
	module_register(&graph_module);
}
