#include <unistd.h>
#include <stdlib.h>
#include <sys/time.h>

#include <iostream>
#include <fstream>
#include <string>
#include <cstdint>
using namespace std;

#include <TFile.h>
#include <TTree.h>


string rootfilename = "hd_iotest.root";

uint64_t USER_BUFF_SIZE = 0;  // kB
double   USER_FILE_SIZE = 20.0; // GB
bool READ_MODE  = false;
bool WRITE_MODE = false;

//char BOLD[] = {0x1B, '[', '1', 'm', 0};
//char NORM[] = {0x1B, '[', '0', 'm', 0};

#ifndef ansi_escape
#define ansi_escape                     ((char)0x1b)
#define ansi_bold                       ansi_escape<<"[1m"
#define ansi_normal                     ansi_escape<<"[0m"
#endif // ansi_escape



//-----------------
// Usage
//-----------------
void Usage(string mess="")
{

	cout << endl;
	cout << "Usage:" << endl;
	cout << "      hd_iotest [options]" << endl;
	cout << endl;
	cout << " options:" << endl;
	cout << "  -h              Print this usage message" << endl;
	cout << "  -B block_size   Set block size in kB" << endl;
	cout << "  -R              Read files, don't write them" << endl;
	cout << "  -W              Write files, don't read them" << endl;
	cout << endl;
	cout << "This will write or read a series of 20GB files to the local" << endl;
	cout << "directory as fast as possible to test the write/read speed" << endl;
	cout << "of the device. Writes and reads are done using a fixed block" << endl;
	cout << "size which can be set using the -B option." << endl;
	cout << "If the -R option is given, files are read rather than" << endl;
	cout << "written. It looks for files with names like \"junk000.tmp\"," << endl;
	cout << "\"junk001.tmp\", ... These can be created by running this" << endl;
	cout << "with the -W option. When reading, it will continuously" << endl;
	cout << "loop over whatever files it finds." << endl;
	cout << "If the -W option is given files will be written. Up to 100" << endl;
	cout << "files will be created before program exit. Use ctl-C to exit" << endl;
	cout << "early." << endl;

	cout << endl << mess << endl << endl;
	
	exit(0);
}

//-----------------
// main
//-----------------
int main(int narg, char *argv[])
{

	for(int i=1; i<narg; i++){
		string arg = argv[i];
	
		if(arg.find("-h")==0) Usage();
		if(arg.find("-B")==0) USER_BUFF_SIZE = atoi(arg.substr(2).c_str());
		if(arg.find("-R")==0) READ_MODE = true;
		if(arg.find("-W")==0) WRITE_MODE = true;
	}
	
	if( !(READ_MODE || WRITE_MODE) ) Usage();
	if(  (READ_MODE && WRITE_MODE) ) Usage("Only one of -R or -W may be specified!");


	TFile *f = new TFile(rootfilename.c_str(), "RECREATE");
	if(!f->IsOpen()){
		cout << "Unable to open ROOT file \""<<rootfilename<<"\"!" << endl;
		return -1;
	}
	
	double rate;
	double buff_size;
	TTree *t = new TTree("rate","I/O write rates");
	t->Branch("rate", &rate, "rate/D");
	t->Branch("buff_size", &buff_size, "buff_size/D");
	
	// Iterate over files
	int ifile = 0;
	buff_size = USER_BUFF_SIZE==0 ? 128.0:USER_BUFF_SIZE;
	for(int j=0; j<100; j++, ifile++){
	
		cout << "Preparing for " << (int)buff_size << "kB buffer " << (READ_MODE ? "read":"write") << " test ..." << endl;
		
		uint64_t buff_size_bytes = (uint32_t)buff_size*1024;
		char *buff = new char[buff_size_bytes];
		
		ofstream *ofs = NULL;
		ifstream *ifs = NULL;
		char fname[256];
		sprintf(fname, "junk%03d.tmp", ifile);
		if(READ_MODE){
			ifs = new ifstream(fname);
			if(!ifs->is_open()) {
				cerr << "ERROR: Unable to open " << fname << " for reading ..!" << endl;
				
				// If this is first file then exit. Otherwise
				// reset file counter and try again
				if(ifile==0) exit(-1);
				ifile=-1; // continue will increment it to 0
				continue;
			}
			cout << "Opened " << fname << " for reading ... "  << endl;
		
		}else{
		
			// Open file for writing
			ofs = new ofstream(fname);
			if(!ofs->is_open()) { cerr << "ERROR: Unable to open " << fname << " for writing ..!" << endl; exit(-1);}
			cout << "Opened " << fname << " for writing ... "  << endl;
		}
		
		// Time of start
		struct timeval tv_start;
		struct timeval tv_now;
		struct timeval tv_diff;
		gettimeofday(&tv_start, NULL);
		
		// Write 1 file worth of data
		uint64_t Nbuffs = 0;
		uint64_t N = (double)(USER_FILE_SIZE*pow(1024.0,3.0))/(double)buff_size_bytes;
		uint64_t nMB = (1.0*pow(1024.0,2.0))/(double)buff_size_bytes;
		if(nMB<50.0) nMB=50.0;
		bool done = false;
		for(uint64_t i=1; i<=N; i++, Nbuffs++){
		
			if(READ_MODE){
				ifs->read(buff, buff_size_bytes);
				if(!ifs->good()) done = true;
			}else{
				ofs->write(buff, buff_size_bytes);
				ofs->flush();
			}
			
			if( (i%nMB) == 0){
				t->Fill();
				
				gettimeofday(&tv_now, NULL);
				timersub(&tv_now, &tv_start, &tv_diff);
				uint64_t usec = tv_diff.tv_sec*(uint64_t)1000000 + tv_diff.tv_usec;

				double percent = 100.0*(double)i/(double)N;
				cout << "   " << percent << "% complete   ("<<usec<<" us = " << (double)usec*1.0E-6 << " sec)            \r";
				cout.flush();
			}
			
			if(done) break;
		}
		
		cout << endl;
		cout << "Closing file ... " ; cout.flush();
		if(ofs) { ofs->close(); delete ofs; }
		if(ifs) { ifs->close(); delete ifs; }

		// Include time to close file since it seems some flushing is
		// still done here regardless of the calls to flush above.
		gettimeofday(&tv_now, NULL);
		timersub(&tv_now, &tv_start, &tv_diff);
		uint64_t usec = tv_diff.tv_sec*(uint64_t)1000000 + tv_diff.tv_usec;
		double sec = (double)usec/1.0E6;
		cout << ansi_bold << (double)Nbuffs*(double)buff_size_bytes/pow(1024.0,2.0)/sec << " MB/s" << ansi_normal << endl;

		if(buff) delete buff;
	}
	
	
	f->Write();
	f->Close();
	delete f;

	return 0;
}