#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <errno.h>
#include <time.h>
#include <signal.h>
#include <pthread.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/uio.h>
#include <sys/time.h>
#include <sys/resource.h>

#include <cstdio>
#include <iostream>
#include <fstream>
#include <string>
#include <sstream>
using namespace std;

uint32_t maxkB = 10240;
string OUTPUTFILE = "";
int PROGRAM_ARGV_IDX = -1;
bool ECHO_TO_STDOUT = false;
bool DELETE_EXISTING_LOG_FILE = false;
bool PRINT_COMMAND_STR = false;
int  MAX_SPAWNS = 1;
int  SECS_BETWEEN_SPAWNS = 15;
bool CHILD_SIGNALED = false;
bool ADD_TIMESTAMP_PREFIX = true;
string COMMAND_STR; // string of command to execute (used just for printing)

pid_t CHILD_PID = 0;
int WAIT_RETURN_STATUS = 1;
rlim_t COREDUMPSIZE = 0; 

void RunProcess(int narg, char *argv[], int spawn_attempt);
string GetChildResourceUsageString(void);

void TruncateOutputFile(ofstream &ofs, bool make_headroom=true);
void ParseCommandLineArguments(int narg, char *argv[]);
void Usage(bool quit);

//.....................
// sigHandler
void sigHandler(int sigNum)
{
	if(CHILD_PID != 0){
		CHILD_SIGNALED = true;
		kill(CHILD_PID, sigNum);
	}else{
		// Should never get here since child will have replaced this process
		if(sigNum != SIGCHLD) exit(0);
	}
}

//.....................
// waitThread
void* waitThread(void *arg)
{
	// thread that will sit in
	// "wait" so we can capture
	// resource usage of child
	while(true){
		wait(&WAIT_RETURN_STATUS);

		// Keep waiting until child is terminated
		if(WIFEXITED(WAIT_RETURN_STATUS)) break;
		if(WIFSIGNALED(WAIT_RETURN_STATUS)) break;
		//cout << "STOP OR CONTINUE RECEIVED!!" << endl;
	}
		
	pthread_exit(NULL);
}

//-----------------------------
// main
//-----------------------------
int main(int narg, char *argv[])
{
	ParseCommandLineArguments(narg, argv);
	cout <<"-------------hdlog starting--------------" << endl;
	cout <<"          command: " << COMMAND_STR << endl;
	cout <<"  max. file size : " << maxkB << "kB" << endl;
	cout <<"    log file name: " << OUTPUTFILE << endl;
	cout <<"         log mode: " << (DELETE_EXISTING_LOG_FILE ? "overwrite existing":"append") << endl;
	cout <<"    max. respawns: " << MAX_SPAWNS << endl;

	// Loop as many times as we should respawn the program.
	for(int i=0; i<MAX_SPAWNS; i++){
	
		// Set environment variable to let process know how many
		// times it has been respawned. This is used by hdmon to
		// pass the info on to hdmongui.
		stringstream ss;
		ss << i;
		setenv("HDLOG_NUM_RESPAWNS", ss.str().c_str(), 1);

		// Record start time
		auto t_started = time(NULL);

		// Run the process and wait for it to end
		RunProcess(narg, argv, i);
		
		// If it exited cleanly, then we are done
		if( WIFEXITED(WAIT_RETURN_STATUS) ){
			// process exited normally by calling exit or _exit
			// (not by seg. fault etc.)
			if( WEXITSTATUS(WAIT_RETURN_STATUS) == 0 ) break;
		}
		
		// If we passed a signal to the child then assume someone really
		// wants it to quit and it didn't just die on its own. Don't
		// respawn.
		if( CHILD_SIGNALED ) break;
		
		// Pause between respawns 
		auto t_ended = time(NULL);
		int t_diff = (int)(t_ended - t_started);
		if( t_diff<SECS_BETWEEN_SPAWNS ) {
			cout << "Pausing " << SECS_BETWEEN_SPAWNS << " seconds before respawning ..." << endl;
			sleep(SECS_BETWEEN_SPAWNS);
		}
	}

	return 0;
}

//-----------------------------
// RunProcess
//
// This will fork a child process to run the actual process and
// capture the output to the log file. A new thread is launched
// to wait for the child to exit so its resource usage can be
// written to the log. This will not return until the child
// process has ended. It can therefore be called multiple times
// if the child should be respawned.
//-----------------------------
void RunProcess(int narg, char *argv[], int spawn_attempt)
{

	// We use the standard Unix trick of creating a pipe,
	// forking a child process, replacing stdout and stderr
	// in the child process with the pipe ends, and then
	// exec-ing the child process into a new process (i.e.
	// the program we are running)

	// Create pipes
	int fd_stdout[2];
	int fd_stderr[2];
	pipe(fd_stdout);
	pipe(fd_stderr);
	
	// Set the core dump size if one was specified. The
	// child should inherit this and then we can reset
	// the parent to whatever the default is.
	struct rlimit save_limit;
	getrlimit(RLIMIT_CORE, &save_limit);
	if( COREDUMPSIZE > 0 ){
		// Set soft limit to smaller of COREDUMPSIZE or
		// current max limit. Keep current max limit.
		struct rlimit new_limit = { COREDUMPSIZE, save_limit.rlim_max };
		if( COREDUMPSIZE < new_limit.rlim_max ) new_limit.rlim_cur =  new_limit.rlim_max;
		cout <<"Setting coredumpsize to " << COREDUMPSIZE << " (was " << save_limit.rlim_cur << ")" << endl;
		setrlimit(RLIMIT_CORE, &new_limit);
	}
	
	// Fork child process
	pid_t pid = fork();
	if(pid == -1){
		cerr << "Error forking !!!" << endl;
		exit(-2);
	}
	
	if(pid == 0){
		// ------- child
		close(fd_stdout[0]);
		close(fd_stderr[0]);

		// Set both stdout and stderr to feed the fd_stdout pipe 
		dup2(fd_stdout[1], STDOUT_FILENO);
		dup2(fd_stdout[1], STDERR_FILENO);
		
		execvp(argv[PROGRAM_ARGV_IDX], &argv[PROGRAM_ARGV_IDX]);
		
	}else{
		// ------- parent
		close(fd_stdout[1]);
		close(fd_stderr[1]);
		
		// Dispatch signals to child process
		CHILD_PID = pid;
		signal(SIGINT, sigHandler);
		signal(SIGTERM, sigHandler);
		
		// Reset coredump file size for parent (if set to something else above)
		if( COREDUMPSIZE > 0 ){
			setrlimit(RLIMIT_CORE, &save_limit);
		}
		
		// Launch a thread to wait on the child process and
		// capture resource usage for it
		pthread_t thread;
		pthread_create(&thread, NULL, waitThread, NULL);
		
		ios_base::openmode mode = ofstream::app;
		if(DELETE_EXISTING_LOG_FILE) mode = ofstream::out;
		ofstream ofs(OUTPUTFILE.c_str(), mode);
		time_t t = time(NULL);
		string tstr(ctime(&t));
		tstr.erase(tstr.length()-1);
		ofs << endl;
		ofs << "------- hdlog starting: " << tstr << " ----------" << endl;
		if(PRINT_COMMAND_STR) ofs << "command: " << COMMAND_STR  << endl << endl;

		// Record if we had to respawn child
		if( spawn_attempt > 0 ){
			cout << "--------- respawning " << spawn_attempt+1 << "/" << MAX_SPAWNS << "--------------" << endl;
			ofs  << "respawning: " << spawn_attempt+1 << "/" << MAX_SPAWNS << endl;
		}

		uint32_t buff_size = 16386;
		char *buff = new char[buff_size];
		long filesize = ofs.tellp();
		
		// Loop until done
		do{
			ssize_t Nbytes = read(fd_stdout[0], buff, buff_size);
			if(Nbytes<=0) break;
			
			string prefix = "";
			if(ADD_TIMESTAMP_PREFIX){
				t = time(NULL);
				prefix = ctime(&t);
				prefix.erase(prefix.length()-1);
				prefix = string("[") + prefix + "] ";
			}
			
			if(ECHO_TO_STDOUT){
				if( !prefix.empty() ) cout << prefix;
				cout.write(buff, Nbytes);
				cout.flush();
			}
			
			if( !prefix.empty() ) ofs << prefix;
			ofs.write(buff, Nbytes);
			ofs.flush();
			filesize += Nbytes;
			
			if(filesize/1024 > maxkB){
				TruncateOutputFile(ofs);
				filesize = ofs.tellp(); // get new file size
			}
			
		}while(true);

		delete[] buff;
		
		t = time(NULL);
		tstr = ctime(&t);
		tstr.erase(tstr.length()-1);
		ofs << endl;
		ofs << "------- hdlog ending: " << tstr << " ----------" << endl;

		void *ptr;
		pthread_join(thread, &ptr);

		// Final check on max file size
		filesize = ofs.tellp();
		if(filesize/1024 > maxkB){
			TruncateOutputFile(ofs, false);
		}
		filesize = ofs.tellp();
		cout <<" final log file size : " << filesize/1024 << "kB" << endl;

		// Get child resource usage
		string usage_str = GetChildResourceUsageString();
		ofs << usage_str;
		cout << usage_str;
		
		// Print info on respawning
		ofs << "                 CHILD_SIGNALED = " << CHILD_SIGNALED <<endl;
		ofs << "             WAIT_RETURN_STATUS = " << WAIT_RETURN_STATUS <<endl;
		ofs << "  WIFEXITED(WAIT_RETURN_STATUS) = " << WIFEXITED(WAIT_RETURN_STATUS) <<endl;
		ofs << "WEXITSTATUS(WAIT_RETURN_STATUS) = " << WEXITSTATUS(WAIT_RETURN_STATUS) <<endl;
		
		// Check if a core dump occured and report it.
#ifdef WCOREDUMP
		if( WIFSIGNALED(WAIT_RETURN_STATUS) ){
			auto core_dumped = WCOREDUMP(WAIT_RETURN_STATUS);
			ofs << "  WCOREDUMP(WAIT_RETURN_STATUS) = " << core_dumped << endl;
			if( core_dumped ){
				ofs << "  ( core file likely named: ./core."<<pid<<" )" << endl;
			}
		}
#endif
		 
		if( (spawn_attempt+1)<MAX_SPAWNS ){
			if( WEXITSTATUS(WAIT_RETURN_STATUS) == 0 ){
				ofs << "Return status 0 so will not respawn." << endl;
			}
			if(CHILD_SIGNALED){
				ofs << "Child signaled so will not respawn." << endl;
			}
		}else{
			ofs << "Process spawned max. times (" << MAX_SPAWNS << ")" << endl;
		}

		ofs.close();
		close(fd_stdout[0]);
		close(fd_stderr[0]);
	}

	cout <<"-----------------------------------------" << endl;
	
}

//-----------------------------
// GetChildResourceUsageString
//-----------------------------
string GetChildResourceUsageString(void)
{
	stringstream ss;

	// Get resource usage of child process
	struct rusage usage;
	getrusage(RUSAGE_CHILDREN, &usage);
	double t_user = (double)usage.ru_utime.tv_sec + 1.0E-6*(double)usage.ru_utime.tv_usec;
	double t_sys  = (double)usage.ru_stime.tv_sec + 1.0E-6*(double)usage.ru_stime.tv_usec;
	double mem_usage = (double)(usage.ru_maxrss)/1024.0;
	ss <<"            user time: " << t_user << " sec" << endl;
	ss <<"          system time: " << t_sys << " sec" << endl;
	ss <<"       max. mem usage: " << mem_usage << " kB" << endl;
	
	return ss.str();	
}

//-----------------------------
// TruncateOutputFile
//-----------------------------
void TruncateOutputFile(ofstream &ofs, bool make_headroom)
{
	// This will close the current output file stream,
	// chop off the top part of the file, and then reopen
	// the output stream so we can keep appending to it.
	// This operation is a little tricky since the only
	// practical way of doing this is to copy the part of
	// the log we wish to keep into a temporary file and
	// then delete the original before renaming the temporary
	// file. This operation can be costly, especially for
	// very large files so we don't necessarily want to do it
	// every time a new line is written once the file reaches
	// the limit. Thus, we delete enough to give us at least
	// 10% or 2kB of space before hitting the limit, whichever
	// is larger. In addition, we don't want to chop at the
	// middle of a line so we only get rid of chunks where the
	// last character is a line feed.

	// Close current log file
	ofs.close();

	// Open temporary file
	string temp_fname = string(".tmp_") + OUTPUTFILE;
	ofstream tofs(temp_fname.c_str()); // overwrite any existing temp file

	// Open current log file for reading
	ifstream ifs(OUTPUTFILE.c_str());

	// Determine how much of current log file to keep
	long headroom = (maxkB*1024)/10; // 10%
	if(headroom < 2048) headroom = 2048;
	if(!make_headroom) headroom = 0;
	long max_size = (maxkB*1024) - headroom;
	char buff[1024];
	ifs.seekg(-max_size, ifs.end); // start looking for newline at target size of file
	ifs.read(buff, 1024); // read in 1kB of data
	int Nbytes = ifs.gcount();
	int iextra = 0;
	for(; iextra<Nbytes; iextra++){
		if(buff[iextra] == '\n') break;
	}
	// chop off partial line. If no newline was found in the 1kB
	// buffer, then don't chop anything off
	if(iextra < Nbytes) max_size -= iextra;

	// Set point back to where we want to start reading from
	ifs.seekg(-max_size, ifs.end);
	
	// Copy current log file to temporary file
	ifs.clear();
	while(!ifs.eof()){
		ifs.read(buff, 1024);
		int Nbytes = ifs.gcount();
		if(Nbytes <= 0)break;
		tofs.write(buff,Nbytes);
	}
	
	// Close both files
	ifs.close();
	tofs.close();

	// Remove original log file and rename temporary file
	remove(OUTPUTFILE.c_str());
	rename(temp_fname.c_str(), OUTPUTFILE.c_str());

	// Re-open the output file
	ofs.open(OUTPUTFILE.c_str(), ofstream::app);
}

//-----------------------------
// ParseCommandLineArguments
//-----------------------------
void ParseCommandLineArguments(int narg, char *argv[])
{
	if(narg<2) Usage(true);
	
	// Look for hdlog options
	int iarg=1;
	for(; iarg < narg; iarg++){
		string arg = argv[iarg];
		string arg2 = (iarg+1) < narg ? argv[iarg+1]:"";
		if( (arg=="-h") || (arg=="--help") ){
			Usage(true);
		}else if(arg == "-o") {
			if(arg2 == ""){
				cout << endl << "You must provide a filename with \"-o\"!" << endl;
				Usage(true);
			}
			OUTPUTFILE = arg2;
			iarg++;
		}else if(arg == "-s") {
			if(arg2 == ""){
				cout << endl << "You must provide a max size with \"-s\"!" << endl;
				Usage(true);
			}
			maxkB = atoi(arg2.c_str());
			iarg++;
		}else if(arg == "-e"){
			ECHO_TO_STDOUT = true;
		}else if(arg == "-d"){
			DELETE_EXISTING_LOG_FILE = true;
		}else if(arg == "-c"){
			PRINT_COMMAND_STR = true;
		}else if(arg == "-p"){
			ADD_TIMESTAMP_PREFIX = false;
		}else if(arg == "-r"){
			MAX_SPAWNS = atoi(arg2.c_str());
			iarg++;
		}else if(arg == "--coredumpsize"){
			COREDUMPSIZE = atoi(arg2.c_str());
			iarg++;
		}else{
			break;
		}
	}

	// Look for program name and arguments
	if(iarg >= narg){
		cout << endl << "You must provide the name of a program to run!" << endl;
		Usage(true);
	}
	
	// Record start of actual command user wants us to run
	PROGRAM_ARGV_IDX = iarg;

	// Optionally generate output file name
	if(OUTPUTFILE == ""){
		const char *host   = getenv("HOST");
		const char *logdir = getenv("HDLOG");
		string shost = (host==NULL ? "":host);
		size_t pos = shost.find_first_of(".");
		if(pos != shost.npos) shost.erase(pos);
		string slogdir = (logdir==NULL ? "":logdir);
		if(slogdir != "") slogdir += "/";
		
		OUTPUTFILE = slogdir + argv[iarg] + "_" + shost + ".log";
	}

	// Create string of command line we're going to execute
	for(int i=PROGRAM_ARGV_IDX; i<narg; i++)COMMAND_STR += string(argv[i]) + " ";
}

//-----------------------------
// Usage
//-----------------------------
void Usage(bool quit)
{
	cout << endl;
	cout << "  Usage:" << endl;
	cout << endl;
	cout << "      hdlog [hdlog options] program arg1 arg2 ...." << endl;
	cout << endl;
	cout << " Simple program that runs a program and captures its" << endl;
	cout << "stdin and stderr, writing them to an output log file. " << endl;
	cout << "The log file name is derived from the program name and " << endl;
	cout << "the host unless it is set explicitly via the -o option." << endl;
	cout << "The size of the file is capped at 1MB unless set to" << endl;
	cout << "something else via the -s option." << endl;
	cout << endl;
	cout << "options:" <<endl;
	cout << endl;
	cout << "    -h , --help          print this usage statement" << endl;
	cout << "    -o file.log          set the output file name " << endl;
	cout << "    -s maxkB             set max file size in kB" << endl;
	cout << "    --coredumpsize size  set the max core dump size for the process (bytes)" << endl;
	cout << "    -e                   echo output to both stdout and log file" << endl;
	cout << "    -c                   print command to be executed to log file" << endl;
	cout << "    -d                   overwrite any existing log file rather than append" << endl;
	cout << "    -p                   do NOT prefix each line with a timestamp" << endl;
	cout << "    -r                   respawns  max times to respawn process if it exits abnormally" << endl;
	cout << endl;
	

	if(quit) exit(0);
}