#include <iostream>
#include <iomanip>
#include <sstream>
#include <string>
#include <fstream>
extern "C" {
	#include <unistd.h>
	#include <errno.h>
	#include <signal.h>
	#include <sys/types.h>
	#include <sys/time.h>
	#include <sys/resource.h>
	#include <sys/wait.h>
	#include <sys/param.h>
}
using namespace std;

#define JIFFIES_TO_NS(TIME) ((TIME) * (1000000000 / HZ))
#define SLEEP_TIME_USEC 100 * 1000 // 100 ms

typedef long mseconds_t;
typedef int kbytes_m;

struct runstat
{
	mseconds_t usetime;
	kbytes_m usemem;
};

// Limits
string execfile;
mseconds_t timelimit;
kbytes_m memlimit;

// Path of statfile
string statfile;

// Get run stat
void getrunstat(int pid, runstat& rs)
{
	string buffer; // For read statfile
	
	// Read statfile
	ifstream ifs(statfile.c_str());
	getline(ifs, buffer);
	ifs.close();
	
	// Find the information
	unsigned spacecount = 0;
	unsigned ut_b, ut_e, rss_b, rss_e;
	for (unsigned i = buffer.find(')') + 1; i != buffer.length(); ++i)
		if (buffer[i] == ' ')
		{
			++spacecount;
			if (spacecount == 12) // Found begin of user time
				ut_b = i + 1;
			else if (spacecount == 13) // Found end of user time
				ut_e = i;
			else if (spacecount == 22) // Found begin of rss
				rss_b = i + 1;
			else if (spacecount == 23) // Found end of rss
				rss_e = i;
			
			// Needn't continue
			if (spacecount >= 23)
				break;
		}
	
	// Pick out the information
	rs.usetime = static_cast<mseconds_t>(JIFFIES_TO_NS(
		static_cast<long long>(
			atoi(buffer.substr(ut_b, ut_e - ut_b).c_str())
		)
	) / 1000000);
	rs.usemem = atoi(buffer.substr(rss_b, rss_e - rss_b).c_str()) << 2;
}

inline mseconds_t tv2ms(timeval tv)
{
	return tv.tv_sec * 1000 + tv.tv_usec / 1000;
}

int main(int argc, char** argv)
{
	// Check the argument
	if (argc < 3)
	{
		cout << "ujudger-base 0.1" << endl;
		cout << "Usage:" << endl 
			<< "\tujudger-base [Program] [Time Limit (ms)] [Memory Limit (KB)]" << endl;
		return -1;
	}
	
	// Get the Program Filename & Limits
	execfile = argv[1];
	timelimit = atol(argv[2]);
	memlimit = atoi(argv[3]);
	
	// Create subprocess
	int pid = fork();
	
	if (pid == 0) // If it's subprocess
	{
		// Execute the program
		execl(execfile.c_str(), "", NULL);
	}
	else if (pid == -1) // If an error occurred
	{
		cout << "An error occurred when create the subprocess!" << endl;
		cout << "Error Number: " << errno << endl;
	}
	else // If it's parent-process
	{
		int status; // For get the return status
		rusage usage; // For get the final information
		runstat rs;  // For get run stat
		kbytes_m maxrss = 0; // Remember the max usage of memory
		mseconds_t timeused;
		
		// Initialize the path of statfile
		ostringstream ostr;
		ostr << "/proc/" << pid << "/stat";
		statfile = ostr.str();
		
		// Start wait for the program
		while (wait4(RUSAGE_CHILDREN, &status, WNOHANG, &usage) == 0)
		{
			bool killchild = false; // Clear the flag
			
			// Get run stat
			getrunstat(pid, rs);
			// Get & check rss using now
			maxrss = MAX(maxrss, rs.usemem);
			if (maxrss > memlimit) killchild = true; // MLE
			// Check time used
			if (rs.usetime > timelimit) killchild = true; //TLE
			
			// Check if out of limit
			if (killchild)
			{
				kill(pid, SIGKILL);
				break;
			}
			
			// Continue waiting
			usleep(SLEEP_TIME_USEC);
		}
		
		// Wait for the child
		wait4(RUSAGE_CHILDREN, &status, 0, &usage);
		
		// Check for the limits
		timeused = tv2ms(usage.ru_utime);
		cout << "Time used: " << timeused << " ms" << endl;
		if (timeused > timelimit) //TLE
			cout << "Time limit exceed!" << endl;
		cout << endl;
		cout << "Max memory: " << maxrss << " KB" << endl;
		if (maxrss > memlimit) // MLE
			cout << "Memory limit exceed!" << endl;
		cout << endl;
	}
	
	return 0;
}