#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <string.h>
#include <lua.h>
#include <lauxlib.h>

#define CKE(val, name) do { if (val == -1) { perror(name); exit(1); } } while (0)

#ifdef DEBUG
#define DP(...) fprintf(stderr, __VA_ARGS__);
#else
#define DP(...)
#endif

int filefork(char *argv[], int input, int output) {
	int pid = fork();
	switch (pid) {
	case -1: perror("fork"); exit(1); break;
	case 0:
		if (input != -1) {
			CKE(close(STDIN_FILENO),"close stdin");
			CKE(dup2(input, STDIN_FILENO),"dup2 stdin");
		}
		if (output != -1) {
			CKE(close(STDOUT_FILENO),"close stdout");
			CKE(dup2(output, STDOUT_FILENO),"dup2 stdout");
		}
		CKE(execvp(argv[0], argv), "exec");
		exit(1);
		break;
	default: return pid;
	}
}

int echofork(char *str, size_t len, int fd) {
	int pid = fork();
	switch (pid) { 
	case -1: perror("fork"), exit(1); break;
	case 0:
		size_t remaining = len;
		char *cur = str;
		do {
			ssize_t res = write(fd, cur, remaining);
			if (res == -1) { perror("write"); exit(1); }
			cur += res;
			remaining -= res;
		} while (remaining > 0);
		exit(0);
	default: return pid;
	}
}

void empipe(int pipefd[2]) {
	CKE(pipe(pipefd), "pipe");
	fcntl(pipefd[0], F_SETFD, FD_CLOEXEC);
	fcntl(pipefd[1], F_SETFD, FD_CLOEXEC);
}

char *run(int count, char **progs[], size_t inputlen, char *inputstr, size_t *outlen) {
	if (inputstr) count++;
	int (*pipefds)[2] = malloc(count * sizeof(int[2]));
	pid_t *pids = malloc(count * sizeof(pid_t));
	for (int i = 0; i < count; i++) empipe(pipefds[i]);
	for (int i = 0; i < count; i++) {
		int input  = (i == 0) ? -1 : pipefds[i-1][0];
		int output = pipefds[i][1];
		if (!inputstr) pids[i] = filefork(progs[i  ], input, output);
		else if (i==0) pids[i] = echofork(inputstr, inputlen, output);
		     else      pids[i] = filefork(progs[i-1], input, output);

		DP("[%d] %d\n",i,pids[i]);
	}
	// pipes need to be closed everywhere they aren't being used.
	// close them in the parent process here; CLOEXEC will take care of them everwhere else
	// (duped fds won't have CLOEXEC set so those ones, ie the ones being used, won't be closed)
	for (int i = 0; i < count-1; i++) {
		CKE(close(pipefds[i][0]), "close");
		CKE(close(pipefds[i][1]), "close");
	}
	CKE(close(pipefds[count-1][1]), "close");

	// if the program produces too much data, you will be killed
	size_t readed = 0;
	size_t cap = 1024*1024;
	char *buf = malloc(cap*sizeof(char));
	size_t remaining = cap;
	size_t amt = 0;
	do {
		amt = read(pipefds[count-1][0], buf+readed, remaining);
		readed += amt;
		remaining -= amt;
		if (remaining == 0) {
			buf = realloc(buf, cap + cap);
			remaining += cap;
			cap += cap;
		}
	} while (amt != 0);
	CKE(close(pipefds[count-1][0]), "close");

	for (int i = 0; i < count; i++) {
		int status;
		pid_t rpid = waitpid(pids[i], &status, 0);
		DP("-> %d %d\n",rpid,WEXITSTATUS(status));
	}

	free(pipefds);
	free(pids);

	*outlen = readed;
	return buf;
}

#define FAIL(msg) do { luaL_where(L, 1); lua_pushliteral(L, msg); lua_concat(L, 2); fail = 1; goto finish; } while (0)
int do_the_thing(lua_State *L) {
	int fail = 0;
	char ***argvs = NULL;
	char *outbuf = NULL;
	size_t nprogs = 0;
	if (lua_gettop(L) != 1) FAIL("run wants exactly one argument");
	if (lua_type(L, 1) != LUA_TTABLE) FAIL("run needs table argumnt");
	nprogs = lua_rawlen(L, 1);
	if (nprogs == 0) FAIL("need at least one program to run");
	argvs = calloc(nprogs, sizeof(char**));
	for (int i = 0; i < nprogs; i++) {
		int ty = lua_rawgeti(L, 1, i + 1);
		if (ty != LUA_TTABLE) FAIL("program arglist must be table");
		size_t argc = lua_rawlen(L, 2);
		if (argc == 0) FAIL("need at least one argument to program");
		argvs[i] = calloc(argc + 1, sizeof(char*));
		for (int j = 0; j < argc; j++) {
			int ty = lua_rawgeti(L, 2, j + 1);
			if (ty != LUA_TSTRING) FAIL("program argument must be string");
			size_t slen = 0; char *str = lua_tolstring(L, 3, &slen);
			argvs[i][j] = calloc(slen + 1, sizeof(char));
			memcpy(argvs[i][j], str, slen);
			argvs[i][j][slen] = '\0'; // just to make sure
			lua_pop(L, 1);
		}
		argvs[i][argc] = NULL; // just to make sure
		lua_pop(L, 1);
	}

	char *inputstr = NULL; size_t inputlen = 0;
	lua_pushliteral(L, "input");
	int ty = lua_rawget(L, 1);
	switch (ty) {
	case LUA_TNIL: break;
	case LUA_TSTRING:
		inputstr = lua_tolstring(L, 2, &inputlen);
		break;
	default: FAIL("input field must be string if present");
	}

	size_t outlen;
	outbuf = run(nprogs, argvs, inputlen, inputstr, &outlen);
	lua_pushlstring(L, outbuf, outlen);

finish:
	if (argvs != NULL) {
		for (int i = 0; i < nprogs; i++) {
			if (argvs[i] != NULL) {
				for (int j = 0; argvs[i][j] != NULL; j++) {
					free(argvs[i][j]);
				}
				free(argvs[i]);
			}
		}
		free(argvs);
	}
	free(outbuf);
	if (fail) return lua_error(L);
	return 1;
}

int convienience_wrapper(lua_State *L) {
	luaL_checktype(L, 1, LUA_TTABLE);
	if (lua_gettop(L) != 1) return luaL_error(L, "run needs exactly 1 argument");
	size_t tlen = lua_rawlen(L, 1);
	if (tlen == 0) return luaL_error(L, "run's argument can't have 0 length");
	int ty = lua_rawgeti(L, 1, 1);
	lua_pop(L, 1);
	if (ty == LUA_TSTRING) {
		// x1  = {"a","b","c",input="ahaha"}
		// x2 := {x1,input="ahaha"}
		lua_newtable(L);
		lua_pushvalue(L, 1);
		lua_rawseti(L, 2, 1);
		lua_pushliteral(L, "input");
		lua_pushliteral(L, "input");
		lua_rawget(L, 1);
		lua_rawset(L, 2);
		lua_replace(L, 1);
	}
	return do_the_thing(L);
}

int luaopen_run(lua_State *L) {
	lua_pushcfunction(L, &convienience_wrapper);
	return 1;
}