#include <stdio.h>
#include <string.h>
#include <getopt.h>
#include <inttypes.h>
#include <sys/time.h>
#include <time.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/stat.h>
#include <fcntl.h> 

// parameters

char * filename = NULL;
enum testermode {READ, WRITE} mode = WRITE;
float speedlimit = 100; // MiB / s
int burstcount = 10;

// internals

int blocksize = 4096;
int blockcount = 0;

int fh = 0;

char filebuffer[65536];

static struct option long_options[] = 
{
	{"mode",     required_argument, 0,  0 },
	{"m",     required_argument, 0,  0 },
	{"file",    required_argument, 0,  0 },
	{"f",    required_argument, 0,  0 },
	{"limit",    required_argument, 0,  0 },
	{"l",    required_argument, 0,  0 },
	{"burst",    required_argument, 0,  0 },
	{"b",    required_argument, 0,  0 },
	{0,         0,                 0,  0 }
};

void parse_commandline(int argc, char **argv)
{
	int r;
	int option_index = 0;
    
    while (0 <= (r = getopt_long_only(argc, argv, "", long_options, &option_index)))
    {
        const char * optname = long_options[option_index].name;
		if (r == 0)
		{			
			if (strcmp("file", optname) == 0 || strcmp("f", optname) == 0)
			{
				filename = strdup(optarg);
			}
			else if (strcmp("mode", optname) == 0 || strcmp("m", optname) == 0)
			{
				if (strcmp("reader", optarg) == 0 || strcmp("read", optarg) == 0)
				{
					mode = READ;
				}
				else if (strcmp("writer", optarg) == 0 || strcmp("write", optarg) == 0)
				{
					mode = WRITE;
				}
				else
				{
					printf("Invalid mode: \"%s\"\n", optarg);
				}
			}
			else if (strcmp("burst", optname) == 0 || strcmp("b", optname) == 0)
			{
				burstcount = atoi(optarg);
				if (burstcount < 1)  burstcount = 1;
			}
			else if (strcmp("limit", optname) == 0 || strcmp("l", optname) == 0)
			{
				speedlimit = atoi(optarg);
				if (speedlimit < 1)  speedlimit = 1;
			}
			else
			{
				printf("Unhandled option: \"%s\"\n", optname);
			}
		}
		else
		{
			exit(1);
			break;
		}
	}
}

char * getmodename(enum testermode amode)
{
    if (READ == amode)  return "reader";
    if (WRITE == amode)  return "writer";
    return "unknown";
}

int64_t getusecs()
{
	struct timeval tv;
	struct timezone tz;
	gettimeofday(&tv, &tz);
	return (tv.tv_sec * 1000000 + tv.tv_usec);
}

void print_timestamp()
{
	struct tm *tmp;
    time_t t;
    char outstr[200];
    
	t = time(NULL);
    tmp = localtime(&t);
    
	if (strftime(outstr, sizeof(outstr), "%H:%M:%S", tmp) == 0) 
	{
        return;
    }	
    printf("%s", outstr);
}

void read_block(int blockindex)
{
	lseek(fh, blockindex * blocksize, SEEK_SET);
	
	int64_t r = read(fh, &filebuffer, blocksize);
	
	if (r <= 0)
	{
		printf("Error reading block %i\n", blockindex);
		exit(1);
	}
}

void write_block(int blockindex)
{
	lseek(fh, blockindex * blocksize, SEEK_SET);
	
	int64_t r = write(fh, &filebuffer, blocksize);
	
	if (r <= 0)
	{
		printf("Error writing block %i\n", blockindex);
		exit(1);
	}
}

void run_worker()  // generates disk writes
{
	printf("Running %s...\n", getmodename(mode));
	
	int bcnt = 0;
	int64_t bytecnt = 0;
	int n;
	
	int64_t t1 = getusecs();
	int64_t t2 = 0;
	float interval = 0;
	float speed = 0;
	
	int64_t worktime = 0;
	
	int64_t maxblocktime = 0;
	
	while (1)
	{		
		int64_t bwt1 = getusecs();
		
		// some burst read
		for (n = 1; n <= burstcount; n++)
		{
			int bi = rand() % blockcount;			
			int64_t sbt1 = getusecs();
			if (WRITE == mode)
			{
				write_block(bi);			
			}
			else
			{
				read_block(bi);
			}
			int64_t sblocktime = getusecs() - sbt1;
			if (maxblocktime < sblocktime)  maxblocktime = sblocktime;
			bcnt++;
			bytecnt += blocksize;
		}
		
		worktime += (getusecs() - bwt1);
				
		// some wait						
		while (1)
		{		
			t2 = getusecs();
			
			interval = (float) (t2 - t1) / 1000000.;

			speed = (float) bytecnt / (interval * 1024. * 1024.);
			
			if ((speed > speedlimit) && (interval < 1))
			{
				usleep(100);
			}  
			else
			{
				break;
			}
		} 
		
		if (interval >= 1)
		{
			float avgblocktime = (float) worktime / (float) bcnt;
					
			print_timestamp();
			printf(": speed: %8.3f MiB/s, time usage: %5.2f %%, avg. block time: %6.1f us, max. block time: %ld us\n", 
			   speed,
			   100 * ((float) worktime / 1000000.) / interval, 
			   avgblocktime,
			   maxblocktime 
			);
			
			if (interval > 1.1)
			{
				printf("  Work was blocked for a long time! Test interval: %8.3f s\n", interval);
			}
			
			t1 = getusecs();
			bcnt = 0;
			bytecnt = 0;
			worktime = 0;
			maxblocktime = 0;
		}
	}
}

int main(int argc, char **argv)
{
	struct stat sst;
	
	srand(time(NULL));
	
	printf("Linux kernel pdflush tester\n\n");
	printf("Usage: pdflushtest --mode read|write --file filename\n");
	printf("Options:\n  --limit x : limit the maximal speed to x MiB/s\n");
	printf("  --burst x : do the work with x block bursts\n\n");
	
    parse_commandline(argc, argv);
           
    printf("Current parameters: \n");    
    printf("  mode: %s\n", getmodename(mode));
    printf("  file: %s\n", filename);
    printf("  burst length: %i\n", burstcount);
    printf("  speed limit: %f MiB/s\n", speedlimit);
    printf("\n");
    
    // getting scratchfile parameters
    
    printf("Working file \"%s\":\n", filename);
    
    if (stat(filename, &sst) == 0)
    {
		blocksize = sst.st_blksize;
		blockcount = sst.st_size / blocksize;
		
		printf("  size: %ld\n", sst.st_size);
		printf("  blocksize: %i\n", blocksize);
		printf("  number of blocks: %i\n", blockcount);
	}
	else
	{
		printf("Error getting info for file \"%s\"\n", filename);
		exit(1);
	}
	
    printf("\n");
    
	fh = open(filename, O_RDWR);	
	if (fh < 0)
	{
		printf("Error opening file \"%s\"\n", filename);
		exit(1);
	}
	
	run_worker();
        
    printf("finished.\n");
       
    return 0;
}