/* server_tcp.c
 *
 * build:
 * epoll: gcc server_tcp.c -o server_tcp.out -DVMA_DEV="ens1f0" -DVMA_API=0 -I/usr/include
 * xtreme: gcc server_tcp.c -o server_tcp.out -DVMA_DEV="ens1f0" -DVMA_API=1 -I<path to mellanox/vma_extra.h>
 *
 * usage:
 * epoll: sudo server_tcp.out 1.1.3.15:17000
 * socketxtreme: sudo env LD_PRELOAD=libvma.so server_tcp.out 1.1.3.15:17000
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <getopt.h>
#include <assert.h>
#include <sys/select.h>
#include <sys/mman.h> /* mlock */
#include <sys/time.h>
#include <time.h>
#include <signal.h>
#include <sys/epoll.h>

#if defined(VMA_API) && (VMA_API == 1)
#include <mellanox/vma_extra.h>
#endif /* VMA_API */

/* Bind to device */
#if !defined(VMA_DEV)
#define IB_DEV "ens3f1"
#else
#define QUOTE(name) #name
#define STR(macro) QUOTE(macro)
#define IB_DEV STR(VMA_DEV)
#endif

/* Number of listeners */
#define SFD_NUM 2

/* Number of peers */
#define FD_NUM 10

#define EXIT_FAILURE 1

#if defined(VMA_API) && (VMA_API == 1)
static struct vma_api_t *_vma_api = NULL;
static int _vma_ring_fd = -1;
#endif /* VMA_API */

static volatile int _done = 0;

static inline char *_addr2str(struct sockaddr_in *addr) {
	static __thread char addrbuf[100];
	inet_ntop(AF_INET, &addr->sin_addr, addrbuf, sizeof(addrbuf));
	sprintf(addrbuf, "%s:%d", addrbuf, ntohs(addr->sin_port));

	return addrbuf;
}

static void _proc_signal(int signal_id)
{
   _done = signal_id;
}

static int _set_noblock(int fd)
{
	int rc = 0;
	int flag;

	flag = fcntl(fd, F_GETFL);
	if (flag < 0) {
		rc = -errno;
		printf("failed to get socket flags %s\n", strerror(errno));
	}
	flag |= O_NONBLOCK;
	rc = fcntl(fd, F_SETFL, flag);
	if (rc < 0) {
		rc = -errno;
		printf("failed to set socket flags %s\n", strerror(errno));
	}

	return rc;
}

static int _tcp_create_and_bind(struct sockaddr_in *addr)
{
	int rc = 0;
	int fd;
	int flag;

	fd = socket(PF_INET, SOCK_STREAM, IPPROTO_IP);
	if (!fd) {
		rc = -EBUSY;
		printf("Failed to create socket\n");
		goto err;
	}

#if defined(IB_DEV)
	rc = setsockopt(fd, SOL_SOCKET, SO_BINDTODEVICE, (void *)IB_DEV, strlen(IB_DEV));
    if (rc < 0) {
            printf("Failed to setsockopt(SO_BINDTODEVICE) for %s: %s\n", IB_DEV, strerror(errno));
            exit(1);
    }
#endif

	flag = 1;
	rc = setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (char *) &flag, sizeof(int));
	if (rc < 0) {
		printf("Failed to setsockopt(SO_REUSEADDR): %s\n", strerror(errno));
		goto err;
	}

	rc = bind(fd, (struct sockaddr *) addr, sizeof(*addr));
	if (rc < 0) {
		rc = -EBUSY;
		printf("Failed to bind socket\n");
		goto err;
	}

	listen(fd, SOMAXCONN);

	printf("Listen  : fd=%d %s\n", fd, _addr2str((struct sockaddr_in *)addr));

err:
	return (rc == 0 ? fd : (-1));
}

int main(int argc, char *argv[])
{
	struct sigaction sa;
	int ret = 0;
	int efd;
	int sfd[SFD_NUM];
	int fd = -1;
	int max_events = 0;
	int max_sfd = 0;
	struct epoll_event ev;
	uint64_t event;
	struct epoll_event *events = NULL;
	struct conn_info {
		int *fds;
		int count;
		char msg[1024];
	} conns;
#if defined(VMA_API) && (VMA_API == 1)
	struct vma_completion_t *vma_comps;
#endif /* VMA_API */
	int flag;
	struct sockaddr_in addr;
	struct sockaddr in_addr;
	socklen_t in_len;
	int i = 0;
	int j = 0;

	/* catch SIGINT to exit */
	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = _proc_signal;
	sa.sa_flags = 0;
	sigemptyset(&(sa.sa_mask));
	if (sigaction(SIGINT, &sa, NULL) != 0) {
        	perror("Failed to create signal handler");
        	exit(EXIT_FAILURE);
	}

	/* Step:1 Initialize VMA API */
#if defined(VMA_API) && (VMA_API == 1)
	_vma_api = vma_get_api();
	if (_vma_api == NULL) {
		printf("VMA Extra API not found\n");
	}
#endif /* VMA_API */

	max_events = FD_NUM + sizeof(sfd) / sizeof(sfd[0]);

	conns.count = 0;
	conns.fds = calloc(max_events, sizeof(*conns.fds));
	assert(conns.fds);

#if defined(VMA_API) && (VMA_API == 1)
	vma_comps = calloc(max_events, sizeof(*vma_comps));
	assert(vma_comps);
#else
	efd = epoll_create1(0);
	assert(efd >= 0);

	events = calloc(max_events, sizeof(*events));
	assert(events);
#endif /* VMA_API */

	printf("Launching <receiver> mode...\n");

	/* Step:2 Create listen socket */
	for (i = 0; (i < SFD_NUM) && (argc > (i + 1)); i++) {
		char *optarg = argv[i + 1];
		char *token1 = NULL;
		char *token2 = NULL;
		const char s[2] = ":";

		token1 = strtok(optarg, s);
		token2 = strtok(NULL, s);

		memset(&addr, 0, sizeof(addr));
		addr.sin_family = PF_INET;
		addr.sin_addr.s_addr = inet_addr(token1);
		addr.sin_port = htons(atoi(token2));
		sfd[i] = _tcp_create_and_bind(&addr);
		if (sfd[i] < 0) {
			perror("Failed to create socket");
			exit(EXIT_FAILURE);
		}
		max_sfd++;
	}

	/* Step:3 Need to get ring or set listen socket */
#if defined(VMA_API) && (VMA_API == 1)
	if (_vma_ring_fd < 0) {
		_vma_api->get_socket_rings_fds(sfd[0], &_vma_ring_fd, 1);
		assert((-1) != _vma_ring_fd);
	}
#else
	for (i = 0; i < max_sfd; i++) {
		ev.events = EPOLLIN;
		ev.data.fd = sfd[i];
		if (epoll_ctl(efd, EPOLL_CTL_ADD, sfd[i], &ev) == -1) {
				perror("epoll_ctl() failed");
			exit(EXIT_FAILURE);
		}
	}
#endif /* VMA_API */

	while (!_done) {
		int n = 0;

		/* Step:4 Get events */
#if defined(VMA_API) && (VMA_API == 1)
		while (0 == n) {
			n = _vma_api->socketxtreme_poll(_vma_ring_fd, vma_comps, max_events, 0);
		}
#else
		n = epoll_wait(efd, events, max_events, 0);
#endif /* VMA_API */
		for (j = 0; j < n; j++) {

#if defined(VMA_API) && (VMA_API == 1)
			event = vma_comps[j].events;
			event |= ( event & VMA_SOCKETXTREME_PACKET ? EPOLLIN : 0);
			fd = (event & VMA_SOCKETXTREME_NEW_CONNECTION_ACCEPTED ? vma_comps[j].listen_fd : vma_comps[j].user_data);
#else
			event = events[j].events;
			fd = events[j].data.fd;
#endif /* VMA_API */

			if ((event & EPOLLERR) || (event & EPOLLHUP) || (event & EPOLLRDHUP)) {
				printf("epoll error\n");
				exit(EXIT_FAILURE);
			}

			/* Step:5 Accept connections */
			for (i = 0; i < max_sfd; i++) {
				if (fd == sfd[i]) break;
			}
			if (i < max_sfd) {
				in_len = sizeof(in_addr);
#if defined(VMA_API) && (VMA_API == 1)
				fd = vma_comps[j].user_data;
				memcpy(&in_addr, &vma_comps[j].src, in_len);
#else
				fd = accept(fd, &in_addr, &in_len);
				if (fd < 0) {
					printf("Accept failed: %s", strerror(errno));
			        	exit(EXIT_FAILURE);
				}
				ev.events = EPOLLIN | EPOLLET;
				ev.data.fd = fd;
				if (epoll_ctl(efd, EPOLL_CTL_ADD, fd, &ev) == -1) {
					printf("epoll_ctl() failed: %s", strerror(errno));
					exit(EXIT_FAILURE);
			    	}
#endif /* VMA_API */

				conns.fds[conns.count] = fd;
				conns.count++;

				printf("Accepted: #%d by sfd=%d fd=%d from %s\n", conns.count, sfd[i], fd, _addr2str((struct sockaddr_in *)&in_addr));

				flag = 1;
				ret = setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, (char *) &flag, sizeof(int));
				if (ret < 0) {
					printf("Failed to disable NAGLE: %s\n", strerror(errno));
			        	exit(EXIT_FAILURE);
				}

				ret = _set_noblock(fd);
				continue;
			}

			/* Step:6 Process data */
			if (event & EPOLLIN) {
#if defined(VMA_API) && (VMA_API == 1)
			  printf("vma_comps[j].packet.num_bufs equal to %d \n", vma_comps[j].packet.num_bufs);
				assert(1 == vma_comps[j].packet.num_bufs);
				assert(sizeof(conns.msg) > vma_comps[j].packet.total_len);
				memcpy(conns.msg, vma_comps[j].packet.buff_lst->payload, vma_comps[j].packet.total_len);
				ret = vma_comps[j].packet.total_len;
				_vma_api->socketxtreme_free_vma_packets(&vma_comps[j].packet, 1);
#else
				ret = recv(fd, conns.msg, sizeof(conns.msg), 0);
#endif /* VMA_API */
				if (ret < 0) {
			        	exit(EXIT_FAILURE);
				}
				if (ret > 0) {
					conns.msg[ret - 1] = '\0';
				} else {
					conns.msg[0] = '\0';
				}
				printf("Received: fd=%d ret=%d %s\n", fd, ret, conns.msg);
			}
		}
	}

err:

	for (i = 0; i < max_sfd; i++) {
		if (sfd[i] > 0) {
			close(sfd[i]);
		}
	}

	for (i = 0; i < conns.count; i++) {
		if (conns.fds[i] > 0) {
#if defined(VMA_API) && (VMA_API == 1)
#else
			epoll_ctl(efd, EPOLL_CTL_DEL, conns.fds[i], NULL);
#endif /* VMA_API */
			close(conns.fds[i]);
		}
	}
	if (conns.fds) {
		free(conns.fds);
	}

#if defined(VMA_API) && (VMA_API == 1)
	if (vma_comps) {
		free(vma_comps);
	}
#else
	if (events) {
		free(events);
	}
#endif /* VMA_API */

	close(efd);

	exit(0);
}