#include <stdio.h>
#include <limits.h>
#include <liburing.h>
#include <netinet/in.h>

#include <osmocom/core/utils.h>
#include <osmocom/core/socket.h>
#include <osmocom/core/sockaddr_str.h>

#define SEND_SLOTS 32000
#define BUF_SIZE 1452

static struct io_uring ring = {};
static unsigned long long num_packets = ULLONG_MAX;
static unsigned long long num_packets_completed = 0;

static int src_fd;
static struct osmo_sockaddr remote_osa;

struct send_slot {
	struct iovec iov;
	uint8_t buf[BUF_SIZE];
	struct msghdr msgh;
};

struct send_slot send_slots[SEND_SLOTS];

static void prepare_tx(struct send_slot *s)
{
	s->iov.iov_base = s->buf;
	s->iov.iov_len = sizeof(s->buf);

	s->msgh = (struct msghdr){
		.msg_name = &remote_osa,
		.msg_namelen = sizeof(remote_osa),
		.msg_iov = &s->iov,
		.msg_iovlen = 1,
	};
}

static void submit_tx(struct send_slot *s)
{
	struct io_uring_sqe *sqe;

	sqe = io_uring_get_sqe(&ring);
	OSMO_ASSERT(sqe);
	io_uring_prep_sendmsg(sqe, src_fd, &s->msgh, 0);
	io_uring_sqe_set_data(sqe, s);
}

static void handle_completion(struct io_uring_cqe *cqe)
{
	struct send_slot *s;
	s = io_uring_cqe_get_data(cqe);

	if (cqe->res < 0) {
		printf("rc = %d\n", cqe->res);
		return;
	}

	io_uring_cqe_seen(&ring, cqe);
	num_packets_completed++;

	/* submit more */
	submit_tx(s);
}

int main(int argc, const char **argv)
{
	int i;
	int rc;
	const char *local_addr_str = "0.0.0.0";
	uint16_t local_port = 42000;
	const char *remote_addr_str = "127.0.0.2";
	uint16_t remote_port = 23000;

	struct osmo_sockaddr_str local_addr = {};
	struct osmo_sockaddr local_osa = {};

	struct osmo_sockaddr_str remote_addr = {};

	struct __kernel_timespec ts_zero = {};
	struct __kernel_timespec ts_1s = { .tv_sec = 1 };

	if (argc >= 2)
		remote_addr_str = argv[1];
	if (argc >= 3)
		remote_port = atoi(argv[2]);
	if (argc >= 4)
		local_port = atoi(argv[3]);
	if (argc >= 5)
		num_packets = atoi(argv[4]);

	if (osmo_sockaddr_str_from_str(&local_addr, local_addr_str, local_port)
	    || osmo_sockaddr_str_to_osa(&local_addr, &local_osa)) {
		printf("ERROR: invalid address or port number: %s:%d\n", local_addr_str, local_port);
		return -1;
	}

	if (osmo_sockaddr_str_from_str(&remote_addr, remote_addr_str, remote_port)
	    || osmo_sockaddr_str_to_osa(&remote_addr, &remote_osa)) {
		printf("ERROR: invalid address or port number: %s:%d\n", remote_addr_str, remote_port);
		return -1;
	}

	/* create and bind socket */
	rc = osmo_sock_init_osa(SOCK_DGRAM, IPPROTO_UDP, &local_osa, NULL, OSMO_SOCK_F_BIND);
	if (rc < 0)
		return -1;
	src_fd = rc;
	printf("bound UDP %s fd=%d\n", osmo_sock_get_name2(src_fd), src_fd);

	/* Set Don't Fragment (DF) bit on IP packets transmitted by socket: */
	int val = IP_PMTUDISC_DO;
	rc = setsockopt(src_fd, IPPROTO_IP, IP_MTU_DISCOVER, &val, sizeof(val));
	if (rc == -1) {
		fprintf(stderr, "ERROR: setsockopt(IPPROTO_IP, IP_DONTFRAG) failed errno=%d\n", errno);
		return -1;
	}

	printf("sending %llu UDP packets to %s\n", num_packets, osmo_sockaddr_to_str(&remote_osa));

	rc = io_uring_queue_init(ARRAY_SIZE(send_slots), &ring, 0);
	if (rc < 0) {
		fprintf(stderr, "ERROR: io_uring_queue_init errno=%d\n", -rc);
		return -1;
	}

	/* Prepare */
	for (i = 0; i < ARRAY_SIZE(send_slots); i++)
		prepare_tx(&send_slots[i]);
	/* fill up tx queue */
	for (i = 0; i < ARRAY_SIZE(send_slots); i++)
		submit_tx(&send_slots[i]);

	while (num_packets_completed < num_packets) {
		uint32_t new_submissions;
		uint32_t new_completions = 0;
		struct io_uring_cqe *cqe;

		/* submit any requests from previous loop */
		new_submissions = io_uring_submit(&ring);

		/* process all pending completions */
		while (io_uring_wait_cqe_timeout(&ring, &cqe, &ts_zero) == 0) {
			handle_completion(cqe);
			new_completions++;

		}

		/* Nothing happened in this loop iteration, so wait a bit longer */
		if (!new_submissions && !new_completions) {
			if (io_uring_wait_cqe_timeout(&ring, &cqe, &ts_1s) == 0)
				handle_completion(cqe);
		}
	}

	printf("done\n");
	return 0;
}