#define _GNU_SOURCE	1
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <assert.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <libaio.h>

#define IOCB_CMD_READAHEAD	12

int io_priority = 0;

struct latency_info {
	unsigned long long	min;
	unsigned long long	max;
	unsigned long long	avg;
	unsigned		nr;
};

void init_latency_info(struct latency_info *info)
{
	info->min = ~0ULL;
	info->max = 0ULL;
	info->avg = 0ULL;
	info->nr = 0;
}

void update_latency(struct latency_info *info, unsigned long long latency)
{
	if (latency < info->min)
		info->min = latency;
	if (latency > info->max)
		info->max = latency;
	info->avg += latency;
	info->nr ++;
}

void show_latency(struct latency_info *info, const char *label)
{
	printf( "%s min: %Lu\n"
		"%s max: %Lu\n"
		"%s avg: %Lu\n"
		"%s nr: %u\n",
		label, info->min, label, info->max, label,
		(info->nr) ? info->avg / info->nr : 0,
		label, info->nr);
}

#define rdtsc(low,high) __asm__ __volatile__("lfence ; mfence ; rdtsc" : "=a" (low), "=d" (high))
static inline unsigned long long rdtscll(void)
{
	unsigned long long result;
	unsigned low, high;
	rdtsc(low, high);
	result = high;
	result <<= 32;
	result |= low;
	return result;
}

struct latency_info io_submit_latency;
struct latency_info write_latency;
struct latency_info read_latency;

int o_direct = 0;
int o_sync = 0;
int no_fsync = 0;

//#define NR_FILES	2
//#define BUFSIZE		(1024*1024*10)
#define NR_FILES	1
#define BUFSIZE		(1024*10)
char buf[BUFSIZE] __attribute__((aligned(4096)));

struct state {
	int		fd;
	enum {
		STATE_READ = 0,
		STATE_WAIT_READ,
		STATE_WRITE = 1,
		STATE_WAIT_WRITE,
		STATE_FSYNC,
		STATE_WAIT_FSYNC,
		STATE_DONE
	}		what;

	unsigned long long	start;
	unsigned long		written;
	unsigned long		read;

	struct iocb	iocb;
} state[NR_FILES];

unsigned nr_iocbs, nr_in_flight;
struct iocb *iocbs[NR_FILES];
io_context_t ctx = 0;

unsigned long max_write_size = sizeof(buf);

void run_state_machine(struct state *st)
{
	unsigned long left;

	switch (st->what) {
	case STATE_READ:
		st->start = rdtscll();
		left = sizeof(buf);
		left -= st->read;
		if (left > max_write_size)
			left = max_write_size;
		io_prep_pread(&st->iocb, st->fd, buf + st->read, left, 0LL);
		st->iocb.aio_lio_opcode = IOCB_CMD_READAHEAD;
		st->iocb.aio_reqprio = io_priority;
		st->what = STATE_WAIT_READ;
		st->iocb.data = st;
		iocbs[nr_iocbs++] = &st->iocb;
		break;

	case STATE_WRITE:
		st->start = rdtscll();
		left = sizeof(buf);
		left -= st->written;
		if (left > max_write_size)
			left = max_write_size;
		io_prep_pwrite(&st->iocb, st->fd, buf + st->written, left, 0LL);
		st->what = STATE_WAIT_WRITE;
		st->iocb.data = st;
		iocbs[nr_iocbs++] = &st->iocb;
		break;

	default:
		fprintf(stderr, "run_state_machine: bad state %d\n", st->what);
		exit(2);
	}
}

void run_state_machine_event(struct io_event *event)
{
	struct state *st = event->data;
	unsigned long long end;
	assert(event->obj == &st->iocb);

	switch (st->what) {
	case STATE_WAIT_READ:
		printf("read: res=%ld\n", (long)event->res);
		fflush(stdout);
		if (event->res <= 0) {
			fprintf(stderr, "state[%ld]: %ld: expected %lu\n",
				st - state, event->res, sizeof(buf));
			if (event->res < 0)
				fprintf(stderr, "event->res: %s\n",
					strerror(-event->res));
			exit(2);
		}

		end = rdtscll();
		update_latency(&read_latency, end - st->start);

		st->read += event->res;
		if (st->read < sizeof(buf)) {
			st->what = STATE_READ;
			run_state_machine(st);
			return;
		}

		st->what = STATE_DONE;
		break;

	case STATE_WAIT_WRITE:
		if (event->res <= 0) {
			fprintf(stderr, "state[%ld]: %ld: expected %lu\n",
				st - state, event->res, sizeof(buf));
			exit(2);
		}

		end = rdtscll();
		update_latency(&write_latency, end - st->start);

		st->written += event->res;
		if (st->written < sizeof(buf)) {
			st->what = STATE_WRITE;
			run_state_machine(st);
			return;
		}

		if (no_fsync) {
			st->what = STATE_DONE;
			return;
		}
		st->what = STATE_WAIT_FSYNC;
		io_prep_fsync(&st->iocb, st->fd);
		st->iocb.data = st;
		iocbs[nr_iocbs++] = &st->iocb;
		break;

	case STATE_WAIT_FSYNC:
		st->what = STATE_DONE;
		printf("fsync(%d) = %ld\n", st->fd, event->res);
		break;

	default:
		fprintf(stderr, "run_state_machine: bad state %d\n", st->what);
		exit(2);
	}
}

void do_io_submit(void)
{
	int err;
	if (nr_iocbs) {
		unsigned long long start, end;
		start = rdtscll();
		err = io_submit(ctx, nr_iocbs, iocbs);
		end = rdtscll();
		if (err != nr_iocbs) {
			if (err < 0) {
				fprintf(stderr,
					"iocbs[0]->aio_lio_opcode = %d\n",
					iocbs[0]->aio_lio_opcode);
				fprintf(stderr, "io_submit: %d: %s\n",
					err, strerror(-err));
				exit(2);
			}
			fprintf(stderr, "io_submit(%u): %d\n",
				nr_iocbs, err);
			exit(2);
		}
		nr_in_flight += err;

		update_latency(&io_submit_latency, end - start);

		nr_iocbs = 0;
	}
}

int main (int argc, char *argv[])
{
	char *file = NULL;
	int events = 65536;
	unsigned i;
	int err;

	memset(buf, 0, sizeof buf);

	init_latency_info(&io_submit_latency);
	init_latency_info(&write_latency);

	for (i=1; i<argc; i++) {
		if (!strncmp("--priority=", argv[i], 11)) {
			io_priority = atoi(argv[i] + 11);
			printf("io_priority = %d\n", io_priority);
		} else if (!strcmp("--O_DIRECT", argv[i]))
			o_direct = 1;
		else if (!strcmp("--O_SYNC", argv[i]))
			o_sync = 1;
		else if (!strcmp("--no-fsync", argv[i]))
			no_fsync = 1;
		else if (!strcmp("--max-write-size=1", argv[i]))
			max_write_size = 1;
		else if (!strncmp("--file=", argv[i], 7))
			file = argv[i] + 7;
		else {
			fprintf(stderr, "unknown arg '%s'\n", argv[i]);
			return 2;
		}
	}

	err = io_setup(events, &ctx);
	if (err) {
		fprintf(stderr, "io_setup(): %d => %s\n", err, strerror(-err));
		return 2;
	}

	for (i=0; i<NR_FILES; i++) {
		int flags = O_CREAT | O_RDONLY;
		char name[32];

		if (o_direct)
			flags |= O_DIRECT;
		if (o_sync)
			flags |= O_SYNC;
		sprintf(name, "testfile.%u", i);
		state[i].fd = open(file ? file : name, flags, 0600);
		if (state[i].fd < 0) {
			fprintf(stderr, "error opening '%s': %s\n", name,
				strerror(errno));
			return 2;
		}
	}

	nr_iocbs = 0;

	for (i=0; i<NR_FILES; i++) {
		run_state_machine(&state[i]);
	}

	do_io_submit();

	while (nr_in_flight) {
		struct io_event events[NR_FILES];

		err = io_getevents(ctx, 1, NR_FILES, events, NULL);
		if (err <= 0) {
			fprintf(stderr, "io_getevents(): %d: %s\n", err,
				strerror(-err));
			exit(2);
		}

		nr_in_flight -= err;
		for (i=0; i<err; i++)
			run_state_machine_event(&events[i]);

		do_io_submit();
	}

	printf("done!\n");
	show_latency(&io_submit_latency, "io_submit");
	show_latency(&write_latency, "write");
	show_latency(&read_latency, "read");
	if (!no_fsync)
		show_latency(&io_submit_latency, "fsync");

	return 0;
}