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

char	*service	= "17766";
char	*laddr	= "*";
char	*iaddr;

uchar	flag[127];
ulong	n;
long	sleepms;

#define dprint(...) {if(flag['d']) 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");
}

int
notef(void*, char *note)
{
	if(strstr(note, "alarm") != nil)
		return 1;
	return 0;
}

char*
·netmkaddr(char *linear, char *net, char *service)
{
	char *r;
	static QLock stupid;

	qlock(&stupid);		/* protect stupid static memory */
	r = netmkaddr(linear, net, service);
	if(r != nil)
		r = strdup(r);
	qunlock(&stupid);

	return r;
}

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

	atnotify(notef, 1);
	sleep(1);
	dprint("client\n");
	dst = ·netmkaddr(iaddr, "tcp", service);
	alarm(3*1000);
	if((fd = dial(dst, 0, 0, 0)) < 0)
		sysfatal("tcptest: %r");
	alarm(0);
	free(dst);
	dprint("dialed\n");
	for (i = 0; i < n; ++i) {
		if(write(fd, b, sizeof b) < 0)
			sysfatal("tcptest: %r");
	}
	close(fd);
	dprint("wrote\n");
}

double
cyclestoseconds(uvlong cycles)
{
	char buf[128], *f[10];
	int fd, n;
	static uvlong fasthz;

	if(fasthz == 0){
		fd = open("/dev/time", OREAD);
		if(fd == 0)
			return -1.;
		n = read(fd, buf, sizeof(buf)-1);
		if(n > 0){
			buf[n] = 0;
			n = tokenize(buf, f, nelem(f));
			if(n == 4)
				fasthz = strtoull(f[3], 0, 0);
		}
		close(fd);
		if(fasthz == 0)
			return -1.;
	}
	return (double)cycles / (double)fasthz;
}
				
void
server(void)
{
	int acfd, lcfd, dfd;
	char adir[40], ldir[40], b[Blen], *a;
	long c;
	uvlong i;
	vlong startns, deltans;
	double ratemb, maxr;
	NetConnInfo *nc;

	dprint("server\n");
	a = ·netmkaddr(laddr, "tcp", service);
	acfd = announce(a, adir);
	if(acfd == -1)
		sysfatal("tcptest: %r");
	free(a);

	dprint("listen\n");
	if((lcfd = listen(adir, ldir)) < 0)
		sysfatal("tcptest: %r");
	dprint("listened\n");
	if((dfd = accept(lcfd, ldir)) < 0)
		sysfatal("tcptest: %r");
	nc = getnetconninfo(ldir, dfd);
	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("%s count %llud; ", nc? nc->raddr: "(unknown)", i);
	free(nc);
	i *= Blen;
	ratemb = i*1000./deltans;
	if(sleepms > 0){
		maxr = Blen*1000./sleepms / 1e6;
		print("%llud bytes in %lld ns @ %.2g MB/s (%ldms; limit %.2g MB/s)\n", i, deltans, ratemb, sleepms, maxr);
	}
	else
		print("%llud bytes in %lld ns @ %.2g MB/s (0ms)\n", i, deltans, ratemb);
}

void
main(int argc, char **argv)
{
	ARGBEGIN {
	case 'p':
		service = EARGF(usage());
		break;
	case 'i':
		iaddr = EARGF(usage());
		break;
	case 'a':
		laddr = EARGF(usage());
		break;
	case 'n':
		n = strtoul(EARGF(usage()), nil, 0);
		break;
	case 's':
		sleepms = strtoul(EARGF(usage()), nil, 0);
		break;
	case 'd':
	case 'l':
		flag[ARGC()] = 1;
		break;
	default:
		usage();
	}ARGEND
	if(argc != 0)
		usage();

	if(flag['l']){			// loop back
		switch (fork()) {
		case -1:
			sysfatal("tcptest: %r");
			break;
		case 0:
			server();
			break;
		default:
			client();
			wait();
			break;
		}
	}else if(iaddr != nil)
		client();
	else
		server();
	exits(nil);
}