#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>

#include <getopt.h>
#include <err.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <unistd.h>
#include <netdb.h>

#include "chargen.h"

enum {
	TCP_MAXLEN = 1024,
	UDP_MAXLEN = 512,
};

void usage_and_die(const char *myname, const char *fmt, ...)
{
	if (fmt)
	{
		va_list ap;
		va_start(ap, fmt);
		vwarnx(fmt, ap);
	}
	errx(EXIT_FAILURE, "Usage: %s [-4 | -6] [-u | -t] [-l listen-addr] [-p listen-port]", myname);
}

enum { MAX_LISTENERS = 8, MAX_CONNECTED = 64 };
int listen_fds[MAX_LISTENERS];
int packet_fds[MAX_LISTENERS];
int stream_fds[MAX_CONNECTED];
unsigned num_listen, num_stream, num_packet;

int main(int argc, char **argv)
{
	int opt;
	int ret;

	char hbuf[1024];
	char sbuf[32];

	struct addrinfo hints = {0};
	struct addrinfo *bind_addrs;

	const char *listen_addr = NULL;
	const char *listen_port = NULL;

	hints.ai_flags = AI_PASSIVE;
	hints.ai_family = AF_UNSPEC;

	while ((opt = getopt(argc, argv, "46l:p:tu")) != -1)
	{
		switch (opt)
		{
			case '4':
			case '6':
				if (hints.ai_family != AF_UNSPEC)
					usage_and_die(argv[0], "-4 and -6 are incompatible (leave out to use both)");
				hints.ai_family = (opt == '4') ? AF_INET : AF_INET6;
				break;
			case 't':
			case 'u':
				if (hints.ai_protocol != 0)
					usage_and_die(argv[0], "-t and -u are incompatible (leave out to use both)");
				hints.ai_protocol = (opt == 't') ? IPPROTO_TCP : IPPROTO_UDP;
				break;
			case 'l':
				if (listen_addr != NULL)
					usage_and_die(argv[0], "Only one listen address allowed!");
				listen_addr = optarg;
				break;
			case 'p':
				if (listen_port != NULL)
					usage_and_die(argv[0], "Only one listen port allowed!");
				listen_port = optarg;
				break;
			default:
				usage_and_die(argv[0], NULL);
				/* not reached */
		}
	}

	argc -= optind;
	argv += optind;

	if (argc != 0)
		usage_and_die(argv[0], NULL);

	if (listen_port == NULL)
		listen_port = "chargen";

	printf("Opening listeners on %s:%s...\n", (listen_addr == NULL) ? "*" : listen_addr, listen_port);
	ret = getaddrinfo(listen_addr, listen_port, &hints, &bind_addrs);
	if (ret)
		errx(EXIT_FAILURE, "getaddrinfo: %s", gai_strerror(ret));
	for (struct addrinfo *a = bind_addrs; a != NULL; a = a->ai_next)
	{
		int s;

		ret = getnameinfo(a->ai_addr, a->ai_addrlen, hbuf, sizeof hbuf, sbuf, sizeof sbuf, NI_NUMERICHOST | NI_NUMERICSERV | (a->ai_socktype == SOCK_STREAM) ? 0 : NI_DGRAM);
		if (ret)
			warn("getnameinfo: %s", gai_strerror(ret));
		else
			printf("Opening %s listener on %s:%s...\n", (a->ai_socktype == SOCK_STREAM) ? "stream" : "packet", hbuf, sbuf);

		s = socket(a->ai_family, a->ai_socktype, a->ai_protocol);
		if (s == -1)
		{
			warn("socket");
			continue;
		}

		if (bind(s, a->ai_addr, a->ai_addrlen) == -1)
		{
			warn("bind: %s:%s", hbuf, sbuf);
			close(s);
			continue;
		}

		if (a->ai_socktype == SOCK_STREAM)
		{
			if (listen(s, 128) == -1)
			{
				warn("listen");
				close(s);
				continue;
			}

			if (num_stream >= MAX_LISTENERS)
			{
				warnx("Listener table full, discarding %s:%s", hbuf, sbuf);
				close(s);
			}
			else
				listen_fds[num_listen++] = s;
		}
		else
		{
			if (num_packet >= MAX_LISTENERS)
			{
				warnx("Listener table full, discarding %s:%s", hbuf, sbuf);
				close(s);
			}
			else
				packet_fds[num_packet++] = s;
		}
	}

	if (num_packet + num_listen == 0)
		errx(EXIT_FAILURE, "No listeners to run on!\n");

	while(1)
	{
		fd_set r, w;
		int maxfd = 0;

		union
		{
			struct sockaddr_in a4;
			struct sockaddr_in6 a6;
		} addr;
		/* 64k = max UDP payload size */
		static unsigned char data[64 * 1024];

		FD_ZERO(&r);
		FD_ZERO(&w);

		if (num_stream < MAX_CONNECTED)
		{
			for (unsigned i = 0; i < num_listen; i++)
			{
				FD_SET(listen_fds[i], &r);
				if (listen_fds[i] > maxfd)
					maxfd = listen_fds[i];
			}
		}
		for (unsigned i = 0; i < num_packet; i++)
		{
			FD_SET(packet_fds[i], &r);
			if (packet_fds[i] > maxfd)
				maxfd = packet_fds[i];
		}
		for (unsigned i = 0; i < num_stream; i++)
		{
			FD_SET(stream_fds[i], &r);
			FD_SET(stream_fds[i], &w);
			if (stream_fds[i] > maxfd)
				maxfd = stream_fds[i];
		}

		ret = select(maxfd+1, &r, &w, NULL, NULL);
		if (ret == -1)
			err(EXIT_FAILURE, "select");
		if (ret == 0)
		{
			printf("Huh, that's weird, select says nothing ready but I didn't ask for a timeout\n");
			continue;
		}

		for (unsigned i = 0; i < num_listen; i++)
		{
			int l = listen_fds[i];
			if (FD_ISSET(l, &r))
			{
				struct sockaddr *sa = (struct sockaddr *)&addr;
				socklen_t len = sizeof addr;
				int s = accept(l, sa, &len);
				if (s == -1)
				{
					warn("accept");
					continue;
				}
				stream_fds[num_stream++] = s;

				ret = getnameinfo(sa, len, hbuf, sizeof hbuf, sbuf, sizeof sbuf, NI_NUMERICHOST | NI_NUMERICSERV);
				if (ret)
					warn("getnameinfo: %s", gai_strerror(ret));
				else
					printf("Accepting stream connection from %s:%s...\n", hbuf, sbuf);
			}

			if (num_stream >= MAX_CONNECTED)
				break;
		}

		for (unsigned i = 0; i < num_packet; i++)
		{
			int s = packet_fds[i];
			if (FD_ISSET(s, &r))
			{
				ssize_t rval;

				struct sockaddr *sa = (struct sockaddr *)&addr;
				socklen_t len = sizeof addr;

				rval = recvfrom(s, data, sizeof data, 0, sa, &len);
				if (rval == -1)
				{
					warn("recvfrom");
					continue;
				}

				ret = getnameinfo(sa, len, hbuf, sizeof hbuf, sbuf, sizeof sbuf, NI_NUMERICHOST | NI_NUMERICSERV);
				if (ret)
					warn("getnameinfo: %s", gai_strerror(ret));
				else
					printf("Responding to packet from %s:%s...\n", hbuf, sbuf);

				//rval = sendto(s, "chargen data\r\n", 14, 0, sa, len);
				char character[UDP_MAXLEN];
				size_t datalen = genChar(character, sizeof character);
				rval = sendto(s, character, datalen, 0, sa, len);
				if (rval == -1)
					warn("sendto");
			}
		}
		for (unsigned i = 0; i < num_stream; i++)
		{
			int s = stream_fds[i];
			if (FD_ISSET(s, &r))
			{
				ssize_t rval;
				rval = read(s, data, sizeof data);
				if (rval <= 0)
				{
					if (rval == -1)
						warn("read");
					printf("Closing stream connection\n");
					close(s);

					/*
					 * We will miss checking the socket we move down to fill the gap we're leaving.
					 * That's OK, if it's ready we'll get it next time.
					 */
					stream_fds[i] = stream_fds[--num_stream];
					continue;
				}
				/* Ignore data from a successful read */
			}
			if (FD_ISSET(s, &w))
			{
				if (s == 0)
				{
					printf("WTF?  i=%u num_stream=%u stream_fds[i]=%d s=%d\n", i, num_stream, stream_fds[i], s);
					continue;
				}
				char character[TCP_MAXLEN];
				size_t datalen = genChar(character, sizeof character);
				/* TODO: Detect short writes and save the rest of the buffer for later */
				if (write(s, character, datalen) == -1)
				{
					warn("write");
					printf("Closing stream connection\n");
					close(s);

					/*
					 * We will miss checking the socket we move down to fill the gap we're leaving.
					 * That's OK, if it's ready we'll get it next time.
					 */
					stream_fds[i] = stream_fds[--num_stream];
					continue;
				}
			}
		}
	}

	/* not reached */
}