#include <err.h>
#include <poll.h>
#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>

struct server_conn {
	int fd;
	struct sockaddr addr;
	socklen_t addrlen;
	/* some other stuff */
};

#define VEC_NAME server_conns
#define VEC_TYPE struct server_conn
#include <conts/vec.h>

#define VEC_NAME conn_polls
#define VEC_TYPE struct pollfd
#include <conts/vec.h>

struct server_ctx {
	int tcp_sock;
	struct server_conns conns;
	struct conn_polls polls;
	int timeout;
};

/* buf is sized by count */
static ssize_t readx(int fd, void *buf, size_t count)
{
	while (1) {
		ssize_t r = read(fd, buf, count);
		if (r == -1) {
			if (errno == EINTR)
				/* read interrupted, try again */
				continue;

			if (errno == EWOULDBLOCK)
				/* no more data to read for now */
				return 0;
		}

		/* we've read as many bytes as are currently available */
		return r;
	}

	/* should never be reached, but keep here to avoid warnings */
	return 0;
}

static ssize_t writex(int fd, void *buf, size_t count)
{
	while (1) {
		ssize_t r = write(fd, buf, count);
		if (r == -1) {
			if (errno == EINTR)
				/* write interrupted, try again */
				continue;

			if (errno == EWOULDBLOCK)
				/* no more data to read for now */
				return 0;
		}

		return r;
	}

	/* should never be reached, but keep here to avoid warnings */
	return 0;
}

static int new_connection(struct server_ctx *ctx, int new_conn_fd)
{
	struct sockaddr addr;
	socklen_t addrlen;

	int ret = 0;

	/** @todo retry count? */
	while (1) {
		ret = accept(new_conn_fd, &addr, &addrlen);
		if (ret == -1) {
			/* try again */
			if (errno == EINTR)
				continue;

			/* some other error, throw it at the user */
			return ret;
		}

		/* we have a valid fd */
		break;
	}

	int fd = ret;
	ret = fcntl(fd, F_SETFL, O_NONBLOCK);
	if (ret == -1) {
		close(fd);
		return ret;
	}

	struct server_conn conn = {
		.fd      = fd,
		.addr    = addr,
		.addrlen = addrlen
	};

	server_conns_append(&ctx->conns, conn);
	return 0;
}

/* I'll assume that each individual subcommant has ensured all its resources are
 * freed by now. Invalidates the conn pointer! */
static int kill_connection(struct server_ctx *ctx, struct server_conn *conn)
{
	close(conn->fd);

	assert(server_conns_len(&ctx->conns) > 1);

	if (server_conns_len(&ctx->conns) == 1) {
		server_conns_shrink(&ctx->conns, 0);
		return 0;
	}

	/* overwrite our location in the vector with whaterver is at the back of
	 * the vector */
	*conn = *server_conns_pop(&ctx->conns);
	return 0;
}

static int event_loop(int tcp_sock)
{
	struct server_ctx ctx = {
		.tcp_sock  = tcp_sock,
		.conns     = server_conns_create(0),
		.polls     = conn_polls_create(0),
		.timeout   = -1
	};

	int ret = 0;
	while (1) {
		/* build up set of connections to poll */
		conn_polls_reset(&ctx.polls);

		struct pollfd tcp_poll = {
			.fd      = ctx.tcp_sock,
			.events  = POLLIN,
			.revents = 0
		};
		conn_polls_append(&ctx.polls, tcp_poll);

		foreach(server_conns, c, &ctx.conns) {
			struct pollfd conn_poll = {
				.fd      = c->fd,
				.events  = POLLIN,
				.revents = 0
			};
			conn_polls_append(&ctx.polls, conn_poll);
		}

		/* wait for something to happen on incoming connections, or a
		 * new connection. Eventually might also want to wait on
		 * outgoing connections to become available for writing? */
		ret = poll(ctx.polls.buf, conn_polls_len(&ctx.polls), ctx.timeout);
		if (ret == 0) {
			/* timeout */
			if (server_conns_len(&ctx.conns) == 0) {
				/* nothing to do, quit for now */
				break;
			}

			/* some command just timed out, handle it by fallthrough */
		}

		if (ret == -1) {
			/* some kind of error */
			if (errno == EINTR) {
				/* safe to try again */
				continue;
			}

			/* exit loop due to unknown error*/
			err(EXIT_FAILURE, "poll failed");
		}

		/* otherwise, there's some event we should react to */
		size_t l = conn_polls_len(&ctx.polls);
		assert(l >= 1);

		/* is it a new connection? */
		struct pollfd *new_conn = conn_polls_at(&ctx.polls, 0);
		if (new_conn->revents != 0) {
			if (new_connection(&ctx, new_conn->fd))
				err(EXIT_FAILURE, "tcp socket exploded");

			new_conn->revents = 0;
		}

		for (size_t poll_idx = 1; poll_idx < l; ++poll_idx) {
			size_t conn_idx = poll_idx - 1;

			struct pollfd *pollfd = conn_polls_at(&ctx.polls, poll_idx);
			struct server_conn *conn = server_conns_at(&ctx.conns, conn_idx);

			if (pollfd->revents & (POLLHUP | POLLNVAL | POLLERR)) {
				kill_connection(&ctx, conn);
				continue;
			}

			if (!(pollfd->revents & POLLIN))
				continue;

			/* if there's no error, only POLLIN should be set, as it's the one we
			 * requested */
			assert((pollfd->revents & POLLIN) == POLLIN);
			while (1) {
				/* for testing, just echo everything */
				char buf[4096];
				ssize_t r = readx(conn->fd, buf, 4096);
				if (r == -1) {
					/* no more data to read for now */
					if (errno == EWOULDBLOCK)
						break;

					/* something funky with the socket, kill the
					 * connection and let the client try again later
					 * if it wishes to */
					kill_connection(&ctx, conn);
					break;
				}

				ssize_t w = writex(conn->fd, buf, r);
				if (w == -1) {
					/* likewise, socket is funked up, kill
					 * connection */
					kill_connection(&ctx, conn);
					break;
				}

				/* if readx or writex are broken, this shouldn't
				 * hold */
				assert(w == r);

				/* if we filled the buffer, check if there's
				 * more to read */
				if (r == 4096)
					continue;

				/* otherwise, break out of echo */
				break;
			}

			pollfd->revents = 0;
		}
	}

	conn_polls_destroy(&ctx.polls);
	server_conns_destroy(&ctx.conns);
	return ret;
}

int main(int argc, char *argv[argc])
{
	int ret = 0;

	int tcp_sock = socket(AF_INET, SOCK_STREAM, 0);
	if (tcp_sock == -1) {
		err(EXIT_FAILURE, "failed opening tcp socket");
	}

	struct sockaddr_in addr;
	memset(&addr, 0, sizeof(addr));
	addr.sin_family      = AF_INET;
	addr.sin_addr.s_addr = htonl(INADDR_ANY);
	/* should probably come up with a reasonable port at some point */
	addr.sin_port        = htons(24242);

	if (bind(tcp_sock, (struct sockaddr *)&addr, sizeof(addr)) == -1) {
		err(EXIT_FAILURE, "failed binding tcp socket");
	}

	if (listen(tcp_sock, 64) == -1) {
		err(EXIT_FAILURE, "failed listening on tcp socket");
	}

	ret = event_loop(tcp_sock);
	close(tcp_sock);
	return ret;
}