#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>

#include <sys/socket.h>
#include <sys/un.h>
#include <poll.h>

enum action {
	QUIT,
	ABORT,
	NEXT
};

/* vector of pollfds */
struct fds {
	size_t n;
	struct pollfd *buf;
};

/* simple hashmap */
struct ids {
	size_t n, s;
	char **buf;
};

static struct fds fds_create()
{
	return (struct fds) {
		.n = 0,
		.buf = NULL
	};
}

static int fds_add(struct fds *fds, int fd)
{
	for (size_t i = 0; i < fds->n; ++i) {
		/* reuse existing slot */
		if (fds->buf[i].fd == -1) {
			fds->buf[i].fd = fd;
			return 0;
		}
	}

	/* resize vector */
	size_t size = (fds->n + 1) * sizeof(struct pollfd);
	struct pollfd *new_buf = realloc(fds->buf, size);
	if (!new_buf)
		return -1;

	fds->buf = new_buf;
	fds->buf[fds->n].fd = fd;
	fds->buf[fds->n].events = POLLIN;
	fds->buf[fds->n].revents = 0;
	fds->n++;
	return 0;
}

static void fds_remove(struct fds *fds, size_t i)
{
	assert(i < fds->n);
	if (fds->buf[i].fd != -1) {
		close(fds->buf[i].fd);
		fds->buf[i].fd = -1;
	}
}

static void fds_destroy(struct fds *fds)
{
	for (size_t i = 0; i < fds->n; ++i)
		fds_remove(fds, i);

	free(fds->buf);
}

static struct ids ids_create()
{
	return (struct ids){
		.n = 0,
		.s = 0,
		.buf = NULL
	};
}

static void ids_destroy(struct ids *ids)
{
	for (size_t i = 0; i < ids->s; ++i)
		free(ids->buf[i]);
}

static inline size_t djb2(const char *s)
{
	size_t l = strlen(s);
	size_t h = 5381;
	for (size_t i = 0; i < l; ++i)
		h = ((h << 5) + h) + s[i];

	return h;
}

static char **ids_insert(struct ids *ids, char *id)
{
	/* when over 50% full, rebuild into new pow2 buffer */
	if (2 * ids->n >= ids->s) {
		size_t old_s = ids->s;
		char **old_buf = ids->buf;

		size_t new_s = ids->s ? 2 * ids->s : 1;
		char **new_buf = calloc(new_s, sizeof(char *));
		if (!new_buf)
			return NULL;

		ids->n = 0;
		ids->s = new_s;
		ids->buf = new_buf;

		/* place all items in old buffer into new buffer, this should
		 * always succeed */
		for (size_t i = 0; i < old_s; ++i) {
			if (!old_buf[i])
				continue;

			char **ok = ids_insert(ids, old_buf[i]);
			assert(ok);

			/* maybe unused when compiled with NDEBUG, avoid warning */
			(void)ok;
		}

		free(old_buf);
	}

	size_t h = djb2(id);
	for (size_t i = 0; i < ids->s; ++i) {
		/* fast pow2 modulo */
		size_t idx = (h + i) & (ids->s - 1);

		char **slot = &ids->buf[idx];
		/* insert into empty slot */
		if (!*slot) {
			ids->n++;
			*slot = id;
			return slot;
		}

		/* does this item already exist? */
		if (strcmp(*slot, id) == 0)
			return slot;
	}

	/* shouldn't happen */
	return NULL;
}

static enum action handle_req(struct ids *ids, int fd)
{
	const size_t len = 4096;
	char *buf = malloc(len);
	if (!buf) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed allocating req buf\n");
		free(buf);
		return ABORT;
	}

	ssize_t r = read(fd, buf, len - 1);
	if (r == 0) {
		/* EOF for socket */
		free(buf);
		return NEXT;
	}

	if (r == -1) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed reading req: %s\n",
				strerror(errno));
		free(buf);
		return ABORT;
	}

	/* ensure null termination */
	buf[r] = '\0';

	/* does the client want us to exit? */
	if (strcmp(buf, "quit\n") == 0) {
		free(buf);
		return QUIT;
	}

	/* check if we've already seen this location */
	char **item = ids_insert(ids, buf);
	if (!item) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed inserting location\n");
		free(buf);
		return ABORT;
	}

	int exists = *item != buf;
	if (exists)
		free(buf);

	ssize_t w = write(fd, &exists, sizeof(exists));
	if (w == -1) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed writing response: %s\n",
				strerror(errno));
		return ABORT;
	}

	if (w != sizeof(exists)) {
		fprintf(stderr, "\tCOVSRV: SRV -- partial write\n");
		return ABORT;
	}

	return NEXT;
}

static int handle_new(struct fds *fds, int sock)
{
	int client = accept(sock, NULL, NULL);
	if (client == -1) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed accepting connection: %s\n",
				strerror(errno));
		return -1;
	}

	if (fds_add(fds, client)) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed inserting new connection\n");
		close(client);
		return -1;
	}

	return 0;
}

static enum action handle_in(struct fds *fds, struct ids *ids, size_t i, int sock)
{
	if (fds->buf[i].fd == sock) {
		if (handle_new(fds, sock))
			return ABORT;

		return NEXT;
	} else {
		enum action a = handle_req(ids, fds->buf[i].fd);
		if (a != NEXT)
			return a;
	}

	return NEXT;
}

static enum action handle_event(struct fds *fds, struct ids *ids, int sock)
{
	for (size_t i = 0; i < fds->n; ++i) {
		if (fds->buf[i].revents & POLLIN) {
			enum action a = handle_in(fds, ids, i, sock);
			if (a != NEXT)
				return a;
		}
		else if (fds->buf[i].revents & (POLLHUP | POLLERR)) {
			/* shouldn't happen...? */
			if (fds->buf[i].fd == sock)
				return ABORT;

			fds_remove(fds, i);
		}
		else if (fds->buf[i].revents) {
			/* should never happen */
			abort();
		}

		/* because we directly add to the queue, we might end up going
		 * over an fd with zero, but there's no need to react to it in
		 * any way */
	}

	return NEXT;
}

static enum action loop(struct ids *ids, int sock)
{
	struct fds fds = fds_create();
	if (fds_add(&fds, sock)) {
		close(sock);
		return ABORT;
	}

	int ret = 0;
	enum action a = ABORT;
	while ((ret = poll(fds.buf, fds.n, -1)) > 0) {
		a = handle_event(&fds, ids, sock);
		if (a == NEXT)
			continue;

		break;
	}

	fds_destroy(&fds);
	return ret <= 0 ? ABORT : a;
}

static void usage(FILE *f, char *name)
{
	fprintf(f, "usage: COVSRV_SOCKET=<socket> %s <quit>\n", name);
}

static int quit(char *argv[], const char *name, int sock, struct sockaddr_un addr)
{
	if (strcmp(argv[1], "quit") != 0) {
		fprintf(stderr, "\tCOVSRV: SRV -- unknown command '%s'\n", argv[1]);
		usage(stderr, argv[0]);
		close(sock);
		return -1;
	}

	if (connect(sock, (const struct sockaddr *)&addr, sizeof(addr)) == -1) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed opening connection to %s: %s\n",
				name, strerror(errno));
		close(sock);
		return -1;
	}

	const char *msg = "quit\n";
	ssize_t w = write(sock, msg, strlen(msg));
	if (w != (ssize_t)strlen(msg)) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed sending quit command: %s\n",
				strerror(errno));
		close(sock);
		return -1;
	}

	close(sock);
	return 0;
}

int main(int argc, char *argv[])
{
	if (argc > 2) {
		usage(stderr, argv[0]);
		return -1;
	}

	char *name = getenv("COVSRV_SOCKET");
	if (!name) {
		fprintf(stderr, "\tCOVSRV: SRV -- COVSRV_SOCKET not defined\n");
		return -1;
	}

	int sock = socket(AF_UNIX, SOCK_SEQPACKET, 0);
	if (sock == -1) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed opening %s: %s\n",
				name, strerror(errno));
		return -1;
	}

	struct sockaddr_un addr;
	memset(&addr, 0, sizeof(addr));

	addr.sun_family = AF_UNIX;
	strncpy(addr.sun_path, name, sizeof(addr.sun_path) - 1);

	if (argc == 2)
		return quit(argv, name, sock, addr);

	if (bind(sock, (const struct sockaddr *)&addr, sizeof(addr)) == -1) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed binding %s: %s\n",
				name, strerror(errno));
		close(sock);
		return -1;
	}

	/* should probably be something like k * $(nproc), but 128 is probably
	 * fine for now */
	if (listen(sock, 128) == -1) {
		fprintf(stderr, "\tCOVSRV: SRV -- failed listening to %s: %s\n",
				name, strerror(errno));
		unlink(name);
		close(sock);
		return -1;
	}

	struct ids ids = ids_create();
	enum action a = loop(&ids, sock);
	ids_destroy(&ids);

	/* sock should've been closed by loop but we still need to unbind it */
	unlink(name);
	return a != QUIT;
}