#include <stdlib.h>
#include <string.h>
#include <netinet/in.h>
#include <stdio.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <netdb.h>
#include <sys/time.h>


#define PORT 9930

struct message {
	double ts;
	char data[512];
};

struct udpp_options {
	int port;
	char *host;
	int pings;
};

static void
die(char *s)
{
	perror(s);
	exit(1);
}

static inline double
ts()
{
	struct timeval tv;

	gettimeofday(&tv, NULL);

	return (double)(tv.tv_sec % 86400) + (double)(tv.tv_usec) / 1E6;
}

static int
socket_server(struct udpp_options *opts, struct sockaddr_in *sa)
{
	int s;

	if ((s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == -1)
		die("socket");

	memset((char*)sa, 0, sizeof(*sa));
	sa->sin_family = AF_INET;
	sa->sin_port = htons(opts->port);
	sa->sin_addr.s_addr = htonl(INADDR_ANY);
	if (bind(s, (const struct sockaddr*)sa, sizeof(*sa)) == -1)
		die("bind");

	return s;
}

static int
socket_client(struct udpp_options *opts, struct sockaddr_in *sa)
{
	int s;
	struct hostent *server;

	if ((s = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == -1)
		die("socket");

	server = gethostbyname(opts->host);
	if (server == NULL)
		die("ERROR, no such host");

	memset((char*)sa, 0, sizeof(*sa));
	sa->sin_family = AF_INET;
	sa->sin_port = htons(opts->port);
	memcpy((char*)&sa->sin_addr.s_addr, (char*)server->h_addr, server->h_length);

	return s;
}

static void
pprint(struct sockaddr_in *client, struct message *msg, double end, char ping)
{
	double dur;

	dur = end - msg->ts;
	printf("\n\tReceived p%cng packet from %s:%d ", ping,
	       inet_ntoa(client->sin_addr), ntohs(client->sin_port));
	printf("(%f s)", dur);
	printf("\n\tData: %s\n\n", msg->data);
}

static void
client(struct udpp_options *opts)
{
	struct sockaddr_in ping;
	struct sockaddr_in pong;
	struct sockaddr_in client;
	double end;
	int s;
	int p;
	int i;
	socklen_t slen = sizeof(ping);
	struct message msg;

	s = socket_client(opts, &ping);
	opts->port++;
	p = socket_server(opts, &pong);

	for (i = 0; i < opts->pings; i++) {
		printf("-->\n\tSending ping packet %d\n", i);
		sprintf(msg.data, "This is ping packet %d\n", i);
		msg.ts = ts();
		if (sendto(s, &msg, sizeof(msg), 0, (const struct sockaddr*)&ping, slen) == -1)
			die("sendto()");
		if (recvfrom(p, &msg, sizeof(msg), 0, (struct sockaddr*)&client, &slen) == -1)
			die("recvfrom()");
		end = ts();
		pprint(&client, &msg, end, 'o');
		printf("<--\n");
		sleep(1);
	}

	close(s);
	exit(0);
}

static void
server(struct udpp_options *opts)
{
	struct sockaddr_in ping;
	struct sockaddr_in pong;
	struct sockaddr_in client;
	double end;
	int s;
	int p;
	int i = 0;
	socklen_t slen = sizeof(client);
	struct message msg;

	s = socket_server(opts, &ping);
	opts->port++;
	printf("Starting pong mode\n\n");

	while (1) {
		if (recvfrom(s, &msg, sizeof(msg), 0, (struct sockaddr*)&client, &slen) == -1)
			die("recvfrom()");
		end = ts();
		printf("-->\n");
		pprint(&client, &msg, end, 'i');
		printf("\tSending pong packet %d\n<--\n", i);
		free(opts->host);
		opts->host = strdup(inet_ntoa(client.sin_addr));
		p = socket_client(opts, &pong);
		sprintf(msg.data, "This is pong packet %d\n", i++);
		msg.ts = ts();
		if (sendto(p, &msg, sizeof(msg), 0, (const struct sockaddr*)&pong, slen) == -1)
			die("sendto()");
		close(p);
	}

	close(s);
	exit(0);
}

static void
usage()
{
	fprintf(stderr, "udpp - udp packet ping pong\n");
	fprintf(stderr, "\n");
	fprintf(stderr, "Options:\n");
	fprintf(stderr, "  -h                Show this help screen.\n");
	fprintf(stderr, "  -r <host:port>    Connect to server running on host:port.\n");
	fprintf(stderr, "  -s                Run in server mode.\n");
	fprintf(stderr, "  -n <number>       Number of pings and pongs (defaults to 10).\n");

	exit(-1);
}

int
main(int argc, char *argv[])
{
	int opt;
	struct udpp_options opts;
	int srv = 0;

	opts.port = PORT;
	opts.host = NULL;
	opts.pings = 10;


	while ((opt = getopt(argc, argv, "hsr:n:")) != -1) {
		switch (opt) {
		case 'h':
			usage();
		case 'n':
			opts.pings = atoi(optarg);
			break;
		case 'r':
			opts.host = strchr(optarg, ':');
			if (!opts.host)
				die("derp");
			/* port needs to be at least one digit long */
			if (strlen(optarg) < strchr(optarg, ':') - optarg + 2 )
				die("derp derp");
			opts.port = atoi(strchr(optarg, ':') + 1);
			optarg[strchr(optarg, ':') - optarg] = 0;
			opts.host = strdup(optarg);
			break;
		case 's':
			srv = 1;
			break;
		default:
			usage();
		}
	}

	if (!opts.host && !srv)
		usage();

	if (!srv)
		client(&opts);

	opts.host = strdup("localhost");
	server(&opts);

	return 0;
}