/*
 * pixelflut-bench - UDP send-throughput benchmark for the Pixelflut renderer.
 *
 * It measures how fast we can actually push Pixelflut datagrams out of the
 * kernel, which is the number we care about: "it is a waste to not be sending
 * packets". Each sender thread owns its own socket and a pre-built ring of
 * packets, then loops on sendmmsg() as hard as it can. A monitor prints the
 * achieved packet rate and payload throughput once per second.
 *
 * By default it runs a self-contained loopback test: an in-process sink binds
 * the target port and drains packets so the senders never block, letting us
 * measure the pure generate+syscall ceiling of this machine. Point it at a
 * real server with -H to benchmark an actual link.
 *
 *   ./pixelflut-bench                 # loopback self-test, 1 thread, 5 s
 *   ./pixelflut-bench -t 8 -d 10      # 8 sender threads for 10 s
 *   ./pixelflut-bench -H fcda:f100::d -p 5005 -t 8   # hammer a real display
 */

#define _GNU_SOURCE /* sendmmsg(), recvmmsg(), struct mmsghdr */

#include <errno.h>
#include <inttypes.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#include <arpa/inet.h>
#include <netdb.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/types.h>

#include <pthread.h>

#include "pixelflut.h"

/* Portable host -> little-endian 16-bit conversion (matches the wire format). */
#if defined(__linux__)
#  include <endian.h>
#elif defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__)
#  include <sys/endian.h>
#elif defined(__APPLE__)
#  include <libkern/OSByteOrder.h>
#  define htole16(v) OSSwapHostToLittleInt16(v)
#endif

/* Frame geometry mirrored from the renderer, so -g reproduces its real
 * per-pixel generation cost (xorshf96 + divisions + frame lookups). */
#define FRAME_WIDTH     640
#define FRAME_HEIGHT    360
#define FRAME_CHANNELS  3
#define FRAME_PIXELS    (FRAME_WIDTH * FRAME_HEIGHT)
#define FRAME_BYTES     (FRAME_PIXELS * FRAME_CHANNELS)
#define CANVAS_OFFSET_X 1200
#define CANVAS_OFFSET_Y 40

#define PAYLOAD_BYTES ((uint64_t)sizeof(struct pixelflut_packet))

/* Rough per-packet wire overhead (IPv6 + UDP headers) for a link estimate. */
#define WIRE_OVERHEAD_BYTES 48

/* sendmmsg()/recvmmsg() accept at most this many messages per call. */
#define MAX_BATCH 1024

/* Per-thread counters, each on its own cache line to avoid false sharing. */
typedef struct {
	_Alignas(64) _Atomic uint64_t packets;
	_Atomic uint64_t bytes;
	_Atomic uint64_t calls;
	_Atomic uint64_t errors;
} counters_t;

struct config {
	struct sockaddr_storage addr;
	socklen_t addrlen;
	int family;
	int batch;          /* messages per sendmmsg() call and ring size */
	int sndbuf;         /* SO_SNDBUF, 0 = leave kernel default */
	int rcvbuf;         /* SO_RCVBUF for the sink */
	int port;
	bool generate;      /* regenerate pixel payload each pass (like the renderer) */
	const uint8_t *frame; /* shared read-only source frame for -g */
};

struct sender_arg {
	const struct config *cfg;
	counters_t *ctr;
	int index;
};

static _Atomic bool g_running = true;

static double now_sec(void)
{
	struct timespec ts;
	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (double)ts.tv_sec + (double)ts.tv_nsec / 1e9;
}

static void set_bufsize(int fd, int optname, int size, const char *what)
{
	if (size > 0 && setsockopt(fd, SOL_SOCKET, optname, &size, sizeof size) == -1)
		fprintf(stderr, "warning: setsockopt(%s, %d) failed: %s\n", what,
		        size, strerror(errno));
}

/* Fill a ring of packets with a valid header and arbitrary pixel data. */
static void build_packets(struct pixelflut_packet *pkts, int n)
{
	for (int i = 0; i < n; i++) {
		pkts[i].cmd = PIXELFLUT_CMD_SET;
		pkts[i].subcmd = PIXELFLUT_SUBCMD_RGB;
		for (int j = 0; j < PIXELS_PER_PACKET; j++) {
			struct pixelflut_pixel *px = &pkts[i].pixels[j];
			px->x = (uint16_t)(i + j);
			px->y = (uint16_t)(i - j);
			px->r = (uint8_t)i;
			px->g = (uint8_t)j;
			px->b = 0x80;
		}
	}
}

/* xorshf96 PRNG, per-thread state (same generator the renderer uses). */
static unsigned long xorshf96(unsigned long *x, unsigned long *y, unsigned long *z)
{
	*x ^= *x << 16;
	*x ^= *x >> 5;
	*x ^= *x << 1;
	unsigned long t = *x;
	*x = *y;
	*y = *z;
	*z = t ^ *x ^ *y;
	return *z;
}

/* Repaint a ring of packets from random frame pixels, exactly like the
 * renderer's hot loop, to measure generate+send throughput under -g. */
static void generate_ring(struct pixelflut_packet *pkts, int n,
                          const uint8_t *frame, unsigned long *x,
                          unsigned long *y, unsigned long *z)
{
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < PIXELS_PER_PACKET; j++) {
			uint64_t src = xorshf96(x, y, z) % FRAME_PIXELS;
			uint16_t fx = (uint16_t)(src % FRAME_WIDTH);
			uint16_t fy = (uint16_t)(src / FRAME_WIDTH);
			const uint8_t *rgb = &frame[src * FRAME_CHANNELS];

			struct pixelflut_pixel *px = &pkts[i].pixels[j];
			px->x = htole16((uint16_t)(CANVAS_OFFSET_X + fx));
			px->y = htole16((uint16_t)(CANVAS_OFFSET_Y + fy));
			px->r = rgb[0];
			px->g = rgb[1];
			px->b = rgb[2];
		}
	}
}

static void *sender_main(void *arg)
{
	const struct sender_arg *sa = arg;
	const struct config *cfg = sa->cfg;
	counters_t *ctr = sa->ctr;
	const int n = cfg->batch;

	/* Distinct PRNG seeds per thread so they paint different pixels. */
	unsigned long rng_x = 123456789UL ^ (2654435761UL * (unsigned long)(sa->index + 1));
	unsigned long rng_y = 362436069UL + (unsigned long)(sa->index + 1);
	unsigned long rng_z = 521288629UL;

	int fd = socket(cfg->family, SOCK_DGRAM, 0);
	if (fd == -1) {
		perror("socket");
		return NULL;
	}
	set_bufsize(fd, SO_SNDBUF, cfg->sndbuf, "SO_SNDBUF");
	if (connect(fd, (const struct sockaddr *)&cfg->addr, cfg->addrlen) == -1) {
		perror("connect");
		close(fd);
		return NULL;
	}

	struct pixelflut_packet *pkts = malloc((size_t)n * sizeof *pkts);
	struct mmsghdr *msgs = calloc((size_t)n, sizeof *msgs);
	struct iovec *iov = calloc((size_t)n, sizeof *iov);
	if (pkts == NULL || msgs == NULL || iov == NULL) {
		fprintf(stderr, "sender: out of memory\n");
		goto out;
	}
	build_packets(pkts, n);
	for (int i = 0; i < n; i++) {
		iov[i].iov_base = &pkts[i];
		iov[i].iov_len = sizeof pkts[i];
		msgs[i].msg_hdr.msg_iov = &iov[i];
		msgs[i].msg_hdr.msg_iovlen = 1;
	}

	uint64_t packets = 0, bytes = 0, calls = 0, errors = 0;
	int next = 0;
	while (atomic_load_explicit(&g_running, memory_order_relaxed)) {
		/* Repaint the whole ring before each pass, like the renderer. */
		if (cfg->generate && next == 0)
			generate_ring(pkts, n, cfg->frame, &rng_x, &rng_y, &rng_z);

		int sent = sendmmsg(fd, &msgs[next], (unsigned int)(n - next), 0);
		if (sent < 0) {
			if (errno == EINTR)
				continue;
			errors++;
			next = 0;
			continue;
		}
		packets += (uint64_t)sent;
		bytes += (uint64_t)sent * PAYLOAD_BYTES;
		calls++;
		next += sent;
		if (next >= n)
			next = 0;

		/* Publish cumulative counts for the monitor (cheap, per call). */
		atomic_store_explicit(&ctr->packets, packets, memory_order_relaxed);
		atomic_store_explicit(&ctr->bytes, bytes, memory_order_relaxed);
		atomic_store_explicit(&ctr->calls, calls, memory_order_relaxed);
		atomic_store_explicit(&ctr->errors, errors, memory_order_relaxed);
	}

out:
	free(pkts);
	free(msgs);
	free(iov);
	close(fd);
	return NULL;
}

/* Loopback sink: bind the target port and drain packets so senders never
 * block on a full socket buffer. */
static void *sink_main(void *arg)
{
	const struct sender_arg *sa = arg;
	const struct config *cfg = sa->cfg;
	counters_t *ctr = sa->ctr;
	const int n = (cfg->batch < MAX_BATCH) ? cfg->batch : MAX_BATCH;

	int fd = socket(cfg->family, SOCK_DGRAM, 0);
	if (fd == -1) {
		perror("sink socket");
		return NULL;
	}
	int one = 1;
	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof one);
	set_bufsize(fd, SO_RCVBUF, cfg->rcvbuf, "SO_RCVBUF");
	/* Wake up periodically so we can notice g_running going false. */
	struct timeval tv = {.tv_sec = 0, .tv_usec = 200000};
	setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof tv);
	if (bind(fd, (const struct sockaddr *)&cfg->addr, cfg->addrlen) == -1) {
		perror("sink bind");
		close(fd);
		return NULL;
	}

	const size_t slot = 2048;
	uint8_t *bufs = malloc((size_t)n * slot);
	struct mmsghdr *msgs = calloc((size_t)n, sizeof *msgs);
	struct iovec *iov = calloc((size_t)n, sizeof *iov);
	if (bufs == NULL || msgs == NULL || iov == NULL) {
		fprintf(stderr, "sink: out of memory\n");
		goto out;
	}
	for (int i = 0; i < n; i++) {
		iov[i].iov_base = bufs + (size_t)i * slot;
		iov[i].iov_len = slot;
		msgs[i].msg_hdr.msg_iov = &iov[i];
		msgs[i].msg_hdr.msg_iovlen = 1;
	}

	uint64_t packets = 0, bytes = 0;
	while (atomic_load_explicit(&g_running, memory_order_relaxed)) {
		int got = recvmmsg(fd, msgs, (unsigned int)n, 0, NULL);
		if (got < 0) {
			if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
				continue;
			perror("recvmmsg");
			break;
		}
		packets += (uint64_t)got;
		for (int i = 0; i < got; i++)
			bytes += msgs[i].msg_len;
		atomic_store_explicit(&ctr->packets, packets, memory_order_relaxed);
		atomic_store_explicit(&ctr->bytes, bytes, memory_order_relaxed);
	}

out:
	free(bufs);
	free(msgs);
	free(iov);
	close(fd);
	return NULL;
}

static uint64_t sum_packets(const counters_t *c, int n)
{
	uint64_t t = 0;
	for (int i = 0; i < n; i++)
		t += atomic_load_explicit(&c[i].packets, memory_order_relaxed);
	return t;
}

static uint64_t sum_calls(const counters_t *c, int n)
{
	uint64_t t = 0;
	for (int i = 0; i < n; i++)
		t += atomic_load_explicit(&c[i].calls, memory_order_relaxed);
	return t;
}

static uint64_t sum_errors(const counters_t *c, int n)
{
	uint64_t t = 0;
	for (int i = 0; i < n; i++)
		t += atomic_load_explicit(&c[i].errors, memory_order_relaxed);
	return t;
}

static void print_rate(const char *label, uint64_t packets, double secs)
{
	double pps = (double)packets / secs;
	double gbit = (double)packets * (double)PAYLOAD_BYTES * 8.0 / secs / 1e9;
	double wire = (double)packets * (double)(PAYLOAD_BYTES + WIRE_OVERHEAD_BYTES) *
	              8.0 / secs / 1e9;
	printf("%-8s %8.3f Mpps  %8.3f Gbit/s payload  %8.3f Gbit/s wire\n", label,
	       pps / 1e6, gbit, wire);
}

static int resolve(const char *host, int port, bool passive, struct config *cfg)
{
	char portstr[16];
	snprintf(portstr, sizeof portstr, "%d", port);
	struct addrinfo hints;
	memset(&hints, 0, sizeof hints);
	hints.ai_family = AF_UNSPEC;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_flags = passive ? AI_PASSIVE : 0;

	struct addrinfo *res = NULL;
	int rc = getaddrinfo(host, portstr, &hints, &res);
	if (rc != 0) {
		fprintf(stderr, "getaddrinfo(%s): %s\n", host ? host : "*",
		        gai_strerror(rc));
		return -1;
	}
	memcpy(&cfg->addr, res->ai_addr, res->ai_addrlen);
	cfg->addrlen = res->ai_addrlen;
	cfg->family = res->ai_family;
	freeaddrinfo(res);
	return 0;
}

static void usage(const char *prog)
{
	fprintf(stderr,
	        "usage: %s [-H host] [-p port] [-t threads] [-d secs] "
	        "[-s batch] [-B sndbuf] [-R rcvbuf] [-g]\n"
	        "  -H host    target host (default: loopback self-test with sink)\n"
	        "  -p port    target UDP port (default 5005)\n"
	        "  -t threads sender threads (default 1)\n"
	        "  -d secs    measurement duration (default 5)\n"
	        "  -s batch   messages per sendmmsg() call, 1..%d (default %d)\n"
	        "  -B bytes   SO_SNDBUF per sender socket (default: kernel)\n"
	        "  -R bytes   SO_RCVBUF for the loopback sink (default 4 MiB)\n"
	        "  -g         regenerate pixel payload each pass (renderer workload)\n",
	        prog, MAX_BATCH, MAX_BATCH);
}

int main(int argc, char *argv[])
{
	const char *host = NULL;
	int port = 5005;
	int threads = 1;
	int duration = 5;
	struct config cfg = {.batch = MAX_BATCH, .sndbuf = 0, .rcvbuf = 4 << 20};

	int opt;
	while ((opt = getopt(argc, argv, "H:p:t:d:s:B:R:gh")) != -1) {
		switch (opt) {
		case 'H': host = optarg; break;
		case 'p': port = atoi(optarg); break;
		case 't': threads = atoi(optarg); break;
		case 'd': duration = atoi(optarg); break;
		case 's': cfg.batch = atoi(optarg); break;
		case 'B': cfg.sndbuf = atoi(optarg); break;
		case 'R': cfg.rcvbuf = atoi(optarg); break;
		case 'g': cfg.generate = true; break;
		case 'h': usage(argv[0]); return EXIT_SUCCESS;
		default: usage(argv[0]); return EXIT_FAILURE;
		}
	}
	if (threads < 1 || duration < 1 || cfg.batch < 1 || cfg.batch > MAX_BATCH) {
		usage(argv[0]);
		return EXIT_FAILURE;
	}
	cfg.port = port;

	bool loopback = (host == NULL);
	const char *target = loopback ? "::1" : host;
	if (resolve(target, port, false, &cfg) != 0)
		return EXIT_FAILURE;

	char addrbuf[INET6_ADDRSTRLEN] = "?";
	void *sin_addr =
	    cfg.family == AF_INET6
	        ? (void *)&((struct sockaddr_in6 *)&cfg.addr)->sin6_addr
	        : (void *)&((struct sockaddr_in *)&cfg.addr)->sin_addr;
	inet_ntop(cfg.family, sin_addr, addrbuf, sizeof addrbuf);

	printf("target      %s port %d (%s)\n", addrbuf, port,
	       loopback ? "loopback self-test" : "external");
	printf("threads     %d\n", threads);
	printf("batch       %d msgs/sendmmsg\n", cfg.batch);
	printf("packet      %" PRIu64 " bytes payload (%d pixels)\n", PAYLOAD_BYTES,
	       PIXELS_PER_PACKET);
	printf("workload    %s\n", cfg.generate ? "generate + send" : "send only");
	printf("duration    %d s\n\n", duration);

	/* Shared read-only source frame for -g (contents are irrelevant). */
	uint8_t *frame = NULL;
	if (cfg.generate) {
		frame = malloc(FRAME_BYTES);
		if (frame == NULL) {
			fprintf(stderr, "out of memory\n");
			return EXIT_FAILURE;
		}
		memset(frame, 0x7f, FRAME_BYTES);
		cfg.frame = frame;
	}

	/* Optional loopback sink that drains what we send. */
	counters_t *rctr = NULL;
	pthread_t sink_thread = 0;
	struct sender_arg sink_arg;
	if (loopback) {
		rctr = aligned_alloc(64, sizeof *rctr);
		if (rctr == NULL) {
			fprintf(stderr, "out of memory\n");
			return EXIT_FAILURE;
		}
		memset(rctr, 0, sizeof *rctr);
		sink_arg = (struct sender_arg){.cfg = &cfg, .ctr = rctr};
		if (pthread_create(&sink_thread, NULL, sink_main, &sink_arg) != 0) {
			perror("pthread_create sink");
			return EXIT_FAILURE;
		}
		usleep(100000); /* let the sink bind before senders connect */
	}

	counters_t *sctr = aligned_alloc(64, (size_t)threads * sizeof *sctr);
	struct sender_arg *sargs = calloc((size_t)threads, sizeof *sargs);
	pthread_t *tids = calloc((size_t)threads, sizeof *tids);
	if (sctr == NULL || sargs == NULL || tids == NULL) {
		fprintf(stderr, "out of memory\n");
		return EXIT_FAILURE;
	}
	memset(sctr, 0, (size_t)threads * sizeof *sctr);
	for (int i = 0; i < threads; i++) {
		sargs[i] = (struct sender_arg){.cfg = &cfg, .ctr = &sctr[i], .index = i};
		if (pthread_create(&tids[i], NULL, sender_main, &sargs[i]) != 0) {
			perror("pthread_create sender");
			return EXIT_FAILURE;
		}
	}

	/* Monitor: sample once per second and print the delta. */
	double t0 = now_sec();
	uint64_t prev_pkts = 0;
	double prev_t = t0;
	for (int s = 1; s <= duration; s++) {
		struct timespec req = {.tv_sec = 1, .tv_nsec = 0};
		nanosleep(&req, NULL);
		double t = now_sec();
		uint64_t pkts = sum_packets(sctr, threads);
		char label[16];
		snprintf(label, sizeof label, "[%2ds]", s);
		print_rate(label, pkts - prev_pkts, t - prev_t);
		prev_pkts = pkts;
		prev_t = t;
	}

	atomic_store_explicit(&g_running, false, memory_order_relaxed);
	for (int i = 0; i < threads; i++)
		pthread_join(tids[i], NULL);

	double elapsed = now_sec() - t0;
	uint64_t tx_pkts = sum_packets(sctr, threads);
	uint64_t tx_calls = sum_calls(sctr, threads);
	uint64_t tx_errs = sum_errors(sctr, threads);

	printf("\n=== totals over %.2f s ===\n", elapsed);
	print_rate("send", tx_pkts, elapsed);
	printf("send     %" PRIu64 " packets, %" PRIu64 " sendmmsg calls "
	       "(avg %.1f msgs/call), %" PRIu64 " errors\n",
	       tx_pkts, tx_calls, tx_calls ? (double)tx_pkts / (double)tx_calls : 0.0,
	       tx_errs);

	if (loopback) {
		uint64_t rx_pkts = atomic_load(&rctr->packets);
		print_rate("recv", rx_pkts, elapsed);
		if (tx_pkts > 0)
			printf("recv     %" PRIu64 " packets (%.1f%% of sent reached the "
			       "sink)\n",
			       rx_pkts, 100.0 * (double)rx_pkts / (double)tx_pkts);
		pthread_join(sink_thread, NULL);
		free(rctr);
	}

	free(sctr);
	free(sargs);
	free(tids);
	free(frame);
	return EXIT_SUCCESS;
}