/*
 * Copy me if you can.
 * by 20h
 */

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

long
writen(int fd, char* buf, long len)
{
	int cnt, rlen = 0; char* b = buf;
	for (;;) {
		cnt = write(fd, b, len);
		if (cnt <= 0)
			break;
		rlen += cnt;
		len -= cnt;
		if (len == 0)
			break;
		b += cnt;
	}
	return rlen;
}

void
pattern(char* buf, int buflen)
{
	int i;
	char ch = ' ';
	char* b = buf;
	for (i = 0; i < buflen; i++) {
		*b++ = ch++;
		if (ch == 127)
			ch = ' ';
	}
}

void
reader(char* addr, char* port, int buflen, int nbuf, int sink)
{
	char *buf, adir[40], ldir[40];
	int fd, cnt, acfd, lcfd;
	vlong nbytes = 0;
	vlong now;
	double elapsed;
	int pd;
	char peer[100];


	acfd = announce(netmkaddr(addr, "tcp", port), adir);
	if(acfd < 0)
		sysfatal("announce: %r");

	print("ttcp-r: buflen=%d, nbuf=%d, port=%s tcp\n",
		buflen, nbuf, port);
	buf = malloc(buflen);

	lcfd = listen(adir, ldir);
	if(lcfd < 0)
		sysfatal("listen: %r");

	fd = accept(lcfd, ldir);
	if(fd < 0)
		return;

	sprint(peer, "%s/remote", ldir);
	pd = open(peer, OREAD);
	cnt = read(pd, peer, 100);
	close(pd);

	print("ttcp-r: accept from %*.*s", cnt, cnt, peer);
	now = nsec();
	if (sink) {
		while((cnt = readn(fd, buf, buflen)) > 0)
			nbytes += cnt;
	} else {
		while((cnt = readn(fd, buf, buflen)) > 0 &&
		      write(1, buf, cnt) == cnt)
			nbytes += cnt;
	}
	elapsed = (nsec() - now)/1E9;

	print("ttcp-r: %lld bytes in %.2f real seconds = %.2f KB/sec\n",
	      nbytes, elapsed, nbytes/elapsed/1024);
}

void
writer(char* addr, char* port, int buflen, int nbuf, int src)
{
	char *buf;
	int fd, cnt;
	vlong nbytes = 0;
	vlong now;
	double elapsed;

	print("ttcp-t: buflen=%d, nbuf=%d, port=%s tcp\n",
		buflen, nbuf, port);

	buf = malloc(buflen);
	fd = dial(netmkaddr(addr, "tcp", port), 0, 0, 0);
	if(fd < 0)
		sysfatal("dial: %r");

	print("ttcp-t: connect\n");

	now = nsec();
	if (src) {
		pattern(buf, buflen);
		while (nbuf-- && writen(fd, buf, buflen) == buflen)
			nbytes += buflen;
	} else {
		while ((cnt = read(0, buf, buflen)) > 0 &&
		       writen(fd, buf, cnt) == cnt)
			nbytes += cnt;
	}
	elapsed = (nsec() - now)/1E9;

	print("ttcp-t: %lld bytes in %.2f real seconds = %.2f KB/sec\n",
	      nbytes, elapsed, nbytes/elapsed/1024);
}

void
usage(void)
{
	print("usage:\tttcp -t [options] host [ < in]\n"
	      "\t\tttcp -r [options] [ > out]\n"
	      " options:\n"
	      "  -l\t\tlength of buf (default 8192)\n"
	      "  -p port\tport number (default 5001)\n"
	      "  -n num\tnumber of bufs written (default 2048)\n"
	      "  -s\t\tsource or sink all data from the network\n"
	      );
	exits(0);
}

void
main(int argc, char *argv[])
{
	int buflen = 8192;
	int nbuf = 2048;
	int srcsink = 0;
	char* port = "5001";
	enum {none, recv, xmit} mode = none;

	ARGBEGIN {
	case 'n':
		nbuf = atoi(EARGF(usage()));
		break;
	case 'p':
		port = EARGF(usage());
		break;
	case 'l':
		buflen = atoi(EARGF(usage()));
		break;
	case 't':
		mode = xmit;
		break;
	case 'r':
		mode = recv;
		break;
	case 's':
		srcsink = 1;
		break;
	default:
		usage();
	} ARGEND;

	USED(buflen);
	switch (mode) {
	case none:
		usage();
		break;
	case xmit:
		if (argv[0] == nil)
			usage();
		writer(argv[0], port, buflen, nbuf, srcsink);
		break;
	case recv:
		reader("*", port, buflen, nbuf, srcsink);
		break;
	}
	exits(0);
}