#define _POSIX_C_SOURCE 200112L
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <stdint.h>
#include <malloc.h>
#include <inttypes.h>
#include <errno.h>

#include <getopt.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/udp.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <netdb.h>

#define ARGV0 "udpecho"


struct {
	int bufsize[2];
	struct {
		bool help:1;
		bool verbose:1;
		bool v6only:1;
	} flags;
} opts;

struct {
	struct addrinfo *ai;
	int fd;
} g;

static void init_opts (void) {
	opts.bufsize[0] = opts.bufsize[1] = -1;
}

static void init_g (void) {
	g.fd = -1;
}

static void deinit_g (void) {
	if (g.ai != NULL) {
		freeaddrinfo(g.ai);
		g.ai = NULL;
	}

	if (g.fd >= 0) {
		close(g.fd);
		g.fd = -1;
	}
}

static void print_help (void) {
	printf("Usage: "ARGV0" [-hv46] [-p PORT/SERVICE] [LISTEN_ADDR]\n");
}

static int do_getaddrinfo (
	const char *node,
	const char *service,
	const int af,
	struct addrinfo **res)
{
	struct addrinfo hints = { 0, };

	hints.ai_family = af;
	hints.ai_socktype = SOCK_DGRAM;
	hints.ai_flags = AI_PASSIVE;

	return getaddrinfo(node, service, &hints, res);;
}

static bool parse_args (const int argc, const char **argv) {
	int fr;
	int af = AF_UNSPEC;
	const char *node = NULL;
	const char *service = "1114";

	while (true) {
		fr = getopt(argc, (char*const*)argv, "hv46p:");
		if (fr < 0) {
			break;
		}

		switch (fr) {
		case 'h': opts.flags.help = true; break;
		case 'v': opts.flags.verbose = true; break;
		case '4':
			af = AF_INET;
			opts.flags.v6only = false;
			break;
		case '6':
			af = AF_INET6;
			opts.flags.v6only = true;
			break;
		case 'p': service = optarg; break;
		case '?': return false;
		default: abort();
		}
	}

	if (argc > optind + 1) {
		fprintf(stderr, ARGV0": too many args\n");
		return false;
	}
	if (argc > optind) {
		node = argv[optind];
	}

	fr = do_getaddrinfo(node, service, af, &g.ai);
	if (fr != 0) {
		const char *msg = gai_strerror(fr);

		if (strchr(service, ':')) {
			fprintf(stderr, ARGV0": [%s]:%s: %s\n", node, service, msg);
		}
		else {
			fprintf(stderr, ARGV0": %s:%s: %s\n", node, service, msg);
		}

		return false;
	}

	return true;
}

static const struct addrinfo *pick_af (
	const struct addrinfo *list,
	const int af)
{
	while (list != NULL) {
		if (list->ai_family == af) {
			return list;
		}
		list = list->ai_next;
	}

	return NULL;
}

static int mksock (const int af) {
	int fd;
	int fr;
	int sv;
	const struct addrinfo *ai = pick_af(g.ai, af);

	if (ai == NULL) {
		errno = ENOENT;
		return -1;
	}

	fd = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
	if (fd < 0) {
		perror(ARGV0": socket()");
		goto ERR;
	}

	sv = 1;
	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &sv, sizeof(sv));
#ifdef IPV6_V6ONLY // Linux-specific flag
	if (opts.flags.v6only) {
		sv = 1;
		setsockopt(fd, IPPROTO_IPV6, IPV6_V6ONLY, &sv, sizeof(sv));
	}
#endif

	fr = bind(fd, ai->ai_addr, ai->ai_addrlen);
	if (fr != 0) {
		perror(ARGV0": bind()");
		goto ERR;
	}

	return fd;
ERR:
	close(fd);
	return -1;
}

static void report_remote (
	const struct sockaddr *sa,
	const void *data,
	const size_t len)
{
	uint16_t hport;
	const void *addr_buf;
	char addr_str[INET6_ADDRSTRLEN];
	const char *fr;

	switch (sa->sa_family) {
	case AF_INET:
		addr_buf = &((const struct sockaddr_in*)sa)->sin_addr;
		hport = ntohs(((const struct sockaddr_in*)sa)->sin_port);
		break;
	case AF_INET6:
		addr_buf = &((const struct sockaddr_in6*)sa)->sin6_addr;
		hport = ntohs(((const struct sockaddr_in6*)sa)->sin6_port);
		break;
	default: abort();
	}

	fr = inet_ntop(sa->sa_family, addr_buf, addr_str, sizeof(addr_str));
	if (fr != NULL) {
		switch (sa->sa_family) {
		case AF_INET:
			printf("%5zu bytes from %s:%"PRIu16"\n", len, addr_str, hport);
			break;
		case AF_INET6:
			printf("%5zu bytes from [%s]:%"PRIu16"\n", len, addr_str, hport);
			break;
		}
	}
}

static int server_loop (void) {
	// If the flag doesn't get optimised by the compiler, the bitfield access is
	// extremely slow. Might as well cache it.
	const bool verbose = opts.flags.verbose;
	const int send_flags = MSG_NOSIGNAL
#ifdef MSG_DONTWAIT // Linux-specific flag
		| MSG_DONTWAIT
#endif
		;
	uint8_t buf[UINT16_MAX]; // as per UDP header
	union {
		struct sockaddr sa;
		struct sockaddr_in sin4;
		struct sockaddr_in6 sin6;
	} addr;
	ssize_t iofr;
	socklen_t sl;

	while (true) {
		sl = sizeof(addr);
		iofr = recvfrom(g.fd, buf, sizeof(buf), 0, &addr.sa, &sl);
		if (iofr < 0) {
			switch (errno) {
			case ECONNREFUSED: // ICMP messages
				break;
			default:
				perror(ARGV0": recvfrom()");
				abort();
			}
		}

		if (verbose) {
			report_remote(&addr.sa, buf, (size_t)iofr);
		}

		sendto(g.fd, buf, (size_t)iofr, send_flags, &addr.sa, sl);
	}
}

int main (const int argc, const char **argv) {
	static const int AF_AFF[] = { AF_INET6, AF_INET };
	static int ec = 0;

	init_opts();
	init_g();

	if (!parse_args(argc, argv)) {
		ec = 2;
		goto END;
	}

	if (opts.flags.help) {
		print_help();
		ec = 0;
		goto END;
	}

	for (size_t i = 0; i < sizeof(AF_AFF) / sizeof(AF_AFF[0]); i += 1) {
		g.fd = mksock(AF_AFF[i]);
		if (g.fd >= 0) {
			break;
		}
	}
	if (g.fd < 0) {
		ec = 1;
		goto END;
	}

	ec = server_loop();
END:
	deinit_g();
	return ec;
}