#include <u.h>
#include <libc.h>

int port = 17766;
char *iarg;
ulong debug, n, lflag;
long sleepms;

#define dprint(...) {if (debug) print(__VA_ARGS__);}

enum {
	Blen = 8*1024,
};

void
usage(void) 
{
	fprint(2, "usage: %s [-p port] [-i ip -n count] [-l]\n", argv0);
	exits("usage");
}

void
client(void)
{
	char *dst, b[Blen];
	int fd, i;

	sleep(1);
	dprint("client\n");
	if (!(dst = smprint("tcp!%s!%d", iarg, port)))
		sysfatal("%r");
	alarm(1*1000);
	if ((fd = dial(dst, 0, 0, 0)) < 0)
		sysfatal("%r");
	alarm(0);
	dprint("dialed\n");
	for (i = 0; i < n; ++i) {
		if (write(fd, b, sizeof b) < 0)
			sysfatal("%r");
	}
	close(fd);
	dprint("wrote\n");
}

void
server(void)
{
	int acfd, lcfd, dfd;
	char adir[40], ldir[40], b[Blen];
	long c;
	uvlong i;
	vlong startns, deltans;
	double ratemb;

	dprint("server\n");
	acfd = announce(smprint("tcp!*!%d", port), adir);
	if (acfd < 0)
		sysfatal("%r");

	dprint("listen\n");
	if ((lcfd = listen(adir, ldir)) < 0)
		sysfatal("%r");
	dprint("listened\n");
	if ((dfd = accept(lcfd, ldir)) < 0)
		sysfatal("%r");
	dprint("accepted\n");
	startns = nsec();
	for (i = 0; ; ++i) {
		c = readn(dfd, b, sizeof b);
		if (c != sizeof b)
			break;
		if (sleepms)
			sleep(sleepms);
	}
	deltans = nsec() - startns;
	close(dfd);
	close(lcfd);
	print("read count is %ulld\n", i);
	i *= Blen;
	ratemb = i*1000./deltans;
	print("%ulld bytes in %lld ns @ %.2g MB/s\n", i, deltans, ratemb);
}

void
main(int argc, char **argv)
{
	ARGBEGIN {
	case 'd':
		debug = 1;
		break;
	case 'p':
		port = strtoul(EARGF(usage()), nil, 0);
		break;
	case 'i':
		iarg = EARGF(usage());
		break;
	case 'n':
		n = strtoul(EARGF(usage()), nil, 0);
		break;
	case 's':
		sleepms = strtoul(EARGF(usage()), nil, 0);
		break;
	case 'l':
		lflag = 1;
		break;
	default:
		usage();
	} ARGEND

	if (lflag) {			// loop back
		switch (fork()) {
		case -1:
			sysfatal("rfork failed: %r");
			break;
		case 0:
			server();
			break;
		default:
			client();
			wait();
			break;
		}
	} else if (iarg)
		client();
	else
		server();
	exits(nil);
}