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

#include "config.h"
#include "datapath.h"
#include "graph.h"
#include "log.h"
#include "module.h"
#include "port.h"
#include "sys_queue.h"
#include "vec.h"
#include "worker.h"

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

#include <numa.h>
#include <rte_common.h>
#include <rte_ethdev.h>
#include <rte_graph.h>
#include <rte_malloc.h>

#include <errno.h>
#include <pthread.h>
#include <sched.h>
#include <stdatomic.h>
#include <sys/queue.h>
#include <unistd.h>

LOG_TYPE("worker");

struct workers workers = STAILQ_HEAD_INITIALIZER(workers);

int worker_create(unsigned cpu_id) {
	struct worker *worker = rte_zmalloc(__func__, sizeof(*worker), 0);
	pthread_attr_t attr;
	cpu_set_t cpuset;
	int ret = ENOMEM;

	if (worker == NULL)
		goto end;

	worker->cpu_id = cpu_id;
	worker->lcore_id = LCORE_ID_ANY;
	pthread_mutex_init(&worker->wakeup.lock, NULL);
	pthread_cond_init(&worker->wakeup.cond, NULL);

	CPU_ZERO(&cpuset);
	CPU_SET(cpu_id, &cpuset);
	pthread_attr_init(&attr);
	if (!!(ret = pthread_attr_setaffinity_np(&attr, sizeof(cpuset), &cpuset)))
		goto end;

	if (!!(ret = pthread_create(&worker->thread, &attr, gr_datapath_loop, worker)))
		goto end;

	STAILQ_INSERT_TAIL(&workers, worker, next);

	// wait until thread has initialized lcore_id
	for (unsigned i = 0; !atomic_load(&worker->started); i++) {
		if (i >= 10000) { // 10000 * 500us -> 5s
			ret = ETIMEDOUT;
			break;
		}
		usleep(500);
	}

end:
	pthread_attr_destroy(&attr);

	if (ret == 0) {
		LOG(INFO, "worker %u started", cpu_id);
	} else {
		if (worker != NULL) {
			STAILQ_REMOVE(&workers, worker, worker, next);
			pthread_cancel(worker->thread);
			pthread_cond_destroy(&worker->wakeup.cond);
			pthread_mutex_destroy(&worker->wakeup.lock);
			rte_free(worker);
		}
		LOG(ERR, "worker %u start failed: %s", cpu_id, strerror(ret));
	}

	return errno_set(ret);
}

int worker_destroy(unsigned cpu_id) {
	struct worker *worker = worker_find(cpu_id);

	if (worker == NULL)
		return errno_log(ENOENT, "worker_find");

	STAILQ_REMOVE(&workers, worker, worker, next);

	atomic_store(&worker->shutdown, true);
	worker_wakeup(worker);
	pthread_join(worker->thread, NULL);
	pthread_cond_destroy(&worker->wakeup.cond);
	pthread_mutex_destroy(&worker->wakeup.lock);
	worker_graph_free(worker);
	vec_free(worker->rxqs);
	vec_free(worker->txqs);
	rte_free(worker);

	LOG(INFO, "worker %d destroyed", cpu_id);
	return 0;
}

void worker_wait_wakeup(struct worker *w) {
	pthread_mutex_lock(&w->wakeup.lock);
	while (!w->wakeup.set)
		pthread_cond_wait(&w->wakeup.cond, &w->wakeup.lock);
	w->wakeup.set = false;
	pthread_mutex_unlock(&w->wakeup.lock);
}

void worker_wakeup(struct worker *w) {
	pthread_mutex_lock(&w->wakeup.lock);
	w->wakeup.set = true;
	pthread_cond_signal(&w->wakeup.cond);
	pthread_mutex_unlock(&w->wakeup.lock);
}

unsigned worker_count(void) {
	struct worker *worker;
	unsigned count = 0;

	STAILQ_FOREACH (worker, &workers, next)
		count++;

	return count;
}

struct worker *worker_find(unsigned cpu_id) {
	struct worker *worker;
	STAILQ_FOREACH (worker, &workers, next) {
		if (worker->cpu_id == cpu_id)
			return worker;
	}
	errno = ENOENT;
	return NULL;
}

int port_unplug(struct iface_info_port *p) {
	vec struct iface_info_port **ports = NULL;
	struct iface *iface = NULL;
	struct queue_map *qmap;
	struct worker *worker;
	int changed = 0;
	int ret;

	STAILQ_FOREACH (worker, &workers, next) {
		vec_foreach_ref (qmap, worker->rxqs) {
			if (qmap->port_id == p->port_id) {
				qmap->enabled = false;
				changed++;
			}
		}
		vec_foreach_ref (qmap, worker->txqs) {
			if (qmap->port_id == p->port_id) {
				qmap->enabled = false;
				changed++;
			}
		}
	}
	if (changed == 0)
		return 0;

	while ((iface = iface_next(GR_IFACE_TYPE_PORT, iface)) != NULL) {
		struct iface_info_port *port = iface_info_port(iface);
		if (port->port_id != p->port_id)
			vec_add(ports, port);
	}

	ret = worker_graph_reload_all(ports);
	vec_free(ports);

	LOG(INFO, "port %u unplugged", p->port_id);

	return ret;
}

int port_plug(struct iface_info_port *p) {
	struct queue_map *qmap;
	struct worker *worker;
	int changed = 0;

	STAILQ_FOREACH (worker, &workers, next) {
		vec_foreach_ref (qmap, worker->rxqs) {
			if (qmap->port_id == p->port_id) {
				qmap->enabled = true;
				changed++;
			}
		}
		vec_foreach_ref (qmap, worker->txqs) {
			if (qmap->port_id == p->port_id) {
				qmap->enabled = true;
				changed++;
			}
		}
	}
	if (changed == 0)
		return errno_set(ENODEV);

	LOG(INFO, "port %u plugged", p->port_id);

	vec struct iface_info_port **ports = NULL;
	struct iface *iface = NULL;
	bool found = false;
	while ((iface = iface_next(GR_IFACE_TYPE_PORT, iface)) != NULL) {
		struct iface_info_port *port = iface_info_port(iface);
		if (port->port_id == p->port_id)
			found = true;
		vec_add(ports, port);
	}
	if (!found)
		vec_add(ports, p);

	int ret = worker_graph_reload_all(ports);
	vec_free(ports);
	return ret;
}

static void worker_txq_distribute(vec struct iface_info_port **ports) {
	struct iface_info_port *port;
	struct worker *worker;

	// clear all TX queue assignments
	STAILQ_FOREACH (worker, &workers, next)
		vec_free(worker->txqs);

	// assign TX queues only to workers that have RX queues
	vec_foreach (port, ports) {
		uint16_t txq_idx = 0;
		STAILQ_FOREACH (worker, &workers, next) {
			if (vec_len(worker->rxqs) == 0)
				continue;
			assert(port->n_txq > 0);
			struct queue_map txq = {
				.port_id = port->port_id,
				.queue_id = txq_idx % port->n_txq,
				.enabled = port->started,
			};
			vec_add(worker->txqs, txq);
			txq_idx++;
		}
	}
}

int worker_rxq_assign(uint16_t port_id, uint16_t rxq_id, uint16_t cpu_id) {
	struct worker *src_worker, *dst_worker;
	struct queue_map *qmap;
	int ret;

	if (CPU_ISSET(cpu_id, &gr_config.control_cpus))
		return errno_set(EBUSY);

	if (!CPU_ISSET(cpu_id, &gr_config.datapath_cpus))
		return errno_set(ERANGE);

	STAILQ_FOREACH (src_worker, &workers, next) {
		vec_foreach_ref (qmap, src_worker->rxqs) {
			if (qmap->port_id != port_id)
				continue;
			if (qmap->queue_id != rxq_id)
				continue;
			if (src_worker->cpu_id == cpu_id) {
				// rxq already assigned to the correct worker
				return 0;
			}
			goto move;
		}
	}
	return errno_set(ENODEV);
move:
	// prepare destination worker
	dst_worker = worker_find(cpu_id);
	assert(dst_worker != NULL);

	// unassign from src_worker
	for (size_t i = 0; i < vec_len(src_worker->rxqs); i++) {
		struct queue_map *qmap = &src_worker->rxqs[i];
		if (qmap->port_id != port_id)
			continue;
		if (qmap->queue_id != rxq_id)
			continue;
		vec_del_swap(src_worker->rxqs, i);
		break;
	}

	// assign rxq to dst_worker
	struct queue_map rx_qmap = {
		.port_id = port_id,
		.queue_id = rxq_id,
		.enabled = true,
	};
	vec_add(dst_worker->rxqs, rx_qmap);

	// reassign TX queues and reload all graphs
	vec struct iface_info_port **ports = NULL;
	struct iface *iface = NULL;
	while ((iface = iface_next(GR_IFACE_TYPE_PORT, iface)) != NULL)
		vec_add(ports, iface_info_port(iface));

	worker_txq_distribute(ports);

	ret = worker_graph_reload_all(ports);

	vec_free(ports);
	return ret;
}

int worker_queue_distribute(const cpu_set_t *affinity, vec struct iface_info_port **ports) {
	struct iface_info_port *port;
	vec unsigned *cpus = NULL;
	struct worker *worker, *tmp;
	char buf[BUFSIZ];
	int ret = 0;
	unsigned i;

	if (cpuset_format(buf, sizeof(buf), affinity) < 0) {
		LOG(WARNING, "failed to format new cpu affinity: %s", strerror(errno));
		buf[sizeof(buf) - 1] = '\0';
	}

	STAILQ_FOREACH_SAFE (worker, &workers, next, tmp) {
		if (CPU_ISSET(worker->cpu_id, affinity)) {
			// Remove all RXQ/TXQ from that worker to have a clean slate.
			vec_free(worker->rxqs);
			vec_free(worker->txqs);
			if (CPU_COUNT(affinity) != CPU_COUNT(&gr_config.datapath_cpus)) {
				// Affinity was changed and contains a different number of CPUs.
				// Ports will need to be stopped to be reconfigured.
				// Pause worker to make it stop polling RXQs.
				if ((ret = worker_graph_reload(worker, NULL)) < 0) {
					errno_log(errno, "worker_graph_reload");
					goto end;
				}
			}
		} else {
			// This CPU is out of the affinity mask.
			if ((ret = worker_destroy(worker->cpu_id)) < 0) {
				errno_log(errno, "worker_destroy");
				goto end;
			}
		}
	}

	for (unsigned cpu = 0; cpu < CPU_SETSIZE; cpu++) {
		if (CPU_ISSET(cpu, affinity)) {
			worker = worker_find(cpu);
			if (worker == NULL) {
				// Start workers for new CPUs of the mask
				if ((ret = worker_create(cpu)) < 0) {
					errno_log(-ret, "worker_create");
					goto end;
				}
			}
			vec_add(cpus, cpu);
		}
	}

	// assign all rxqs to new workers in the new affinity mask
	i = 0;
	vec_foreach (port, ports) {
		int socket_id = SOCKET_ID_ANY;

		if (numa_available() != -1)
			socket_id = rte_eth_dev_socket_id(port->port_id);

		if (CPU_COUNT(affinity) != CPU_COUNT(&gr_config.datapath_cpus)) {
			// Affinity was changed and contains a different number of CPUs.
			// The number of TXQs must be adjusted accordingly.
			bool was_started = port->started;
			if (port->started && (ret = rte_eth_dev_stop(port->port_id)) < 0) {
				errno_log(-ret, "rte_eth_dev_stop");
				goto end;
			}
			port->started = false;
			if ((ret = port_configure(port, CPU_COUNT(affinity))) < 0) {
				errno_log(-ret, "port_configure");
				goto end;
			}
			if (was_started && (ret = rte_eth_dev_start(port->port_id)) < 0) {
				errno_log(-ret, "rte_eth_dev_start");
				goto end;
			}
			port->started = was_started;
		}

		for (uint16_t rxq = 0; rxq < port->n_rxq; rxq++) {
			// find CPU in the affinity where to assign RXQ
			unsigned j = 0;
			while (socket_id != SOCKET_ID_ANY && socket_id != numa_node_of_cpu(cpus[i])
			       && j < vec_len(cpus)) {
				if (++i >= vec_len(cpus))
					i = 0;
				j++;
			}
			if (j == vec_len(cpus)) {
				LOG(WARNING,
				    "no socket %d CPU found in new affinity %s for port %d",
				    socket_id,
				    buf,
				    port->port_id);
			}

			worker = worker_find(cpus[i]);
			assert(worker != NULL);

			struct queue_map q = {
				.port_id = port->port_id,
				.queue_id = rxq,
				.enabled = port->started,
			};
			vec_add(worker->rxqs, q);

			if (++i >= vec_len(cpus))
				i = 0;
		}
	}

	worker_txq_distribute(ports);

	ret = worker_graph_reload_all(ports);
end:
	vec_free(cpus);
	return ret;
}

static struct gr_stat *find_stat(vec struct gr_stat *stats, const char *name) {
	struct gr_stat *s;

	vec_foreach_ref (s, stats) {
		if (strncmp(s->name, name, sizeof(s->name)) == 0)
			return s;
	}

	return errno_set_null(ENOENT);
}

vec struct gr_stat *worker_dump_stats(uint16_t cpu_id) {
	uint64_t loop_cycles = 0, node_cycles = 0, n_loops = 0, pkts = 0;
	char xname[MEMBER_SIZE(struct gr_stat, name)];
	vec struct gr_stat *stats = NULL;
	const struct gr_node_info *info;
	struct worker *worker;
	struct gr_stat *s;

	STAILQ_FOREACH (worker, &workers, next) {
		const struct worker_stats *w_stats = atomic_load(&worker->stats);
		if (w_stats == NULL)
			continue;
		if (cpu_id != UINT16_MAX && worker->cpu_id != cpu_id)
			continue;
		for (unsigned i = 0; i < w_stats->n_stats; i++) {
			const struct node_stats *n = &w_stats->stats[i];
			const char *name = rte_node_id_to_name(n->node_id);

			info = gr_node_info_get(n->node_id);
			if (info == NULL)
				info = gr_node_info_get(n->parent_id);
			assert(info != NULL);

			s = find_stat(stats, name);
			if (s != NULL) {
				s->packets += n->packets;
				s->batches += n->batches;
				s->cycles += n->cycles;
			} else {
				struct gr_stat stat = {
					.packets = n->packets,
					.batches = n->batches,
					.cycles = n->cycles,
					.topo_order = n->topo_order,
				};
				gr_strcpy(stat.name, sizeof(stat.name), name);
				vec_add(stats, stat);
			}
			for (uint8_t x = 0; x < n->nb_xstats; x++) {
				snprintf(
					xname,
					sizeof(xname),
					"%s.%s",
					name,
					info->node->xstats->xstat_desc[x]
				);
				s = find_stat(stats, xname);
				if (s != NULL) {
					s->packets += n->xstats[x];
					s->batches += n->batches;
				} else {
					struct gr_stat stat = {
						.packets = n->xstats[x],
						.batches = n->batches,
						.topo_order = n->topo_order,
					};
					gr_strcpy(stat.name, sizeof(stat.name), xname);
					vec_add(stats, stat);
				}
			}

			if (strncmp(name, "port_rx-", strlen("port_rx-")) == 0
			    || strcmp(name, "control_input") == 0)
				pkts += n->packets;
			node_cycles += n->cycles;
		}
		s = find_stat(stats, "idle");
		if (s != NULL) {
			s->batches += w_stats->n_sleeps;
			s->cycles += w_stats->sleep_cycles;
		} else {
			struct gr_stat stat = {
				.packets = 0,
				.batches = w_stats->n_sleeps,
				.cycles = w_stats->sleep_cycles,
				.topo_order = UINT64_MAX,
			};
			gr_strcpy(stat.name, sizeof(stat.name), "idle");
			vec_add(stats, stat);
		}
		loop_cycles += w_stats->loop_cycles - w_stats->sleep_cycles;
		n_loops += w_stats->n_loops;
	}

	struct gr_stat stat = {
		.packets = pkts,
		.batches = n_loops,
		.cycles = loop_cycles - node_cycles,
		.topo_order = UINT64_MAX - 1,
	};
	gr_strcpy(stat.name, sizeof(stat.name), "overhead");
	vec_add(stats, stat);

	return stats;
}

static void worker_init(struct event_base *) {
	if (worker_queue_distribute(&gr_config.datapath_cpus, NULL) < 0)
		ABORT("initial worker start failed");
}

static void worker_fini(struct event_base *) {
	struct worker *w, *tmp;

	STAILQ_FOREACH_SAFE (w, &workers, next, tmp)
		worker_destroy(w->cpu_id);

	STAILQ_INIT(&workers);
}

static struct module worker_module = {
	.name = "worker",
	.depends_on = "control_queue",
	.init = worker_init,
	.fini = worker_fini,
};

RTE_INIT(control_infra_init) {
	module_register(&worker_module);
}
