#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <unistd.h>
#include <asm/unistd.h>
#include <sys/ptrace.h>
#include <sys/wait.h>
#include <sys/resource.h>
#include <sys/user.h>
#include <fcntl.h>
#include <cstring>
#include <string>
#include <vector>
#include <set>
#include <argp.h>
#include "uoj_env.h"
using namespace std;

struct RunResult {
	int result;
	int ust;
	int usm;
	int exit_code;

	RunResult(int _result, int _ust = -1, int _usm = -1, int _exit_code = -1)
			: result(_result), ust(_ust), usm(_usm), exit_code(_exit_code) {
		if (result != RS_AC) {
			ust = -1, usm = -1;
		}
	}
};

struct RunProgramConfig
{
	int time_limit;
	int real_time_limit;
	int memory_limit;
	int output_limit;
	int stack_limit;
	string input_file_name;
	string output_file_name;
	string error_file_name;
	string result_file_name;
	string work_path;
	string type;
	vector<string> extra_readable_files;
	vector<string> extra_writable_files;
	bool allow_proc;
	bool safe_mode;
	bool need_show_trace_details;

	string program_name;
	string program_basename;
	vector<string> argv;

};

int put_result(string result_file_name, RunResult res) {
	FILE *f;
	if (result_file_name == "stdout") {
		f = stdout;
	} else if (result_file_name == "stderr") {
		f = stderr;
	} else {
		f = fopen(result_file_name.c_str(), "w");
	}
	fprintf(f, "%d %d %d %d\n", res.result, res.ust, res.usm, res.exit_code);
	if (f != stdout && f != stderr) {
		fclose(f);
	}
	if (res.result == RS_JGF) {
		return 1;
	} else {
		return 0;
	}
}

char self_path[PATH_MAX + 1] = {};

#include "run_program_conf.h"

argp_option run_program_argp_options[] =
{
	{"tl"                 , 'T', "TIME_LIMIT"  , 0, "Set time limit (in second)"                            ,  1},
	{"rtl"                , 'R', "TIME_LIMIT"  , 0, "Set real time limit (in second)"                       ,  2},
	{"ml"                 , 'M', "MEMORY_LIMIT", 0, "Set memory limit (in mb)"                              ,  3},
	{"ol"                 , 'O', "OUTPUT_LIMIT", 0, "Set output limit (in mb)"                              ,  4},
	{"sl"                 , 'S', "STACK_LIMIT" , 0, "Set stack limit (in mb)"                               ,  5},
	{"in"                 , 'i', "IN"          , 0, "Set input file name"                                   ,  6},
	{"out"                , 'o', "OUT"         , 0, "Set output file name"                                  ,  7},
	{"err"                , 'e', "ERR"         , 0, "Set error file name"                                   ,  8},
	{"work-path"          , 'w', "WORK_PATH"   , 0, "Set the work path of the program"                      ,  9},
	{"type"               , 't', "TYPE"        , 0, "Set the program type (for some program such as python)", 10},
	{"res"                , 'r', "RESULT_FILE" , 0, "Set the file name for outputing the result            ", 10},
	{"add-readable"       , 500, "FILE"        , 0, "Add a readable file"                                   , 11},
	{"add-writable"       , 505, "FILE"        , 0, "Add a writable file"                                   , 11},
	{"unsafe"             , 501, 0             , 0, "Don't check dangerous syscalls"                        , 12},
	{"show-trace-details" , 502, 0             , 0, "Show trace details"                                    , 13},
	{"allow-proc"         , 503, 0             , 0, "Allow fork, exec... etc."                              , 14},
	{"add-readable-raw"   , 504, "FILE"        , 0, "Add a readable (don't transform to its real path)"     , 15},
	{"add-writable-raw"   , 506, "FILE"        , 0, "Add a writable (don't transform to its real path)"     , 15},
	{0}
};
error_t run_program_argp_parse_opt (int key, char *arg, struct argp_state *state)
{
	RunProgramConfig *config = (RunProgramConfig*)state->input;

	switch (key)
	{
		case 'T':
			config->time_limit = atoi(arg);
			break;
		case 'R':
			config->real_time_limit = atoi(arg);
			break;
		case 'M':
			config->memory_limit = atoi(arg);
			break;
		case 'O':
			config->output_limit = atoi(arg);
			break;
		case 'S':
			config->stack_limit = atoi(arg);
			break;
		case 'i':
			config->input_file_name = arg;
			break;
		case 'o':
			config->output_file_name = arg;
			break;
		case 'e':
			config->error_file_name = arg;
			break;
		case 'w':
			config->work_path = realpath(arg);
			if (config->work_path.empty()) {
				argp_usage(state);
			}
			break;
		case 'r':
			config->result_file_name = arg;
			break;
		case 't':
			config->type = arg;
			break;
		case 500:
			config->extra_readable_files.push_back(realpath(arg));
			break;
		case 501:
			config->safe_mode = false;
			break;
		case 502:
			config->need_show_trace_details = true;
			break;
		case 503:
			config->allow_proc = true;
			break;
		case 504:
			config->extra_readable_files.push_back(arg);
			break;
		case 505:
			config->extra_writable_files.push_back(realpath(arg));
			break;
		case 506:
			config->extra_writable_files.push_back(arg);
			break;
		case ARGP_KEY_ARG:
			config->argv.push_back(arg);
			for (int i = state->next; i < state->argc; i++) {
				config->argv.push_back(state->argv[i]);
			}
			state->next = state->argc;
			break;
		case ARGP_KEY_END:
			if (state->arg_num == 0) {
				argp_usage(state);
			}
			break;
		default:
			return ARGP_ERR_UNKNOWN;
	}
	return 0;
}
char run_program_argp_args_doc[] = "program arg1 arg2 ...";
char run_program_argp_doc[] = "run_program: a tool to run program safely";

argp run_program_argp = {
	run_program_argp_options,
	run_program_argp_parse_opt,
	run_program_argp_args_doc,
	run_program_argp_doc
};

RunProgramConfig run_program_config;

void parse_args(int argc, char **argv) {
	run_program_config.time_limit = 1;
	run_program_config.real_time_limit = -1;
	run_program_config.memory_limit = 256;
	run_program_config.output_limit = 64;
	run_program_config.stack_limit = 1024;
	run_program_config.input_file_name = "stdin";
	run_program_config.output_file_name = "stdout";
	run_program_config.error_file_name = "stderr";
	run_program_config.work_path = "";
	run_program_config.result_file_name = "stdout";
	run_program_config.type = "default";
	run_program_config.safe_mode = true;
	run_program_config.need_show_trace_details = false;
	run_program_config.allow_proc = false;

	argp_parse(&run_program_argp, argc, argv, ARGP_NO_ARGS | ARGP_IN_ORDER, 0, &run_program_config);

	if (run_program_config.real_time_limit == -1)
		run_program_config.real_time_limit = run_program_config.time_limit + 2;
	run_program_config.stack_limit = min(run_program_config.stack_limit, run_program_config.memory_limit);

	if (!run_program_config.work_path.empty()) {
		if (chdir(run_program_config.work_path.c_str()) == -1) {
			exit(put_result(run_program_config.result_file_name, RS_JGF));
		}
	}

	if (run_program_config.type == "java8" || run_program_config.type == "java11") {
		run_program_config.program_name = run_program_config.argv[0];
	} else {
		run_program_config.program_name = realpath(run_program_config.argv[0]);
	}
	if (run_program_config.work_path.empty()) {
		run_program_config.work_path = dirname(run_program_config.program_name);
		run_program_config.program_basename = basename(run_program_config.program_name);
		run_program_config.argv[0] = "./" + run_program_config.program_basename;

		if (chdir(run_program_config.work_path.c_str()) == -1) {
			exit(put_result(run_program_config.result_file_name, RS_JGF));
		}
	}

	if (run_program_config.type == "python2") {
		string pre[4] = {"/usr/bin/python2", "-E", "-s", "-B"};
		run_program_config.argv.insert(run_program_config.argv.begin(), pre, pre + 4);
	} else if (run_program_config.type == "python3") {
		string pre[3] = {"/usr/bin/python3", "-I", "-B"};
		run_program_config.argv.insert(run_program_config.argv.begin(), pre, pre + 3);
	} else if (run_program_config.type == "java8") {
		string pre[3] = {"/usr/lib/jvm/java-8-openjdk-amd64/bin/java", "-Xmx1024m", "-Xss1024m"};
		run_program_config.argv.insert(run_program_config.argv.begin(), pre, pre + 3);
	} else if (run_program_config.type == "java11") {
		string pre[3] = {"/usr/lib/jvm/java-11-openjdk-amd64/bin/java", "-Xmx1024m", "-Xss1024m"};
		run_program_config.argv.insert(run_program_config.argv.begin(), pre, pre + 3);
	}
}

void set_limit(int r, int rcur, int rmax = -1)  {
	if (rmax == -1)
		rmax = rcur;
	struct rlimit l;
	if (getrlimit(r, &l) == -1) {
		exit(55);
	}
	l.rlim_cur = rcur;
	l.rlim_max = rmax;
	if (setrlimit(r, &l) == -1) {
		exit(55);
	}
}
void run_child() {
	set_limit(RLIMIT_CPU, run_program_config.time_limit, run_program_config.real_time_limit);
	set_limit(RLIMIT_FSIZE, run_program_config.output_limit << 20);
	set_limit(RLIMIT_STACK, run_program_config.stack_limit << 20);

	if (run_program_config.input_file_name != "stdin") {
		if (freopen(run_program_config.input_file_name.c_str(), "r", stdin) == NULL) {
			exit(11);
		}
	}
	if (run_program_config.output_file_name != "stdout" && run_program_config.output_file_name != "stderr") {
		if (freopen(run_program_config.output_file_name.c_str(), "w", stdout) == NULL) {
			exit(12);
		}
	}
	if (run_program_config.error_file_name != "stderr") {
		if (run_program_config.error_file_name == "stdout") {
			if (dup2(1, 2) == -1) {
				exit(13);
			}
		} else {
			if (freopen(run_program_config.error_file_name.c_str(), "w", stderr) == NULL) {
				exit(14);
			}
		}
		
		if (run_program_config.output_file_name == "stderr") {
			if (dup2(2, 1) == -1) {
				exit(15);
			}
		}
	}

	char *env_path_str = getenv("PATH");
	char *env_lang_str = getenv("LANG");
	char *env_shell_str = getenv("SHELL");
	string env_path = env_path_str ? env_path_str : "";
	string env_lang = env_lang_str ? env_lang_str : "";
	string env_shell = env_shell_str ? env_shell_str : "";

	clearenv();
	setenv("USER", "poor_program", 1);
	setenv("LOGNAME", "poor_program", 1);
	setenv("HOME", run_program_config.work_path.c_str(), 1);
	if (env_lang_str) {
		setenv("LANG", env_lang.c_str(), 1);
	}
	if (env_path_str) {
		setenv("PATH", env_path.c_str(), 1);
	}
	setenv("PWD", run_program_config.work_path.c_str(), 1);
	if (env_shell_str) {
		setenv("SHELL", env_shell.c_str(), 1);
	}

	char **program_c_argv = new char*[run_program_config.argv.size() + 1];
	for (size_t i = 0; i < run_program_config.argv.size(); i++) {
		program_c_argv[i] = new char[run_program_config.argv[i].size() + 1];
		strcpy(program_c_argv[i], run_program_config.argv[i].c_str());
	}
	program_c_argv[run_program_config.argv.size()] = NULL;

	if (ptrace(PTRACE_TRACEME, 0, NULL, NULL) == -1) {
		exit(16);
	}
	if (execv(program_c_argv[0], program_c_argv) == -1) {
		exit(17);
	}
}

const int MaxNRPChildren = 50;
struct rp_child_proc {
	pid_t pid;
	int mode;
};
int n_rp_children;
pid_t rp_timer_pid;
rp_child_proc rp_children[MaxNRPChildren];

int rp_children_pos(pid_t pid) {
	for (int i = 0; i < n_rp_children; i++) {
		if (rp_children[i].pid == pid) {
			return i;
		}
	}
	return -1;
}
int rp_children_add(pid_t pid) {
	if (n_rp_children == MaxNRPChildren) {
		return -1;
	}
	rp_children[n_rp_children].pid = pid;
	rp_children[n_rp_children].mode = -1;
	n_rp_children++;
	return 0;
}
void rp_children_del(pid_t pid) {
	int new_n = 0;
	for (int i = 0; i < n_rp_children; i++) {
		if (rp_children[i].pid != pid) {
			rp_children[new_n++] = rp_children[i];
		}
	}
	n_rp_children = new_n;
}

void stop_child(pid_t pid) {
	kill(pid, SIGKILL);
}
void stop_all() {
	kill(rp_timer_pid, SIGKILL);
	for (int i = 0; i < n_rp_children; i++) {
		kill(rp_children[i].pid, SIGKILL);
	}
}

RunResult trace_children() {
	rp_timer_pid = fork();
	if (rp_timer_pid == -1) {
		stop_all();
		return RunResult(RS_JGF);
	} else if (rp_timer_pid == 0) {
		struct timespec ts;
		ts.tv_sec = run_program_config.real_time_limit;
		ts.tv_nsec = 0;
		nanosleep(&ts, NULL);
		exit(0);
	}

	if (run_program_config.need_show_trace_details) {
		cerr << "timerpid " << rp_timer_pid << endl;
	}

	pid_t prev_pid = -1;
	while (true) {
		int stat = 0;
		int sig = 0;
		struct rusage ruse;
		
		pid_t pid = wait4(-1, &stat, __WALL, &ruse);
		if (run_program_config.need_show_trace_details) {
			if (prev_pid != pid) {
				cerr << "----------" << pid << "----------" << endl;
			}
			prev_pid = pid;
		}
		if (pid == rp_timer_pid) {
			if (WIFEXITED(stat) || WIFSIGNALED(stat)) {
				stop_all();
				return RunResult(RS_TLE);
			}
			continue;
		}
		
		int p = rp_children_pos(pid);
		if (p == -1) {
			if (run_program_config.need_show_trace_details) {
				fprintf(stderr, "new_proc  %lld\n", (long long int)pid);
			}
			if (rp_children_add(pid) == -1) {
				stop_child(pid);
				stop_all();
				return RunResult(RS_DGS);
			}
			p = n_rp_children - 1;
		}

		int usertim = ruse.ru_utime.tv_sec * 1000 + ruse.ru_utime.tv_usec / 1000;
		int usermem = ruse.ru_maxrss;
		if (usertim > run_program_config.time_limit * 1000) {
			stop_all();
			return RunResult(RS_TLE);
		}
		if (usermem > run_program_config.memory_limit * 1024) {
			stop_all();
			return RunResult(RS_MLE);
		}

		if (WIFEXITED(stat)) {
			if (run_program_config.need_show_trace_details) {
				fprintf(stderr, "exit     : %d\n", WEXITSTATUS(stat));
			}
			if (rp_children[0].mode == -1) {
				stop_all();
				return RunResult(RS_JGF, -1, -1, WEXITSTATUS(stat));
			} else {
				if (pid == rp_children[0].pid) {
					stop_all();
					return RunResult(RS_AC, usertim, usermem, WEXITSTATUS(stat));
				} else {
					rp_children_del(pid);
					continue;
				}
			}
		}

		if (WIFSIGNALED(stat)) {
			if (run_program_config.need_show_trace_details) {
				fprintf(stderr, "sig exit : %d\n", WTERMSIG(stat));
			}
			if (pid == rp_children[0].pid) {
				switch(WTERMSIG(stat)) {
				case SIGXCPU: // nearly impossible
					stop_all();
					return RunResult(RS_TLE);
				case SIGXFSZ:
					stop_all();
					return RunResult(RS_OLE);
				default:
					stop_all();
					return RunResult(RS_RE);
				}
			} else {
				rp_children_del(pid);
				continue;
			}
		}
		
		if (WIFSTOPPED(stat)) {
			sig = WSTOPSIG(stat);
			
			if (rp_children[p].mode == -1) {
				if ((p == 0 && sig == SIGTRAP) || (p != 0 && sig == SIGSTOP)) {
					if (p == 0) {
						int ptrace_opt = PTRACE_O_EXITKILL | PTRACE_O_TRACESYSGOOD;
						if (run_program_config.safe_mode) {
							ptrace_opt |= PTRACE_O_TRACECLONE | PTRACE_O_TRACEFORK | PTRACE_O_TRACEVFORK;
							ptrace_opt |= PTRACE_O_TRACEEXEC;
						}
						if (ptrace(PTRACE_SETOPTIONS, pid, NULL, ptrace_opt) == -1) {
							stop_all();
							return RunResult(RS_JGF);
						}
					}
					sig = 0;
				}
				rp_children[p].mode = 0;
			} else if (sig == (SIGTRAP | 0x80)) {
				if (rp_children[p].mode == 0) {
					if (run_program_config.safe_mode) {
						if (!check_safe_syscall(pid, run_program_config.need_show_trace_details)) {
							stop_all();
							return RunResult(RS_DGS);
						}
					}
					rp_children[p].mode = 1;
				} else {
					if (run_program_config.safe_mode) {
						on_syscall_exit(pid, run_program_config.need_show_trace_details);
					}
					rp_children[p].mode = 0;
				}
				
				sig = 0;
			} else if (sig == SIGTRAP) {
				switch ((stat >> 16) & 0xffff) {
					case PTRACE_EVENT_CLONE:
					case PTRACE_EVENT_FORK:
					case PTRACE_EVENT_VFORK:
						sig = 0;
						break;
					case PTRACE_EVENT_EXEC:
						rp_children[p].mode = 1;
						sig = 0;
						break;
					case 0:
						break;
					default:
						stop_all();
						return RunResult(RS_JGF);
				}
			}

			if (sig != 0) {
				if (run_program_config.need_show_trace_details) {
					fprintf(stderr, "sig      : %d\n", sig);
				}
			}
			
			switch(sig) {
			case SIGXCPU:
				stop_all();
				return RunResult(RS_TLE);
			case SIGXFSZ:
				stop_all();
				return RunResult(RS_OLE);
			}
		}
		
		ptrace(PTRACE_SYSCALL, pid, NULL, sig);
	}
}

RunResult run_parent(pid_t pid) {
	init_conf(run_program_config);
	
	n_rp_children = 0;

	rp_children_add(pid);
	return trace_children();
}
int main(int argc, char **argv) {
	self_path[readlink("/proc/self/exe", self_path, PATH_MAX)] = '\0';
	parse_args(argc, argv);

	pid_t pid = fork();
	if (pid == -1) {
		return put_result(run_program_config.result_file_name, RS_JGF);
	} else if (pid == 0) {
		run_child();
	} else {
		return put_result(run_program_config.result_file_name, run_parent(pid));
	}
	return put_result(run_program_config.result_file_name, RS_JGF);
}