/*
 * Creates a configurable number of UDP sockets with IPv4 and IPv6 addresses.
 *
 * If run as receiver: counts the number of UDP datagrams received.
 *    - binds on all addresses on a given port and connects the socket to the
 *      corresponding IP address and port.
 *    - Afterwars, if epolls and count the corectly received datagrams.
 *    - not knowing when to stop listening it must be killed with CTRL+C (or
 *      other killing signel). On the first CTRL+C received prints the number of
 *      received datagrams; on the second CTRL+C received it exit()s.
 *
 * If run as sender: sends a configurable number of UDP datagrams to the
 *                   receiver and prints the number of second passed.
 *    - binds on all addresses on a given port and connects the socket to the
 *      corresponding IP address and port.
 *    - epolls for a free sending slot and send a new datagram to the receiver.
 *    - at end prints the time it took to send all the datagrams.
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <termios.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/epoll.h>
#include <time.h>
#include <signal.h>

struct timespec tp;

struct epoll_data_t {
	int fd;
	int count;
	int disabled;
};



static void start_timer(struct timespec *tp)
{
	int rc = clock_gettime(CLOCK_REALTIME, tp);
	if (rc != 0) {
		perror("clock_realtime err in start_timer");
	}
}
static void end_timer(struct timespec *pold, const char * msg)
{
	int rc;
	long unsigned dsec, dnsec;
	struct timespec now;
	rc = clock_gettime(CLOCK_REALTIME, &now);
	if (rc != 0) {
		perror("clock_realtime err in end_timer");
	}

	if (now.tv_nsec < pold->tv_nsec) {
		dnsec = pold->tv_nsec - now.tv_nsec;
		dsec = now.tv_sec - pold->tv_sec - 1;
	} else {
		dnsec = now.tv_nsec - pold->tv_nsec;
		dsec = now.tv_sec - pold->tv_sec;
	}
	printf(">>> %s: %lu s  %lu ns \n", msg, dsec, dnsec);
}


static struct in_addr * make_ip4_addrs(int count, int suffix)
{
	int i = 0;
	char str[INET6_ADDRSTRLEN];
	struct in_addr * ret = malloc(sizeof(struct in_addr) * count);
	for(i = 0; i < count; i++) {
		snprintf(str, INET6_ADDRSTRLEN, "1.%d.%d.%d",
			 (i & 0xFF00) >> 8, i & 0xFF, suffix);
		int s = inet_pton(AF_INET, str, &ret[i]);
		if (s <= 0) {
			if (s == 0)
				fprintf(stderr, "ipv4 Not in presentation format");
			else
				perror("inet_pton");
			exit(EXIT_FAILURE);
		}
	}
	return ret;
}

static struct in6_addr * make_ip6_addrs(int count, int suffix)
{
	int i = 0;
	char str[INET6_ADDRSTRLEN];
	struct in6_addr * ret = malloc(sizeof(struct in6_addr) * count);
	for(i = 0; i < count; i++) {
		snprintf(str, INET6_ADDRSTRLEN, "1::%x:%x:%x",
			 (i / 0x10000), (i & 0xFFFF), suffix);
		int s = inet_pton(AF_INET6, str, &ret[i]);
		if (s <= 0) {
			if (s == 0)
				fprintf(stderr, "ipv6 Not in presentation format");
			else
				perror("inet_pton");
			exit(EXIT_FAILURE);
		}
	}
	return ret;
}



static int * make_sockets_(int count, int type)
{
	int i = 0;
	int * ret = malloc(sizeof(int)*count);
	for (i = 0; i < count; i++) {
		ret[i] = socket(type, SOCK_DGRAM, 0);
		if (ret[i] < 0) {
			perror("make_ipv4_sockets socket() failed");
			exit(EXIT_FAILURE);
		}
	}
	return ret;
}

static int * make_ip4_sockets(int count)
{
	return make_sockets_(count, AF_INET);
}

static int * make_ip6_sockets(int count)
{
	return make_sockets_(count, AF_INET6);
}

static void print_ip_err(void * in_addr, int addr_type)
{
	char str[INET6_ADDRSTRLEN];
	if (inet_ntop(addr_type, in_addr, str, INET6_ADDRSTRLEN) == NULL) {
		perror("inet_ntop");
		exit(EXIT_FAILURE);
	}
	fprintf(stderr, "failed for addr: %s\n", str);
}


static void for_all_set_sockopt(int * s, int count, int level, int optname, int val)
{
	int i = 0;
	int rc;
	for (i = 0; i < count; i++) {
		rc = setsockopt(s[i], level, optname, (int[]){val}, sizeof(int));
		if (rc != 0) {
			perror("setsockopt");
			exit(EXIT_FAILURE);
		}
	}
}

static void for_all_set_reuseaddr(int * s, int count)
{
	for_all_set_sockopt(s, count, SOL_SOCKET, SO_REUSEADDR, 1);
}

static void for_all_set_rcvbuf(int * s, int count, int size)
{
	for_all_set_sockopt(s, count, SOL_SOCKET, SO_RCVBUF, size);
}

static void for_all_set_sndbuf(int * s, int count, int size)
{
	for_all_set_sockopt(s, count, SOL_SOCKET, SO_SNDBUF, size);
}


static void for_all_set_nonblocking(int * s, int count)
{
	int rc, i;
	for (i = 0; i < count; i++) {
		rc = fcntl(s[i], F_SETFL, O_NONBLOCK);
		if (rc != 0) {
			fprintf(stderr, "i = %d\n", i);
			perror("fcntl O_NONBLOCK");
			exit(EXIT_FAILURE);
		}
	}

}

static void for_all_set_ipv6only(int * s, int count)
{
	for_all_set_sockopt(s, count, IPPROTO_IPV6, IPV6_V6ONLY, 1);
}





static void for_all_bind_or_connect_ip4(int * s, struct in_addr * ip4, int count, int port,
			     int (*f)(int sockfd, const struct sockaddr *addr,
				      socklen_t addrlen))
{
	int i, rc;
	struct sockaddr_in sin;
	sin.sin_family = AF_INET;
	sin.sin_port = htons(port);

	for (i = 0; i < count; i++) {
		sin.sin_addr = ip4[i];
		rc = f(s[i], (struct sockaddr*)&sin, sizeof(sin));
		if (rc != 0) {
			if (f == bind)
				perror("bind ipv4");
			else
				perror("connect ipv4");
			fprintf(stderr, "i = %d\n", i);
			print_ip_err(&ip4[i], AF_INET);
			exit(EXIT_FAILURE);
		}
	}
}

static void for_all_bind_or_connect_ip6(int * s, struct in6_addr * ip6, int count, int port,
			     int (*f)(int sockfd, const struct sockaddr *addr,
				      socklen_t addrlen))
{
	int i, rc;
	struct sockaddr_in6 sin6;
	sin6.sin6_family = AF_INET6;
	sin6.sin6_port = htons(port);

	for (i = 0; i < count; i++) {
		sin6.sin6_addr = ip6[i];
		rc = f(s[i], (struct sockaddr*)&sin6, sizeof(sin6));
		if (rc != 0) {
			if (f == bind)
				perror("bind ipv6");
			else
				perror("connect ipv6");
			fprintf(stderr, "i = %d\n", i);
			print_ip_err(&ip6[i], AF_INET6);
			exit(EXIT_FAILURE);
		}
	}
}

static void for_all_bind_ip4(int * s, struct in_addr * ip4, int count, int port)
{
	for_all_bind_or_connect_ip4(s, ip4, count, port, bind);
}
static void for_all_bind_ip6(int * s, struct in6_addr * ip6, int count, int port)
{
	for_all_bind_or_connect_ip6(s, ip6, count, port, bind);
}
static void for_all_connect_ip4(int * s, struct in_addr * ip4, int count, int port)
{
	for_all_bind_or_connect_ip4(s, ip4, count, port, connect);
}
static void for_all_connect_ip6(int * s, struct in6_addr * ip6, int count, int port)
{
	for_all_bind_or_connect_ip6(s, ip6, count, port, connect);
}



static struct epoll_data_t * for_all_add_to_epoll(int *s, int count, int epollfd, int is_sender)
{
	int rc;
	int i;
	struct epoll_event ev;
	struct epoll_data_t * data = calloc(count, sizeof(struct epoll_data_t));
	ev.events = EPOLLET;
	if (is_sender)
		ev.events |= EPOLLOUT;
	else
		ev.events |= EPOLLIN;
	for (i = 0; i < count; i++) {
		data[i].fd = s[i];
		ev.data.ptr = &data[i];
		rc = epoll_ctl(epollfd, EPOLL_CTL_ADD, s[i], &ev);
		if (rc == -1) {
			fprintf(stderr, "i = %d\n", i);
			perror("epoll_ctl: conn_sock");
			exit(EXIT_FAILURE);
		}
	}
	return data;
}

static void remove_from_epoll(int epollfd, int fd, struct epoll_event * ev)
{
	int rc;
	//printf("removing from epoll fd %d\n", fd);
	rc = epoll_ctl(epollfd, EPOLL_CTL_DEL, fd, ev);
	if (rc != 0)
		perror("epoll_ctl EPOLL_CTL_DEL");
}

static void disable_epoll_event_type(int epollfd, int fd, int event_type, struct epoll_event * ev)
{
	int rc;
	//printf("disabling event %d on fd=%d\n", event_type, fd);
	ev->events &= ~event_type;
	rc = epoll_ctl(epollfd, EPOLL_CTL_MOD, fd, ev);
	if (rc != 0)
		perror("epoll_ctl EPOLL_CTL_MOD");
}

static int handle_events(int epollfd, int fds_in_epoll, int is_sender, int max_packets, int datalen)
{
	char * databuf = malloc(datalen);
	struct epoll_event * events = malloc(sizeof(struct epoll_event) * fds_in_epoll);
	int todo = fds_in_epoll;
	int total = 0;

	for (;todo;) {
		int i, nfds;
		nfds = epoll_wait(epollfd, events, fds_in_epoll, -1);
		if (nfds == -1) {
			if (errno == EINTR) {
				printf("signal interrupted epoll_wait: total= %d\n", total);
				continue;
			}
			perror("epoll_wait");
			exit(EXIT_FAILURE);
		}
		for (i = 0; i < nfds; i++) {
			int fd, rc;
			struct epoll_event  *ev = &events[i];
			struct epoll_data_t *ed = ev->data.ptr;
			fd = ed->fd;
			if (ev->events & (EPOLLERR | EPOLLHUP)) {
				remove_from_epoll(epollfd, fd, ev);
				todo--;
				continue;
			}
			if (ev->events & EPOLLIN) {
				rc = recv(fd, databuf, datalen, 0);
				if (rc == -1) {
					//if (errno != ECONNREFUSED)
						perror("recv");
					continue;
				}
			}
			if (ev->events & EPOLLOUT) {
				rc = send(fd, databuf, datalen, 0);
				if (rc == -1) {
					//if (errno != ECONNREFUSED)
						perror("send");
					continue;
				}
			}
			if (total == 0)
				start_timer(&tp);
			total ++;
			ed->count ++;
			if (ed->count == max_packets) {
				remove_from_epoll(epollfd, fd, ev);
				todo--;
			}
		}
	}
	free(events);
	return total;
}

static int signal_was_delivered = 0;
static void handler(int unused)
{
	if (!signal_was_delivered)
		signal_was_delivered = 1;
	else
		exit(0);
}
static void add_signal()
{
	signal(SIGINT, handler);
}

#define stoi atoi
int main(int argc, char * argv[])
{
	int port = 1212;

	if (argc < 8) {
		printf("Usage: %s nr_interfaces4 nr_interfaces6 my_last_suffix their_last_suffix [s|r] max_packets datalen [recvsize]\n", argv[0]);
		exit(EXIT_FAILURE);
	}

	int count4       = stoi(argv[1]);
	int count6       = stoi(argv[2]);
	int my_suffix    = stoi(argv[3]);
	int their_suffix = stoi(argv[4]);
	int is_sender    = argv[5][0] == 's';
	int max_packets  = stoi(argv[6]);
	int datalen      = stoi(argv[7]);
	int recvsize = -1;
	if (argc >= 8)
		recvsize = stoi(argv[8]);
	add_signal();
	printf("args = count4=%d count6=%d my_suffix=%d their_suffix=%d is_sender=%d max_packets=%d datalen=%d recvsize=%d\n",
	       count4, count6, my_suffix, their_suffix, is_sender, max_packets, datalen, recvsize);
	int epollfd;
	struct in_addr  * ip4_src = make_ip4_addrs(count4, my_suffix);
	struct in6_addr * ip6_src = make_ip6_addrs(count6, my_suffix);

	struct in_addr  * ip4_dst = make_ip4_addrs(count4, their_suffix);
	struct in6_addr * ip6_dst = make_ip6_addrs(count6, their_suffix);

	struct epoll_data_t *ep_data_ip4, *ep_data_ip6;

	int * s4 = make_ip4_sockets(count4);
	int * s6 = make_ip6_sockets(count6);
	int total = 0;

	for_all_set_reuseaddr(s4, count4);
	for_all_set_reuseaddr(s6, count6);

	if (recvsize != -1) {
		for_all_set_rcvbuf(s4, count4, recvsize);
		for_all_set_rcvbuf(s6, count6, recvsize);
	}
	//for_all_set_sndbuf(s4, count4, 1);
	//for_all_set_sndbuf(s6, count6, 1);

	start_timer(&tp);
	for_all_bind_ip6(s6, ip6_src, count6, port);
	for_all_bind_ip4(s4, ip4_src, count4, port);
	end_timer(&tp, "bound     IPv4 and IPv6");

	start_timer(&tp);
	for_all_connect_ip4(s4, ip4_dst, count4, port);
	for_all_connect_ip6(s6, ip6_dst, count6, port);
	end_timer(&tp, "connected IPv4 and IPv6");

	free(ip4_src); free(ip6_src);
	free(ip4_dst); free(ip6_dst);

	for_all_set_nonblocking(s4, count4);
	for_all_set_nonblocking(s6, count6);
	printf("done builing sockets.\n");


	epollfd = epoll_create(count4+count6);
	if (epollfd == -1) {
		perror("epoll_create");
		exit(EXIT_FAILURE);
	}
	ep_data_ip4 = for_all_add_to_epoll(s4, count4, epollfd, is_sender);
	ep_data_ip6 = for_all_add_to_epoll(s6, count6, epollfd, is_sender);
	printf("done builing epoll.\n");



	total = handle_events(epollfd, count4 + count6, is_sender, max_packets, datalen);
	end_timer(&tp, "handled all events");

	free(ep_data_ip4); free(ep_data_ip6);
	free(s4); free(s6);
	printf("done total= %d\n", total);
	return 0;
}