#include "compile.h"

static void _child_cleanup(asmsh_asmc_child_t *child);

asmsh_asmc_ctx_t* asmsh_asmc_ctx_default()
{
	char * const args[] = ASMSH_COMPILE_ARGS;
	return asmsh_asmc_ctx(ASMSH_COMPILE_AS, ASMSH_COMPILE_OBJ, args, 1, NULL);
}


asmsh_asmc_ctx_t* asmsh_asmc_ctx(const char *progname, \
		const char *result_path, char* const args[], int pre_spawn,
		struct timeval *ctimeout)
{

	const struct timeval default_timeout = {.tv_sec = 1, .tv_usec = 250000};
	asmsh_asmc_ctx_t *res;

	int argc; // args count
	int errno_bck; // errno will take this value in err labels
	int tmpfd;
	char *temp_name;

	int i;
	char * const *ptr;

	if((res = malloc(sizeof(*res))) == NULL)
	{
		errno_bck =  errno;
		goto err_alloc_res;
	}
	bzero(res, sizeof(res));

	if(asmsh_asmc_buildarg(progname, result_path, args,
			&res->args, &res->progname, &res->respath) < 0)
	{
		errno_bck = errno;
		goto err_buildarg;
	}

	res->child.pid = -1;
	res->pre_spawn = pre_spawn?1:0;
	if(res->pre_spawn)
	{
		asmsh_asmc_spawn(res);
	}

	memcpy(&(res->ctimeout), ctimeout?ctimeout:&default_timeout,
			sizeof(struct timeval));

	free(temp_name);
	return res;

err_buildarg:
	free(res);
err_alloc_res:
	errno = errno_bck;
	return NULL;
}


void asmsh_asmc_ctx_free(asmsh_asmc_ctx_t *ctx)
{
	pid_t cpid;
	int wstatus;
	char **ptr;

	if(ctx->child.pid > 0)
	{
		cpid = waitpid(ctx->child.pid, &wstatus, WNOHANG);
		if(cpid != -1 && cpid != 0)
		{
			if(WCOREDUMP(wstatus))
			{
				dprintf(2, "Child %d segfault\n", cpid);
				ctx->child.pid = 0;
			}
			else if(WIFEXITED(wstatus))
			{
				if(WEXITSTATUS(wstatus))
				{
					dprintf(2, "Child exited with status %d\n",
							WEXITSTATUS(wstatus));
				}
				ctx->child.pid = 0;
			}
		}
		if(ctx->child.pid > 0)
		{
			if(kill(ctx->child.pid, SIGKILL) == -1)
			{
				perror("Unable to kill child process");
			}
		}
		ctx->child.pid = 0;
		close(ctx->child.pipe_stdin);
	}
	if(ctx->respath)
	{
		unlink(ctx->respath);
	}
	
	for(ptr = ctx->args; *ptr; ptr++) { free(*ptr); }
	free(ctx->args);
	free(ctx);
}


int asmsh_asmc_compile(asmsh_asmc_ctx_t *ctx, const char *instr, asmsh_bytecode_t *res)
{
	char *reason = NULL;
	pid_t wpid;
	int wstatus;

	if(asmsh_asmc_syntax(instr, &reason) == -1)
	{
		dprintf(2, "Syntax error %s : '%s'\n",
				reason, instr);
		errno = EINVAL;
		return -1;
	}

	if(!ctx->pre_spawn && asmsh_asmc_spawn(ctx) == -1)
	{
		return -1;
	}

	wpid = waitpid(ctx->child.pid, &wstatus, WNOHANG);
	if(wpid == -1)
	{
		perror("Unable to check child state");
		ctx->child.pid = -1;
		_child_cleanup(&ctx->child);
		return -1;
	}
	else if (wpid)
	{
		if(WCOREDUMP(wstatus))
		{
			dprintf(2, "Child %d segfault\n", ctx->child.pid);
		}
		else if (WIFEXITED(wstatus))
		{
			dprintf(2, "Child %d exited with status %d\n",
					ctx->child.pid,
					WEXITSTATUS(wstatus));
		}
		else
		{
			dprintf(2, "Child %d in unknown status %d\n",
					ctx->child.pid, wstatus);
			_child_cleanup(&ctx->child);
		}
		return -1;
	}


	return asmsh_asmc_compile_unsafe(ctx, instr, res);
}

int asmsh_asmc_compile_unsafe(asmsh_asmc_ctx_t *ctx, const char *_instr,
		asmsh_bytecode_t *res)
{
	int ret;
	size_t instr_sz;
	pid_t wpid;
	int wstatus, sigfd;
	sigset_t sigfdmask, oldmask;
	fd_set sigset;
	struct timeval stimeout;
	char *instr;

	sigemptyset(&sigfdmask);
	sigaddset(&sigfdmask, SIGCHLD);
	sigfillset(&sigfdmask);

	/// TODO : early (context ?) sigprocmask ?
	if((sigprocmask(SIG_BLOCK, &sigfdmask, &oldmask) == -1))
	{
		perror("Unable to sigmask");
		goto err_early;
	}
	if((sigfd = signalfd(-1, &sigfdmask, SFD_NONBLOCK)) == -1)
	{
		perror("Unable to sigfd");
		goto err_sigfd;
	}


	instr_sz = strlen(_instr);
	instr = alloca(instr_sz+1);
	memcpy(instr, _instr, instr_sz);
	instr[instr_sz] = '\n';
	instr[instr_sz+1] = '\0';
	instr_sz++;

	ret = write(ctx->child.pipe_stdin, instr, instr_sz);
	if(ret < 0)
	{
		perror("Unable to send instruction to child process");
		goto err;
	}
	else if(ret < instr_sz)
	{
		asmsh_log_error("Unable to write the whole instruction in one write, bailout !\n");
		close(ctx->child.pipe_stdin);
		goto err;
	}
	close(ctx->child.pipe_stdin);
	close(ctx->child.pipe_stderr);

	FD_ZERO(&sigset);
	FD_SET(sigfd, &sigset);

	stimeout = ctx->ctimeout;
	do
	{
		if(select(sigfd+1, &sigset, NULL, NULL, &ctx->ctimeout) == 0)
		{
			//compilation timeout t_t
			asmsh_log_error("Compilation timeout");
			goto err;
		}

		wpid = waitpid(ctx->child.pid, &wstatus, 0);
	}while(!wpid);

	if(wpid == -1)
	{
		perror("Unable to wait for as to exit");
		goto err_stderr;
	}
	if(WCOREDUMP(wstatus))
	{
		asmsh_log_error("as segfault will compiling");
		goto err_stderr;
	}
	if(!WIFEXITED(wstatus))
	{
		asmsh_log_error("as didn't exited");
		goto err_stderr;
	}
	if(WEXITSTATUS(wstatus))
	{
		asmsh_log_error("GNU as did not exited with 0 status");
		goto err_stderr;
	}
	if((sigprocmask(SIG_BLOCK, &oldmask, NULL) == -1))
	{
		perror("Unable to restore sigmask");
		goto err;
	}

	ctx->child.pid = 0;
	if(ctx->pre_spawn && asmsh_asmc_spawn(ctx) == -1)
	{
		perror("Unable to pre-spawn");
		goto err;
	}
	
	return asmh_asmc_bytecode_from_obj(ctx->respath, res);


err_stderr:
	// TODO read stderr & log it
err:
	close(sigfd);
err_sigfd:
	sigprocmask(SIG_BLOCK, &oldmask, NULL);
err_early:
	_child_cleanup(&ctx->child);
	return -1;
}

int asmsh_asmc_spawn(asmsh_asmc_ctx_t *ctx)
{
	int fpid;
	int pipes[2][2];
	int err;
	int i;

	if(ctx->child.pid > 0)
	{
		errno = EUCLEAN;
		return -1;
	}

	for(i=0; i<2; i++)
	{
		if(pipe(pipes[i]) == -1)
		{
			goto err_pipe;
		}
	}
	fpid = fork();
	if(fpid == -1)
	{
		err = errno;
		goto err_fork;
	}
	else if (fpid == 0)
	{
		asmsh_asmc_child(pipes, ctx->progname, ctx->args);
	}

	close(pipes[0][0]);
	close(pipes[1][1]);

	ctx->child.pid = fpid;
	ctx->child.pipe_stdin = pipes[0][1];
	ctx->child.pipe_stderr = pipes[1][0];

	return 0;

err_fork:
	i=1;
err_pipe:
	for(i; i>=0; i--)
	{
		close(pipes[i][0]);
		close(pipes[i][1]);
	}
	errno = err;
	return -1;
}


void asmsh_asmc_child(const int std_pipes[2][2], const char *progname, char* const args[])
{
	close(std_pipes[0][1]);
	if(dup2(std_pipes[0][0], 0) < 0)
	{
		perror("Unable to pipe stdin");
		exit(1);
	}
	execvp(progname, args);
	perror("execvp failed");
	close(std_pipes[0][0]);
	exit(2);
}


int asmsh_asmc_syntax(const char* instr, char **reason)
{

#define explain(ptr, reason){\
if(ptr) {\
	*ptr = strdup(reason);\
	if(!*ptr) { \
		int e = errno;\
		perror("Unable to strdup to explain syntax error");\
		errno=e;return -1;}}}
	
	const unsigned char *ptr;

	switch(*instr)
	{
		case '.':
			explain(reason, "Cannot starts with '.'");
			goto err;
		default:
			break;
	}

	ptr = instr;
	while(*ptr)
	{
		switch(*ptr)
		{
			case '\n':
			case ';':
			case '#':
				explain(reason, "Contains forbidden char");
				goto err;
			default:
				break;
		}
		if(*ptr < 0x20 || *ptr > 0x7E)
		{
			explain(reason, "Contains char in unexpected value range");
			goto err;
		}
		ptr++;
	}

	*reason = NULL;
	errno = 0;
	return 0;
err:
	errno = EINVAL;
	return -1;
#undef explain
}

static void _child_cleanup(asmsh_asmc_child_t *child)
{
	if(child->pid > 0)
	{
		kill(child->pid, SIGKILL);
	}
	if(child->pipe_stdin > 0)
	{
		close(child->pipe_stdin);
	}
	if(child->pipe_stderr > 0)
	{
		close(child->pipe_stderr);
	}
	child->pid = child->pipe_stdin = child->pipe_stderr = -1;
}


int asmh_asmc_bytecode_from_obj(const char *objfile, asmsh_bytecode_t *bytecode)
{
	/*TODO parse ELF header instead of harcoded addr ?
	  	at least check if the object file has the expected
	  	size for the whole file.
	   Header parsing is better in case the ELF organisation is
	   platform/version dependent ?
	*/
	// ELF 64 constants
	const int bytecode_sz_off = 0x138;
	const int bytecode_sz_sz = 4;
	const int bytecode_off = 0x40;

	int fd, err;
	off_t off;
	ssize_t rret;
	unsigned int bytecode_sz;

	if((fd = open(objfile, O_RDONLY)) < 0)
	{
		perror("Unable to open bytecode's object file");
		return -1;
	}

	if((off = lseek(fd, bytecode_sz_off, SEEK_SET)) != bytecode_sz_off)
	{
		/// TODO check off
		perror("Unable to seek at bytecode size offset");
		err = errno;
		goto err;
	}

	if((rret = read(fd, &bytecode_sz, bytecode_sz_sz)) != bytecode_sz_sz)
	{
		err = errno;
		if(rret < 0)
		{
			perror("Unable to read bytecode size");
		}
		else
		{
			dprintf(2, "Unable to read the 4 bytes of the bytecode size\n");
			err = ENODATA;
		}
		goto err;
	}

	if(bytecode_sz == 0)
	{
		dprintf(2, "Null bytecode size");
		err = ENODATA;
		goto err;
	}
	if(bytecode_sz > sizeof(bytecode->bytes))
	{
		dprintf(2, "Bytecode size invalid, too many bytes (%d)\n",
				bytecode_sz);
		err = ENOBUFS;
		goto err;
	}
	if((off = lseek(fd, bytecode_off, SEEK_SET)) != bytecode_off)
	{
		/// TODO check off
		perror("Unable to seek at bytecode offset");
		err = errno;
		goto err;
	}

	bytecode->size = bytecode_sz;
	if((rret = read(fd, bytecode->bytes, bytecode_sz)) != bytecode_sz)
	{
		err = errno;
		if(rret < 0)
		{
			perror("Unable to read bytecode");
		}
		else
		{
			dprintf(2, "Unable to read the %d bytes of bytecode\n");
			err = ENODATA;
		}
		goto err;
	}

	close(fd);
	return 0;

err:
	close(fd);
	errno = err;
	return -1;
}


int asmsh_asmc_buildarg(const char* progname, const char *result_tpl, \
		char *const args_in[],
		char ***args_o, char **progname_o, char **respath_o)
{
	int argc, tmpfd, err, i;
	char *tempname;
	char* const *args_ptr;

	if((tempname = strdupa(result_tpl)) == NULL)
	{
		return -1;
	}
	if((tmpfd = mkstemp(tempname)) == -1)
	{
		err = errno;
		perror("Unable to create temporary result .o file");
		errno = err;
		return -1;
	}
	close(tmpfd);

	args_ptr = args_in;
	argc = 0;
	for(i=0;i<2;i++) // here we overflow if args_in was malformed
	{
		while(*args_ptr) { args_ptr++; argc++; }
		args_ptr++;
		argc++;
	}
	if((*args_o = malloc(sizeof(char*)*(argc))) == NULL)
	{
		err = errno;
		perror("unable to allocate compiler args array");
		goto err_malloc;
	}

	if(((*args_o)[0] = strdup(progname)) == NULL)
	{
		err = errno;
		perror("unable to allocate compiler first argument");
		goto err_progname;
	}
	*progname_o = (*args_o)[0];
	*respath_o = NULL;

	for(i=1; i<argc-1; i++)
	{
		char *const *tocopy;
		if(args_in[i])
		{
			tocopy = &(args_in[i]);
		}
		else
		{
			tocopy = &tempname;
		}
		(*args_o)[i] = strdup(*tocopy);
		if(!(*args_o)[i])
		{
			perror("Unable to copy compiler args");
			goto err_dupargs;
		}
		if(!args_in[i])
		{
			*respath_o = (*args_o)[i];
		}
	}
	(*args_o)[argc-1] = NULL;

	return 0;

err_dupargs:
	i--;
	while(i>0)
	{
		free((*args_o)[i]);
		i--;
	}

err_progname:
	free(args_o);
	*args_o = NULL;
err_malloc:
	unlink(tempname);
	errno = err;
	return -1;
}