#include <stdlib.h>
#include <stdio.h>
#include <time.h>
#include <sys/mman.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <errno.h>
#include <stdbool.h>

#define NUM_OF(a)   (sizeof(a)/sizeof((a)[0]))

static int update_file(const char *fname, const void *p_data, unsigned data_size, int mmap_prot, int mprotect_prot, bool do_msync_on_change)
{
	void *mem;
	int fd;
	struct stat stat_buf;

	// Open up the file as rd/wr (creating if necessary)
	fd = open(fname, O_RDWR | O_CREAT, 0666);
	if (fd == -1)
	{
		printf("Failed to open %s: %s\n", fname, strerror(errno));
		return -1;
	}

	// Get the existing file size
	if (fstat(fd, &stat_buf) == -1)
	{
		printf("Failed to stat %s: %s", fname, strerror(errno));
		return -1;
	}

	// If the existing file size is different from our new size, truncate the file to the desired length
	if (stat_buf.st_size != data_size)
	{
		if (ftruncate(fd, data_size) == -1)
		{
			printf("Failed to truncate %s to %u: %s\n", fname, data_size, strerror(errno));
			return -1;
		}
	}

	if ((mem = mmap(NULL, data_size, mmap_prot, MAP_SHARED, fd, 0)) == MAP_FAILED)
	{
		printf("Failed to map in %s: %s\n", fname, strerror(errno));
		return -1;
	}

	// memcmp against our mmap'd buffer
	if (memcmp(mem, p_data, data_size) == 0)
	{
		// We don't need to sync the data (it already matches)
	}
	else
	{
		// Update the file. Do an mprotect (if specified) to update permissions.
		if (mprotect_prot != 0)
		{
			if (mprotect(mem, data_size, mprotect_prot) == -1)
			{
				printf("mprotect update failed for %s: %s\n", fname, strerror(errno));
				return -1;
			}
		}

		// Only actually copy and sync if PROT_WRITE was set on  the mprotect prot (or initial prot if no mprotect) and we were supposed to do an msync
		if (do_msync_on_change &&
			( (mprotect_prot & PROT_WRITE) ||
			  ( (mprotect_prot == 0) && (mmap_prot & PROT_WRITE) )
			))
		{
			// Copy to our mmap'd buffer
			memcpy(mem, p_data, data_size);

			// Sync out to disk
			if (msync(mem, data_size, MS_ASYNC) == -1)
			{
				printf("Failed to msync %s: %s\n", fname, strerror(errno));
				return -1;
			}
		}
	}

	// Close handles
	// We leak handles in the above code if there were errors, but since this is a test program, we don't care
	if (munmap(mem, data_size) == -1)
	{
		printf("munmap failed: %s\n", strerror(errno));
		return -1;
	}

	if (close(fd) == -1)
	{
		printf("close failed: %s\n", strerror(errno));
		return -1;
	}

	return 0;
}

int main(int argc, char *argv[])
{

	const struct
	{
		const char *descr;
		int mmap_prot;
		int mprotect_prot;
		bool do_msync_on_change;
	} prot_flags[] = {
		{"mmap PROT_READ | PROT_WRITE, no additional mprotect", PROT_READ | PROT_WRITE, 0, true},
		{"mmap PROT_READ, mprotect PROT_READ | PROT_WRITE", PROT_READ, PROT_READ | PROT_WRITE, true},
		{"mmap PROT_READ, no additional mprotect (file is not actually updated even on a change)", PROT_READ, 0, false},
		{"mmap PROT_READ | PROT_WRITE, mprotect PROT_READ (file is not actually updated even on a change)", PROT_READ | PROT_WRITE, PROT_READ, false},
		{"mmap PROT_READ, mprotect PROT_READ | PROT_WRITE, but no mysnc (file is not actually updated even on a change)", PROT_READ, PROT_READ | PROT_WRITE, false},
	};
	struct timespec ts1, ts2;
	char *fname;
	unsigned i, j;
	char *default_str_data = "hello";
	char *str_data_0 = default_str_data;
	char *str_data_1 = default_str_data;

	if (argc < 2)
	{
		printf("Usage: %s <filename> [str_data_0 str_data_1]\n", argv[0]);
		return -1;
	}
	fname = argv[1];
	if (argc >= 3)
	{
		str_data_0 = argv[2];
	}
	if (argc >= 4)
	{
		str_data_1 = argv[3];
	}

	printf("Toggling between writing %s and %s to %s\n", str_data_0, str_data_1, fname);

	// Forever in a loop, set the data values for each set of flags in sets of 100
	for (;;)
	{
		for (i = 0; i < NUM_OF(prot_flags); i++)
		{
			clock_gettime(CLOCK_MONOTONIC, &ts1);
			for (j = 0; j < 100; j++)
			{
				if (j % 2 == 0)
				{
					if (update_file(fname, str_data_0, strlen(str_data_0), prot_flags[i].mmap_prot, prot_flags[i].mprotect_prot, prot_flags[i].do_msync_on_change) == -1)
					{
						printf("Error with %s... exiting\n", prot_flags[i].descr);
						return -1;
					}
				}
				else
				{
					if (update_file(fname, str_data_1, strlen(str_data_1), prot_flags[i].mmap_prot, prot_flags[i].mprotect_prot, prot_flags[i].do_msync_on_change) == -1)
					{
						printf("Error with %s... exiting\n", prot_flags[i].descr);
						return -1;
					}
				}
			}
			clock_gettime(CLOCK_MONOTONIC, &ts2);
			printf("Elapsed time %5llu msec for 100 cycles of %s\n", (timespec2nsec(&ts2) - timespec2nsec(&ts1)) / 1000000, prot_flags[i].descr);
		}
	}

	return 0;
}