// proxy.cpp
// NOTE: This is a generic file. Actual unit tests are located in
//       unit_tests.cpp.
// By Jack Toole for CS 225 spring 2011

// For strsignal:
#ifndef _GNU_SOURCE
	#define _GNU_SOURCE
#endif

#include <signal.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/time.h>

#include <iomanip>

#include "memcheck.h"
#include "monad_shared.h"
#include "pipestream.h"
#include "proxy.h"
#include "util.h"
#include "valgrind.h"

using std::string;
using std::vector;
using std::pair;

using namespace util;
using namespace monad_shared;

namespace proxy
{
	vector<unit_test> * global_tests = NULL;
	output_check_map * global_output_checks = NULL;
}


OUTPUT_CHECK(equals)
{
	return output == expected;
}


OUTPUT_CHECK(contains)
{
	return output.find(expected) != string::npos;
}


int main(int argc, char ** argv)
{
	using namespace proxy;

	// Set up run-time environment
	RunTimeEnvironment env(global_tests, global_output_checks);

	// Set up the tests
	RunTests runner(argc, argv, env);

	// Execute
	return runner.execute();
}


namespace proxy {

int8_t RunTimeEnvironment::singleton = 0;

// class add_unit_test
add_unit_test::add_unit_test(int8_t MPpart, const char * name, unit_test::function func,
                             int32_t points_in_part, int32_t points_in_total, long timeout,
                             bool is_valgrind)
{
	assertMPpart(MPpart, name);
	lazy_init_global_tests();
	int32_t points = get_points(points_in_total, points_in_part);
	// Add to global tests vector
	global_tests->push_back(unit_test(MPpart, name, func, points, timeout, is_valgrind));
}


// Check to make global tests vector
void add_unit_test::lazy_init_global_tests()
{
	if (global_tests == NULL)
		global_tests = new std::vector<unit_test>;
}
		
// Check that test was legal to compile
void add_unit_test::assertMPpart(int8_t MPpart, const char * name)
{
	if (!MP_PART(MPpart))
	{
		std::cerr << "Internal Error: unit tests should be surrounded by" << std::endl
			 << "#if MP_PART(partNumber) ... #endif" << std::endl
			 << name << "()" << " is defined for part " << (int)MPpart
			 << " but compiled for part " << MP_PART_NUMBER << std::endl;
		exit(-2);
	}
}


// Discriminate which points value to add
int32_t add_unit_test::get_points(int32_t points_in_total, int32_t points_in_part)
{
	#if MP_PART(NO_MP_PART)
		return points_in_total;
	#else
		return points_in_part;
	#endif
}



// class add_output_check
add_output_check::add_output_check(const char * name, output_check func)
{
	if (global_output_checks == NULL)
		global_output_checks = new output_check_map;
	(*global_output_checks)[name] = func;
}



// class Run_Time_Environment
RunTimeEnvironment::RunTimeEnvironment(vector<unit_test> *& init_tests,
                                       output_check_map *& init_output_checks)
	: itimer_number0(ITIMER_PROF),
	  itimer_number1(ITIMER_REAL),
	  timeout_signum0(SIGPROF),
	  timeout_signum1(SIGALRM),
	  max_output_length(8*1024), //arbitrary
	  single_test_passed_string("Result: passed"),
	  heap_tests(init_tests),
	  output_checks(init_output_checks)
{
	// Copy globals to the RunTimeEnvironment space
	// And remove them from the global scope
	EXIT_IF_ERROR(singleton++ != 0, "There may only be one runtime environment");
	EXIT_IF_ERROR(heap_tests == NULL, "No test cases found");
	if (output_checks == NULL)
		output_checks = new output_check_map;
	
	init_tests = NULL;
	init_output_checks = NULL;
}


int RunTimeEnvironment::cleanup_globals()
{
	if (heap_tests    != NULL) delete heap_tests;
	if (output_checks != NULL) delete output_checks;
	heap_tests    = NULL;
	output_checks = NULL;
	return 0;
}



// class RunTests
RunTests::RunTests(int argc, char ** argv, RunTimeEnvironment & env)
	: environment(env)
{
	process_args(argc, argv); // sets up mode and test_arg
	redirect_glibc_to_stderr();
}


void RunTests::redirect_glibc_to_stderr()
{
	// Turn off glibc errors default write-to-terminal behaviour, because
	// it does not get caught by stderr. This instead results in an abort.
	// Unfortunately, this has still-reachable memory leaks under valgrind
	if (RUNNING_ON_VALGRIND == 0)
		setenv("LIBC_FATAL_STDERR_","1",1);
		//setenv("MALLOC_CHECK_","2",1);
}


int32_t RunTests::execute()
{
	int32_t return_code = execute_by_mode();
	environment.cleanup_globals();
	return return_code;
}


int32_t RunTests::execute_by_mode()
{
	if (mode == SINGLE_TEST)
		return run_single_test(test_arg);
	else // if (mode == ALL_TESTS)
		return run_all_tests();
}


void RunTests::process_args(int argc, char ** argv)
{
	if (argc > 2)
	{
		cout << "Usage: " << argv[0] << "[testname]" << endl;
		exit(0);
	}

	if (argc == 2 && strcasecmp(argv[1], "--info") == 0)
	{
		printInfo();
		exit(0);
	}

	if (argc == 1 || strcmp(argv[1], "all") == 0)
		mode = ALL_TESTS;
	else
	{
		mode = SINGLE_TEST;
		test_arg = argv[1];
	}
}


int32_t RunTests::run_single_test(const char * testname)
{
	vector<unit_test> & tests = *environment.heap_tests;

	for (size_t test_i = 0; test_i < tests.size(); test_i++)
		if (strcmp(tests[test_i].name, testname) == 0)
			return run_single_test(tests[test_i]);

	cout << "Test not found" << endl;
	exit(-1);
}


int32_t RunTests::run_single_test(unit_test & curr_test)
{
	cout << "Running " << curr_test.name << " [worth "
		 << curr_test.points << " points, output below]" << endl;

	bool is_parent_process = UnitTestContainer(curr_test, environment).execute(false);
	if (!is_parent_process)
		return environment.cleanup_globals();

	string & error  = curr_test.errormsg;
	
	handle_single_test_output(curr_test.output);

	if (error == "")
		error = "Unexpectedly Aborted";

	if (curr_test.passed())
		cout << environment.single_test_passed_string << endl;
	else
		cout << "Result: FAILED:" << endl << error << endl;

	return curr_test.valgrind_flags;
}


void RunTests::handle_single_test_output(const string & output)
{
	if (output != "")
	{
		cout << output;
		if (output[output.size()-1] != '\n')
			cout << endl;
	}
}


int RunTests::run_all_tests()
{
	vector<unit_test> & tests = *environment.heap_tests;

	output::header("Running tests");

	int32_t points_sum = get_sum_points();
	int32_t max_testname_len = get_max_testname_length();
	int32_t max_points_len   = get_max_points_length();

	if (points_sum < 100)
		output::warning("Unit test scores sum to " + to_string(points_sum) + 
		                ", should be at least 100");

	int32_t score = 0;
	for (size_t test_i = 0; test_i < tests.size(); test_i++)
	{
		unit_test & curr_test = tests[test_i];
		output::testname(curr_test, max_testname_len, max_points_len);

		bool is_parent_process = UnitTestContainer(curr_test, environment).execute(true);
		// Check for the child process
		// This is unfortunately necessary (instead of an exit) to clean up
		// all the memory in use in main and the global space for valgrind
		if (!is_parent_process)
			return environment.cleanup_globals();

		// Check for success
		if (curr_test.passed())
			score += curr_test.points;
		output_single_test_passfail(curr_test);
	}

	cout << endl << endl;
	output_detailed_info_if_any_failed(score);
	output::total_score(score);
	
	return score;
}

int32_t RunTests::get_sum_points()
{
	vector<unit_test> & tests = *environment.heap_tests;
	int32_t points_sum = 0;
	for (size_t test_i = 0; test_i < tests.size(); test_i++)
		points_sum += tests[test_i].points;
	return points_sum;
}

int32_t RunTests::get_max_testname_length()
{
	vector<unit_test> & tests = *environment.heap_tests;
	int32_t max_testname_len = 0;
	for (size_t test_i = 0; test_i < tests.size(); test_i++)
	{
		int32_t currlen = strlen(tests[test_i].name) + (int)tests[test_i].is_valgrind * 11; // strlen(" (valgrind)");

		if (currlen > max_testname_len)
			max_testname_len = currlen;
	}
	return max_testname_len;
}

int32_t RunTests::get_max_points_length()
{
	vector<unit_test> & tests = *environment.heap_tests;
	int32_t max_points_len = 0;
	for (size_t test_i = 0; test_i < tests.size(); test_i++)
	{
		if (tests[test_i].points >= 100)
			max_points_len = 3;
		else if (tests[test_i].points >= 10)
			max_points_len = 2;
	}
	return max_points_len;
}

void RunTests::output_detailed_info_if_any_failed(int32_t score)
{
	vector<unit_test> & tests = *environment.heap_tests;
	
	bool any_failed = false;
	for (size_t test_i = 0; test_i < tests.size(); test_i++)
		if (!tests[test_i].passed())
			any_failed = true;
	
	if (any_failed)
		output_detailed_tests_info(score);
}
	
	
void RunTests::output_detailed_tests_info(int32_t score)
{
	output::total_score(score);
	cout << endl << endl;
	
	output::header("Detailed test output");
	
	vector<unit_test> & tests = *environment.heap_tests;
	for (size_t test_i = 0; test_i < tests.size(); test_i++)
		if (!tests[test_i].passed())
			output::detailed_info(tests[test_i]);
	
	cout << endl << "--------------------------------" << endl;
}


void RunTests::output_single_test_passfail(const unit_test & curr_test)
{
	if (curr_test.passed())
		std::cout << "passed" << endl;
	else
		std::cout << "FAILED: " << curr_test.errormsg << endl;
}



// originally from stackoverflow.com user plinth
// http://stackoverflow.com/questions/2180079/how-can-i-copy-a-file-on-unix-using-c
// Modified by Jack Toole
bool UnitTestContainer::execute(bool enable_valgrind_call)
{
	cout << std::flush;
	bool do_valgrind = enable_valgrind_call && test.is_valgrind;

	// Make our pipes
	pipestream fmsg_pipe; // For error messages
	pipestream cout_pipe; // For stdout/stderr
	pipestream nums_pipe; // for numbers: time, valgrind

	if (!do_valgrind)
		test.checkstream = new pipestream;

	// Fork
	pid_t pid;
	pid = fork();

	if (pid == 0) /* child */
	{
		fmsg_pipe.close_read();
		cout_pipe.close_read();
		nums_pipe.close_read();

		// Redirect stdout/stderr to pipe
		cout_pipe.steal_output(STDOUT_FILENO);
		cout_pipe.steal_output(STDERR_FILENO);

		if (do_valgrind)
		{
			// We're giving up control to valgrind, so we can't
			// Use anything but the cout pipe now
			fmsg_pipe.close_write();
			nums_pipe.close_write();
			child_valgrind();
		}
		else // if (!test.is_valgrind)
		{
			child_test(fmsg_pipe, nums_pipe);
		}

		// Unfortunately necessary to use a return stack instead of
		// exit() to get rid of valgrind errors
		// (which is important if we use valgrind ./proxy recursively)
		return false; // previously exit(0);
	}
	else if (pid < 0)
	{
		perror("Abort: " __FILE__ ":" STR(__LINE__));
		exit(-1);
	}
	else
	{
		/* parent - wait for child - this has all error handling, you
		 * could just call wait() as long as you are only expecting to
		 * have one child process at a time.
		 */

		fmsg_pipe.close_write();
		cout_pipe.close_write();
		nums_pipe.close_write();
		if (test.checkstream != NULL)
			test.checkstream->close_write();

		// Read stdout/stderr pipe while process is running
		cout_pipe >> setmax(environment.max_output_length) >> test.output;
		cout_pipe.close_read();

		// Eat child process
		// Should be instant because of cout_pipe EOF
		int child_status;
		pid_t ws = waitpid( pid, &child_status, 0); //should return immediately
		EXIT_IF_ERROR(ws == -1);

		if (WIFEXITED(child_status)) /* exit code in child_status */
		{
			int8_t return_code = WEXITSTATUS(child_status);

			if (do_valgrind)
			{
				fmsg_pipe.close_read();
				nums_pipe.close_read();
				valgrind_test_exit(return_code);
			}
			else // if (!test.is_valgrind)
			{
				test_exit(fmsg_pipe, nums_pipe);
			}
			return true;
		}
		else if (WIFSIGNALED(child_status)) /* killed */
		{
			fmsg_pipe.close_read();
			nums_pipe.close_read();
			test_signaled(WTERMSIG(child_status));
			return true;
		}
		else
		{
			fmsg_pipe.close_read();
			nums_pipe.close_read();
			test.errormsg = "Unexpectedly Aborted";
			return true;
		}
	}
}


void UnitTestContainer::child_valgrind()
{
	start_timeout();
	execl("/usr/bin/valgrind", "/usr/bin/valgrind", "--trace-children=yes", /*"--log-fd=-1",*/ "-q", "./proxy", test.name, NULL);
	
	// Execl failed
	cerr << "valgrind execute failed" << endl;
	exit(-1);
}


void UnitTestContainer::child_test(pipestream & fmsg_pipe, pipestream & nums_pipe)
{
	test.checkstream->close_read();
	// Execute test
	start_timeout();
	string * error_msg = new unit_test::return_type(test.func(test)); // execute function
	long test_time = end_timeout();

	// Write failure message to pipe
	fmsg_pipe << *error_msg;
	fmsg_pipe.close();

	// write time and valgrind flags to pipe
	bool test_failed = (*error_msg != unit_test::pass_string);
	delete error_msg;
	delete test.checkstream;
	environment.cleanup_globals();
	int32_t valgrind_flags = get_valgrind_flags(test_failed);
	nums_pipe << test_time;
	nums_pipe << valgrind_flags;
	nums_pipe.close();
}


void UnitTestContainer::valgrind_test_exit(int8_t return_code)
{
	size_t last_endl = findNthLast(test.output, '\n', 2);
	if (last_endl == string::npos)
		test.errormsg = "Test did not complete";
	else
	{
		test.errormsg = test.output.substr(last_endl + 1,
							test.output.length() - last_endl - 2);

		if (test.errormsg == "")
			test.errormsg = "Exception Thrown / Aborted";

		test.valgrind_flags = return_code;
		if (test.errormsg == environment.single_test_passed_string)
			test.errormsg = get_valgrind_string(test.valgrind_flags);
	}
}


void UnitTestContainer::test_exit(pipestream & fmsg_pipe, pipestream & nums_pipe)
{
	fmsg_pipe >> test.errormsg;
	fmsg_pipe.close();
	nums_pipe >> test.time;
	nums_pipe >> test.valgrind_flags;
	nums_pipe.close();
	
	// Check for output's correctness, if that was a condition of passing
	if (test.passed())
	{
		while (!test.checkstream->eof())
		{
			string checkname;
			string checkstr;
			*test.checkstream >> checkname;
			if (test.checkstream->eof()) break;
			*test.checkstream >> checkstr;
			if (test.checkstream->eof()) break;

			output_check check_function = (*environment.output_checks)[checkname];
			if (check_function == NULL)
			{
				cerr << "Internal Error: in test " << test.name << ": "
				     << checkname << " is not a registered OUTPUT_CHECK function" << endl;
				exit(-2);
			}

			if (!check_function(test.output, checkstr))
				test.errormsg = "Incorrect Terminal Output";
		}
	}
	delete test.checkstream;
}


void UnitTestContainer::test_signaled(int signum)
{
	if (signum == environment.timeout_signum0 || signum == environment.timeout_signum1)
	{
		test.errormsg = string("Timed out") + " (" + to_string(test.timeout) + "ms)";
		test.time = test.timeout;
	}
	else
		test.errormsg = strsignal(signum);
}


int32_t get_valgrind_flags(bool test_failed)
{
	// Check for valgrind errors or leaks (if running under valgrind)
	unsigned long errors     = 0;
	unsigned long leaked     = 0;
	unsigned long dubious    = 0;
	unsigned long reachable  = 0;
	unsigned long suppressed = 0;

	errors = VALGRIND_COUNT_ERRORS;
	VALGRIND_DO_LEAK_CHECK; //QUICK
	VALGRIND_COUNT_LEAK_BLOCKS(leaked, dubious, reachable, suppressed);

	return bitflags(test_failed, errors, leaked, dubious, reachable);
}


const char * get_valgrind_string(int32_t flags)
{
	if (flags == 0) return unit_test::pass_string;

	bool test_failed = bitflag(flags, 0);
	bool errors      = bitflag(flags, 1);
	bool leaked      = bitflag(flags, 2);
	bool dubious     = bitflag(flags, 3);
	bool reachable   = bitflag(flags, 4);

	if (test_failed) return "Test failed (see output)";
	if (errors)      return "Invalid read/write errors";
	if (leaked)      return "Directly lost memory leaks";
	if (dubious)     return "Possibly lost memory leaks";
	if (reachable)   return "Still-reachable memory leaks";
	return "Unknown memory errors";
}


void UnitTestContainer::start_timeout()
{
	struct itimerval timeout;
	timeout.it_interval.tv_sec  = 0;
	timeout.it_interval.tv_usec = 0;
	timeout.it_value.tv_sec  = test.timeout/1000;
	timeout.it_value.tv_usec = (test.timeout%1000) * 1000;

	EXIT_IF_ERROR(setitimer(environment.itimer_number0, &timeout, NULL));

	// second real time signal in case the student calls a blocking call
	timeout.it_value.tv_sec *= 10;
	EXIT_IF_ERROR(setitimer(environment.itimer_number1, &timeout, NULL));
}


long UnitTestContainer::end_timeout()
{
	struct itimerval timeout;
	timeout.it_interval.tv_sec  = 0;
	timeout.it_interval.tv_usec = 0;
	timeout.it_value.tv_sec  = 0;
	timeout.it_value.tv_usec = 0;
	struct itimerval remaining;

	EXIT_IF_ERROR(setitimer(environment.itimer_number0, &timeout, &remaining));
	EXIT_IF_ERROR(setitimer(environment.itimer_number1, &timeout, NULL));

	// There seems to be a strange -1 error here. I may just be tired,
	// but I can't figure out why right now
	long time = test.timeout - remaining.it_value.tv_sec*1000 - remaining.it_value.tv_usec/1000;
	return (time < 0) ? 0 : time;
}



UnitTestContainer::UnitTestContainer(unit_test & _test, RunTimeEnvironment & env)
	: environment(env), test(_test) { }


	

} // namespace proxy