#include <unistd.h>
#include <sys/select.h>

#include "process.h"
#include "debug.h"

static int child_count = 0;
static int parent_count = 0;
static int process_count = 0;

int get_child_count(void)
{
	return child_count;
}

int get_parent_count(void)
{
	return parent_count;
}

int get_sibling_count(void)
{
	return process_count;
}

int do_fork(void)
{
	int child, error;
	child = fork();
	if (child < 0) {
		error = errno;
		pr_err("fork() failed: %s\n", strerror(error));
		return -1;
	}

	if (child) {
		child_count++;
		pr_info("Fork %d, child %d\n", child_count, child);
		return child;
	}

	process_count = child_count;
	/* reset child's child count */
	child_count = 0;
	parent_count++;
	return 0;
}

int harvest_zombies(int pid)
{
	int status, error;

	if (child_count == 0)
		return 0;

	if (pid)
		pr_info("Waiting on pid %d, children left: %d\n", pid, 
			child_count);

	pid = waitpid(pid, &status, 0);
	if (pid < 0) {
		error = errno;
		pr_err("Error on wait(): %s\n", strerror(error));
		child_count--; /* Decrement child count anyway */
	}
	else {
		child_count--;
		pr_info("pid %d: exit code %d. Children left: %d\n", pid,
			status, child_count);
	}

	return 1;
}

/*
 * Runs a command cmd with params argv, connects stdin and stdout to
 * readfd and writefd
 *
 * Returns the pid of the executed process
 */
int run_piped(const char *cmd, char *const argv[],
	      int *stdinfd, int *stdoutfd, int *stderrfd)
{
	int ifd[2], ofd[2], efd[2], error, pid;

	if (stdinfd && pipe(ifd)) {
		error = errno;
		pr_err("pipe() failed: %s\n", strerror(error));
		return -1;
	}

	if (stdoutfd && pipe(ofd)) {
		error = errno;
		pr_err("pipe() failed: %s\n", strerror(error));
		return -1;
	}

	if (stderrfd && pipe(efd)) {
		error = errno;
		pr_err("pipe() failed: %s\n", strerror(error));
		return -1;
	}

	pid = do_fork();
	if (pid) { /* Parent side */
		if (stdinfd) {
			close(ifd[0]);
			*stdinfd = ifd[0];
		}

		if (stdoutfd) {
			close(ofd[1]);
			*stdoutfd = ofd[0];
		}

		if (stderrfd) {
			close(efd[1]);
			*stderrfd = efd[0];
		}

		return pid;
	}

	if (stdinfd) {
		close(ifd[1]);
		dup2(ifd[0], STDIN_FILENO);
	}

	if (stdoutfd) {
		close(ofd[0]);
		dup2(ofd[1], STDOUT_FILENO);
	}

	if (stderrfd) {
		close(efd[0]);
		dup2(efd[1], STDERR_FILENO);
	}

	/* Now we have redirected standard streams to parent process */
	execvp(cmd, argv);
	error = errno;
	pr_err("Failed to execv command %s: %s\n", cmd, strerror(error));
	exit(1);

	return 0;
}

/*
 * Runs a command cmd with params argv, connects stdin and stdout to
 * readfd and writefd
 *
 * Returns the pid of the executed process
 */
int run_piped_stream(const char *cmd, char *const argv[],
		     FILE **stdinf, FILE **stdoutf, FILE **stderrf)
{
	int ifd, ofd, efd, pid, error;
	int *i, *o, *e;

	if (stdinf)
		i = &ifd;
	else 
		i = 0;
	if (stdoutf)
		o = &ofd;
	else 
		o = 0;
	if (stderrf)
		e = &efd;
	else 
		e = 0;

	pid = run_piped(cmd, argv, i, o, e);

	if (stdinf) {
		*stdinf = fdopen(ifd, "r");
		if (*stdinf == NULL) {
			error = errno;
			pr_err("Error opening file stream for fd %d: %s\n",
			       ifd, strerror(error));
			return -1;
		}
	}

	if (stdoutf) {
		*stdoutf = fdopen(ofd, "r");
		if (*stdoutf == NULL) {
			error = errno;
			pr_err("Error opening file stream for fd %d: %s\n",
			       ofd, strerror(error));
			return -1;
		}
	}

	if (stderrf) {
		*stderrf = fdopen(efd, "r");
		if (*stderrf == NULL) {
			error = errno;
			pr_err("Error opening file stream for fd %d: %s\n",
			       efd, strerror(error));
			return -1;
		}
	}

	return pid;
}

/*
 * Forks a child and executes a command to run on parallel
 */

#define max(a,b) (a) < (b) ? (b) : (a)
#define BUF_SIZE (128*1024)
int run(const char *cmd, char *const argv[])
{
	int child, error;
	int ofd, efd;
	fd_set rfds;
	int maxfd;
	char stdoutstr[32], stderrstr[32], indent[16] = { "                " };

	indent[get_parent_count() + 1] = 0;
		
	if ((child = do_fork()))
	    return child;

	child = run_piped(cmd, argv, NULL, &ofd, &efd);
	snprintf(stdoutstr, 32, "%sstdout", green_color);
	snprintf(stderrstr, 32, "%sstderr", red_color);

	FD_ZERO(&rfds);
	FD_SET(ofd, &rfds);
	FD_SET(efd, &rfds);

	while (1) {
		char *sptr , *eptr;
		char rbuf[BUF_SIZE];
		int bytes;
		char *typestr;

		maxfd = max(ofd, efd);
		error = select(maxfd, &rfds, NULL, NULL, NULL);

		if (FD_ISSET(ofd, &rfds)) {
			typestr = stdoutstr;
			bytes = read(ofd, rbuf, BUF_SIZE);
			goto print;
		}

		if (FD_ISSET(efd, &rfds)) {
			typestr = stderrstr;
			bytes = read(efd, rbuf, BUF_SIZE);
			goto print;
		}

		pr_err("select() returned unknown fd\n");
		break;

print:
		if (bytes < 0) {
			error = errno;
			pr_err("read() failed: %s\n", strerror(error));
			break;
		}
		if (bytes == 0)
			break;

		sptr = eptr = rbuf;
		while(bytes--) {
			if (*eptr == '\n') {
				*eptr = 0;
				printf("%s[%5d %s] %s: %s%s\n", indent,
				       child, cmd, typestr, sptr, normal_color);
				sptr = eptr;
			}
			eptr++;
		}
	}

	harvest_zombies(child);	

	exit(1);
	return 0;
}