#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>
#include <unistd.h>
#include <numaif.h>

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->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:
		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;
	}
}

struct aio_ring {
	unsigned	id;     /* kernel internal index number */
	unsigned	nr;     /* number of io_events */
	unsigned	head;
	unsigned	tail;

	unsigned	magic;
	unsigned	compat_features;
	unsigned	incompat_features;
	unsigned	header_length;  /* size of aio_ring */

	struct io_event	io_events[0];
};

unsigned long get_ring_size(io_context_t ctx)
{
	struct aio_ring *ring = (void*)ctx;
	unsigned long size = sizeof(*ring);

	size += ring->nr * sizeof(struct io_event);
	size += (4096 - (size % 4096)) & 4095;
	return size;
}

int main (int argc, char *argv[])
{
	char *file = NULL;
	int events = 64; //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 (!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[1], 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;
	}

	unsigned long ring_size = get_ring_size(ctx);
	printf("ctx = 0x%lx, ring_size = %lu\n", (unsigned long)ctx, ring_size);
	long ret;
#if 0
	unsigned long oldnode = 0, newnode = 2;
	ret = migrate_pages(0, 2, &oldnode, &newnode);
	if (ret < 0)
		perror("migrate_pages(0 -> 2)");
	printf("migrate_pages(0 -> 2) = %ld\n", ret);
	oldnode = 2;
	newnode = 1;
	ret = migrate_pages(0, 2, &oldnode, &newnode);
	if (ret < 0)
		perror("migrate_pages(2 -> 1)");
	printf("migrate_pages(2 -> 1) = %ld\n", ret);
#else
	void *pages[1] = { (void *)ctx };
	int nodes[1] = { 1 };
	int status[1] = { 0 };
	printf("pages[0] = %p\n", pages[0]);
	ret = move_pages(0, 1, pages, nodes, status, MPOL_MF_MOVE_ALL);
	if (ret < 0)
		perror("move_pages(to 1)");
	printf("to 1: status[0] = %d\n", status[0]);
	nodes[0] = 0;
	ret = move_pages(0, 1, pages, nodes, status, MPOL_MF_MOVE_ALL);
	if (ret < 0)
		perror("move_pages(to 0)");
	printf("to 2: status[0] = %d\n", status[0]);
	nodes[0] = 1;
	ret = move_pages(0, 1, pages, nodes, status, MPOL_MF_MOVE_ALL);
	if (ret < 0)
		perror("move_pages(to 1)");
	printf("to 2: status[0] = %d\n", status[0]);
#endif
#if 0
	sleep(60);
	exit(0);
#endif

	for (i=0; i<NR_FILES; i++) {
		int flags = O_CREAT | O_RDWR;
		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;
}