#include <stdlib.h>
#include <time.h>
#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include "iolib.h"
#include "freedt.h"
#include "config.h"
const char *progname = "ratelimit";
const char *proghelp =
	"Usage: ratelimit [OPTIONS] command ...\n"
	"Run a command limiting stdin/stdout bandwidth and/or dumping "
	"IO to stderr.\n\n"
	"-i N      Limit input bandwidth to N bytes per second\n"
	"-o N      Limit output bandwidth to N bytes per second\n"
	"-r        Print all traffic to stderr\n";

/* The size of the data buffers either way. */
#define BUFFERSIZE 64

pid_t pid;

void recordio(char *data, int length, char *prefix) {
	buffer buf = BUFFER;

	if (length == 0) {
		bformat(&buf, "@i@c[EOF]\n", pid, prefix);
	} else {
		do {
			char c;
	
			bformat(&buf, "@i@c", pid, prefix);
			while (length--) {
				c = *data++;
				if (c == '\n') break;
				bappendc(&buf, c);
			}
			if (c == '\n')
				bappends(&buf, " \n");
			else
				bappends(&buf, "+\n");
		} while (length > 0);
	}
	writeba(fd_err, &buf);
	bfree(&buf);
}

int main(int argc, char **argv) {
	int inlimit = 0, outlimit = 0, record = 0;
	int incount, outcount, inopen = 1, rc;
	int inpipe[2], outpipe[2];
	char inbuf[BUFFERSIZE], outbuf[BUFFERSIZE];
	int inbsize = 0, outbsize = 0;
	time_t instart = -1, outstart = -1;

	pid = getpid();

	while (1) {
		int c = getopt(argc, argv, "+V?i:o:r");
		if (c == -1)
			break;
		switch (c) {
		case 'i':
			inlimit = atoi(optarg);
			break;
		case 'o':
			outlimit = atoi(optarg);
			break;
		case 'r':
			record = 1;
			break;
		case 'V':
			version();
		default:
			help();
		}
	}

	if ((argc - optind) < 1)
		help();

	if (pipe(inpipe) < 0 || pipe(outpipe) < 0)
		die("unable to create pipe");
	set_fd_cloexec(inpipe[0]);
	set_fd_cloexec(inpipe[1]);
	set_fd_cloexec(outpipe[0]);
	set_fd_cloexec(outpipe[1]);

	rc = fork();
	if (rc < 0)
		die("fork failed");
	if (rc > 0) {
		dup2(inpipe[0], 0);
		dup2(outpipe[1], 1);

		execvp(argv[optind], &argv[optind]);
		die2(argv[optind], "unable to exec");
	}
	close(inpipe[0]);
	close(outpipe[1]);

	incount = outcount = 0;
	while (1) {
		int rc, maxfd = -1;
		fd_set infds, outfds;
		struct timeval tv;
		time_t now = time(NULL);

		tv.tv_sec = 1;
		tv.tv_usec = 0;

#define ADD(fd, set) \
	do { \
		FD_SET(fd, set); \
		if (fd > maxfd) \
			maxfd = fd; \
	} while (0)

		FD_ZERO(&infds);
		if (inopen && inbsize == 0)
			ADD(fd_in, &infds);
		if (outbsize == 0)
			ADD(outpipe[0], &infds);

/* This is somewhat messy, but it's less messy than inlining it twice.
   A write is allowed if some data is available to be written,
   and some time has elapsed so that the rate can be measured,
   and the average rate so far doesn't exceed the limit. */
#define CAN_WRITE(bsize, limit, count, start) \
	((bsize > 0) && \
		((limit == 0) \
		 || ((now != start) && (count / (now - start) <= limit))))

		FD_ZERO(&outfds);
		if (inopen && CAN_WRITE(inbsize, inlimit, incount, instart))
			ADD(inpipe[1], &outfds);
		if (CAN_WRITE(outbsize, outlimit, outcount, outstart))
			ADD(fd_out, &outfds);

		rc = select(maxfd + 1, &infds, &outfds, NULL, &tv);
		if (rc < 0)
			die("select failed");

		if (rc == 0)
			continue;

		if (FD_ISSET(fd_in, &infds)) {
			inbsize = read(fd_in, inbuf, BUFFERSIZE);
			if (inbsize < 0)
				die("read from stdin failed");
			if (record)
				recordio(inbuf, inbsize, " < ");
			if (inbsize == 0) {
				close(inpipe[1]);
				inopen = 0;
			}
			incount += inbsize;
			if (instart == -1)
				instart = time(NULL);
		}

		if (FD_ISSET(outpipe[0], &infds)) {
			outbsize = read(outpipe[0], outbuf, BUFFERSIZE);
			if (outbsize < 0)
				die("read from child failed");
			if (record)
				recordio(outbuf, outbsize, " > ");
			if (outbsize == 0) {
				break;
			}
			outcount += outbsize;
			if (outstart == -1)
				outstart = time(NULL);
		}

		if (FD_ISSET(inpipe[1], &outfds)) {
			int n = write(inpipe[1], inbuf, inbsize);
			if (n < 0)
				die("write to child failed");
			inbsize -= n;
			memmove(inbuf, inbuf + n, inbsize);
		}

		if (FD_ISSET(fd_out, &outfds)) {
			int n = write(fd_out, outbuf, outbsize);
			if (n < 0)
				die("write to stdout failed");
			outbsize -= n;
			memmove(outbuf, outbuf + n, outbsize);
		}
	}

	return 0;
}