#include <fcntl.h>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
#include <sys/shm.h>
#include <sys/mman.h>

#ifdef __APPLE__
#	include <sys/sysctl.h>
#endif

#include <Python.h>

#include "hashmap.h"
#include "pointers.h"

#define F_HEAD					0x00000001
#define F_BODY					0x00000002
#define F_TAIL					0x00000004
#define F_ALLOC					0x80000000

#define SHM_RETRY				10

#define PAGEMAP_KEY_SIZE		64
#define PAGEMAP_MAX_SHMID		2048
#define PAGEMAP_INCREMENT		(128 * 1024 * 1024) /* Max usable shared memory = PAGEMAP_MAX_SHMID * PAGEMAP_INCREMENT = 256G, should be enough */

/* adapt for OS X */
#ifdef __APPLE__
#	define MAP_ANONYMOUS		MAP_ANON
#endif

typedef struct _shm_t
{
	int shmid;
	long pageid;
	uint32_t pages;
} shm_t;

typedef struct _pagemap_t
{
	uint32_t lock;
	uint32_t free;
	uint32_t total;
	hashmap_t hashmap;
} pagemap_t;

typedef struct _pagetable_t
{
	int shmid;
	uint32_t flags;
	uint32_t count;
	uint32_t offset;
} pagetable_t;

/* Class struct declatation */

typedef struct _HashMapClass
{
	PyObject_HEAD

	long pageid;
	void *memory;
	uint32_t pages;
	hashmap_t *hashmap;
	PyStringObject *name;
} HashMapClass;

/* Global variables */

static long pageSize = 0;
static long pageCount = 0;

static int shmidCount = 0;
static int shmidBuffer[PAGEMAP_MAX_SHMID] = {0};

static pagemap_t *pagemap = NULL;
static pagetable_t *pagetable = NULL;

/* Method prototypes */

static void HashMap___dealloc__(HashMapClass *self);
static int HashMap___init__(HashMapClass *self, PyObject *args, PyObject *kwargs);

static int HashMap___clear__(HashMapClass *self);
static int HashMap___traverse__(HashMapClass *self, visitproc visit, void *arg);

static PyObject *HashMap_keys(HashMapClass *self, PyObject *args);
static PyObject *HashMap_find(HashMapClass *self, PyObject *args);
static PyObject *HashMap_clear(HashMapClass *self, PyObject *args);
static PyObject *HashMap_insert(HashMapClass *self, PyObject *args);
static PyObject *HashMap_remove(HashMapClass *self, PyObject *args);

static PyObject *HashMap_lock(HashMapClass *self, PyObject *args);
static PyObject *HashMap_unlock(HashMapClass *self, PyObject *args);

static PyObject *HashMap_userData(HashMapClass *self, PyObject *args);
static PyObject *HashMap_setUserData(HashMapClass *self, PyObject *args);

/* Class Declarations */

static PyMethodDef AppMethods[] =
{
	/*	ml_name			ml_meth								ml_flags		ml_doc								*/
	{	"keys",			(PyCFunction)HashMap_keys,			METH_NOARGS,	"Get keys of HashMap"				},
	{	"find",			(PyCFunction)HashMap_find,			METH_VARARGS,	"Lookup HashMap entry"				},
	{	"clear",		(PyCFunction)HashMap_clear,			METH_NOARGS,	"Clear the entire HashMap"			},
	{	"insert",		(PyCFunction)HashMap_insert,		METH_VARARGS,	"Insert items into HashMap"			},
	{	"remove",		(PyCFunction)HashMap_remove,		METH_VARARGS,	"Remove items from HashMap"			},

	{	"lock",			(PyCFunction)HashMap_lock,			METH_NOARGS,	"Lock the HashMap manually"			},
	{	"unlock",		(PyCFunction)HashMap_unlock,		METH_VARARGS,	"Unlock the HashMap manually"		},

	{	"userData",		(PyCFunction)HashMap_userData,		METH_NOARGS,	"Get UserData stored in HashMap"	},
	{	"setUserData",	(PyCFunction)HashMap_setUserData,	METH_VARARGS,	"Set UserData stored in HashMap"	},
	{	NULL,			NULL,								0,				NULL								}
};

static PyTypeObject HashMapType =
{
	PyObject_HEAD_INIT(NULL)

	.ob_size			= 0,
	.ob_refcnt			= 1,

	.tp_doc				= "Shared HashMap",
	.tp_name			= "shared.HashMap",

	.tp_itemsize		= 0,
	.tp_basicsize		= sizeof(HashMapClass),

	.tp_members			= NULL,
	.tp_methods			= AppMethods,

	.tp_new				= PyType_GenericNew,
	.tp_init			= (initproc)HashMap___init__,
	.tp_dealloc			= (destructor)HashMap___dealloc__,

	.tp_clear			= (inquiry)HashMap___clear__,
	.tp_traverse		= (traverseproc)HashMap___traverse__,

	.tp_flags			= Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
};

/* Pagetable management */

static void *acquire_shm(uint32_t size, int *shmid)
{
	if (shmidCount >= PAGEMAP_MAX_SHMID)
		return NULL;

	int key = 0;
	int count = 0;
	int random = open("/dev/urandom", O_RDONLY);

	while (count < SHM_RETRY && read(random, &key, sizeof(int)) == sizeof(int))
	{
		if (shmget(key, size, 0666) == -1)
		{
			if ((*shmid = shmget(key, size, IPC_CREAT | 0666)) == -1)
			{
				count++;
				continue;
			}

			void *memory = shmat(*shmid, NULL, 0);

			close(random);
			shmidBuffer[shmidCount++] = *shmid;
			return memory;
		}
	}

	close(random);
	return MAP_FAILED;
}

static char acquire_new_bucket(void)
{
	long added = PAGEMAP_INCREMENT / pageSize;

	if (added + pagemap->total > pageCount)
		return 0;

	int shmid = -1;
	void *bucket = acquire_shm(PAGEMAP_INCREMENT, &shmid);

	if (bucket == MAP_FAILED)
		return 0;

	for (long i = 0; i < added; i++)
	{
		pagetable[pagemap->total + i].shmid = shmid;
		pagetable[pagemap->total + i].flags = F_ALLOC;
		pagetable[pagemap->total + i].offset = i * pageSize;
	}

	pagemap->free += added;
	pagemap->total += added;
	return 1;
}

static long acquire_shared_pages(uint32_t needed, int *shmid)
{
	long pageid = -1;
	uint32_t count = 0;

	if (pagemap->free < needed)
	{
		if (!acquire_new_bucket())
		{
			*shmid = -1;
			return -1;
		}
	}

	for (long i = 0; i < pageCount && (pagetable[i].flags & F_ALLOC); i++)
	{
		if (*shmid != pagetable[i].shmid)
		{
			count = 0;
			pageid = -1;
			*shmid = pagetable[i].shmid;
		}

		if (pagetable[i].flags & F_HEAD)
		{
			i += pagetable[i].count - 1;

			if (i >= pageCount || !(pagetable[i].flags & F_TAIL))
			{
				*shmid = -2;
				return -1;
			}

			count = 0;
			pageid = -1;
		}
		else if (!(pagetable[i].flags & ~F_ALLOC))
		{
			count++;
			if (pageid < 0)
				pageid = i;
		}
		else
		{
			*shmid = -3;
			return -1;
		}

		if (count == needed)
		{
			for (long j = 0; j < needed; j++)
			{
				pagetable[pageid + j].count = needed - j;
				pagetable[pageid + j].flags |= F_BODY;
				pagetable[pageid + j].flags |= (j == 0 ? F_HEAD : 0);
				pagetable[pageid + j].flags |= (j == needed - 1 ? F_TAIL : 0);
			}

			pagemap->free -= needed;
			return pageid;
		}
	}

	*shmid = acquire_new_bucket() ? -4 : -5;
	return -1;
}

static char release_shared_pages(long pageid, uint32_t count)
{
	for (long i = pageid; i < count && i < pageCount; i++)
	{
		if (!(pagetable[i].flags & F_BODY))
			return -1;

		if (i == pageid && (!(pagetable[i].flags & F_HEAD) || pagetable[i].count != count))
			return -2;

		if (i == count - 1 && (!(pagetable[i].flags & F_TAIL) || pagetable[i].count != 1))
			return -3;

		pagetable[i].count = 0;
		pagetable[i].flags = F_ALLOC;
	}

	pagemap->free += count;
	return 0;
}

/* Module methods */

static void _pagemap_traverse(const shm_t *data, uint32_t length, shm_t *context)
{
	context->shmid = data->shmid;
	context->pages = data->pages;
	context->pageid = data->pageid;
}

static PyObject *shared_acquire(PyObject *self, PyObject *args)
{
	uint32_t maxCount = 0;
	uint32_t maxItemSize = 0;
	PyStringObject *name = NULL;

	if (!PyArg_ParseTuple(args, "SII", &name, &maxCount, &maxItemSize))
		return NULL;

	if (maxCount <= 0)
	{
		PyErr_SetString(PyExc_ValueError, "Assert `maxCount > 0` failed");
		return NULL;
	}

	if (maxItemSize <= 0)
	{
		PyErr_SetString(PyExc_ValueError, "Assert `maxItemSize > 0` failed");
		return NULL;
	}

	shm_t shm;
	char isCreated = 0;

	char *key = PyString_AS_STRING(name);
	long keySize = PyString_GET_SIZE(name);

	if (keySize > PAGEMAP_KEY_SIZE)
	{
		PyErr_SetString(PyExc_OverflowError, "Key is too long");
		return NULL;
	}

	while (__sync_lock_test_and_set(&(pagemap->lock), F_LOCKED));

	if (hashmap_find(&(pagemap->hashmap), key, keySize, (enumerator_t)_pagemap_traverse, &shm) == E_NO_ENTRY)
	{
		uint32_t length = hashmap_size(maxCount, maxItemSize);

		if (length > PAGEMAP_INCREMENT)
		{
			__sync_lock_release(&(pagemap->lock));
			PyErr_SetString(PyExc_OverflowError, "HashMap too large");
			return NULL;
		}

		int shmid = -4;
		long pageid = -1;
		uint32_t pages = (length - 1) / pageSize + 1;

		while (pageid == -1 && shmid == -4)
			pageid = acquire_shared_pages(pages, &shmid);

		if (pageid < 0)
		{
			__sync_lock_release(&(pagemap->lock));

			switch (shmid)
			{
				case -1:
					PyErr_SetString(PyExc_MemoryError, "FATAL: Allocation of new bucket failed");
					return NULL;

				case -5:
					PyErr_SetString(PyExc_MemoryError, "FATAL: Failed to lookup a continous chunk");
					return NULL;

				case -2:
					PyErr_SetString(PyExc_SystemError, "FATAL: PageTable corrupted(F_TAIL expected)");
					return NULL;

				case -3:
					PyErr_SetString(PyExc_SystemError, "FATAL: PageTable corrupted(unexpected tags)");
					return NULL;

				default:
					PyErr_SetString(PyExc_SystemError, "FATAL: Unknown internal error(should never happen)");
					return NULL;
			}
		}

		isCreated = 1;
		shm.shmid = shmid;
		shm.pages = pages;
		shm.pageid = pageid;
		hashmap_insert(&(pagemap->hashmap), key, keySize, &shm, sizeof(shm_t));
	}

	__sync_lock_release(&(pagemap->lock));

	void *memory = shmat(shm.shmid, NULL, 0);
	hashmap_t *hashmap = ptradd(memory, pagetable[shm.pageid].offset);

	if (memory == MAP_FAILED)
	{
		PyErr_SetString(PyExc_MemoryError, "FATAL: shmat() failed");
		return NULL;
	}

	HashMapClass *result = PyObject_GC_New(HashMapClass, &HashMapType);

	result->name = name;
	result->memory = memory;
	result->hashmap = hashmap;

	result->pages = shm.pages;
	result->pageid = shm.pageid;

	if (isCreated)
		hashmap_init(hashmap, maxCount, maxItemSize);

	Py_INCREF(name);
	hashmap_ref(hashmap);
	return (PyObject *)result;
}

/* Constructor and Destructor */

static int HashMap___init__(HashMapClass *self, PyObject *args, PyObject *kwargs)
{
	PyErr_SetString(PyExc_RuntimeError, "shared.HashMap() cannot be constructed by user");
	return -1;
}

static void HashMap___dealloc__(HashMapClass *self)
{
	HashMap___clear__(self);
	self->ob_type->tp_free((PyObject *)self);
}

/* GC specified methods */

static int HashMap___clear__(HashMapClass *self)
{
	while (__sync_lock_test_and_set(&(pagemap->lock), F_LOCKED));

	if (hashmap_unref(self->hashmap) == 0)
	{
		int result = release_shared_pages(self->pageid, self->pages);
		hashmap_remove(&(pagemap->hashmap), PyString_AS_STRING(self->name), PyString_GET_SIZE(self->name), NULL, NULL);

		if (result < 0)
		{
			__sync_lock_release(&(pagemap->lock));

			switch (result)
			{
				case -1:
					PyErr_SetString(PyExc_SystemError, "FATAL: PageTable corrupted(F_BODY expected)");
					return -1;

				case -2:
					PyErr_SetString(PyExc_SystemError, "FATAL: PageTable corrupted(F_HEAD expected)");
					return -1;

				case -3:
					PyErr_SetString(PyExc_SystemError, "FATAL: PageTable corrupted(F_TAIL expected)");
					return -1;

				default:
					PyErr_SetString(PyExc_SystemError, "FATAL: Unknown internal error(should never happen)");
					return -1;
			}
		}
	}

	__sync_lock_release(&(pagemap->lock));

	shmdt(self->memory);
	Py_CLEAR(self->name);
	return 0;
}

static int HashMap___traverse__(HashMapClass *self, visitproc visit, void *arg)
{
	Py_VISIT(self->name);
	return 0;
}

/* Class methods */

static void _HashMap_builder(const void *data, uint32_t length, PyObject **string)
{
	Py_DECREF(*string);
	*string = PyString_FromStringAndSize(data, length);
}

static void _HashMap_enumerator(const void *data, uint32_t length, PyObject *result)
{
	PyList_Append(result, PyString_FromStringAndSize(data, length));
}

static PyObject *HashMap_keys(HashMapClass *self, PyObject *args)
{
	PyObject *result = PyList_New(0);
	hashmap_enum(self->hashmap, (enumerator_t)_HashMap_enumerator, result);
	return result;
}

static PyObject *HashMap_find(HashMapClass *self, PyObject *args)
{
	PyStringObject *key = NULL;

	if (!PyArg_ParseTuple(args, "S", &key))
		return NULL;

	PyObject *result = (Py_INCREF(Py_None), Py_None);
	hashmap_find(self->hashmap, PyString_AS_STRING(key), PyString_GET_SIZE(key), (enumerator_t)_HashMap_builder, &result);
	return result;
}

static PyObject *HashMap_clear(HashMapClass *self, PyObject *args)
{
	hashmap_wipe(self->hashmap);
	Py_RETURN_NONE;
}

static PyObject *HashMap_insert(HashMapClass *self, PyObject *args)
{
	PyStringObject *key = NULL;
	PyStringObject *value = NULL;

	if (!PyArg_ParseTuple(args, "SS", &key, &value))
		return NULL;

	int keySize = PyString_GET_SIZE(key);
	int valueSize = PyString_GET_SIZE(value);

	switch (hashmap_insert(self->hashmap, PyString_AS_STRING(key), keySize, PyString_AS_STRING(value), valueSize))
	{
		case E_OK:
			Py_RETURN_NONE;

		case E_MAP_FULL:
			PyErr_SetString(PyExc_OverflowError, "HashMap full");
			return NULL;

		case E_TOO_LONG:
			PyErr_SetString(PyExc_OverflowError, "Item out of space");
			return NULL;

		default:
			PyErr_SetString(PyExc_SystemError, "FATAL: Unknown internal error(should never happen)");
			return NULL;
	}
}

static PyObject *HashMap_remove(HashMapClass *self, PyObject *args)
{
	PyStringObject *key = NULL;

	if (!PyArg_ParseTuple(args, "S", &key))
		return NULL;

	hashmap_remove(self->hashmap, PyString_AS_STRING(key), PyString_GET_SIZE(key), NULL, NULL);
	Py_RETURN_NONE;
}

static PyObject *HashMap_lock(HashMapClass *self, PyObject *args)
{
	while (__sync_lock_test_and_set(&(self->hashmap->lock), F_LOCKED));
	Py_RETURN_NONE;
}

static PyObject *HashMap_unlock(HashMapClass *self, PyObject *args)
{
	__sync_lock_release(&(self->hashmap->lock));
	Py_RETURN_NONE;
}

static PyObject *HashMap_userData(HashMapClass *self, PyObject *args)
{
	return PyLong_FromLongLong(self->hashmap->data);
}

static PyObject *HashMap_setUserData(HashMapClass *self, PyObject *args)
{
	int64_t data = 0;

	if (!PyArg_ParseTuple(args, "L", &data))
		return NULL;

	self->hashmap->data = data;
	Py_RETURN_NONE;
}

/* Module Initialization */

static PyMethodDef ModuleMethods[] =
{
	/*	ml_name		ml_meth				ml_flags		ml_doc						*/
	{	"acquire",	shared_acquire,		METH_VARARGS,	"Acquire a shared HashMap"	},
	{	NULL,		NULL,				0,				NULL						}
};

static void free_shms(void)
{
	for (int i = 0; i < shmidCount; i++)
		shmctl(shmidBuffer[i], IPC_RMID, NULL);
}

static char init_memory(void)
{
	pageSize = sysconf(_SC_PAGE_SIZE);

#ifndef __APPLE__
	pageCount = sysconf(_SC_PHYS_PAGES);
#else
	size_t memsize;
	size_t memsizelen = sizeof(size_t);
	sysctlbyname("hw.memsize", &memsize, &memsizelen, NULL, 0);
	pageCount = memsize / pageSize;
#endif

	long pageMapSize = sizeof(pagemap_t) - sizeof(hashmap_t) + hashmap_size(pageCount, PAGEMAP_KEY_SIZE + sizeof(shm_t));

	pagemap = (pagemap_t *)mmap(NULL, pageMapSize, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, 0, 0);

	if (pagemap == (pagemap_t *)MAP_FAILED)
	{
		PyErr_SetString(PyExc_MemoryError, "FATAL: mmap() of PageMap failed");
		return 0;
	}

	pagetable = (pagetable_t *)mmap(NULL, pageCount * sizeof(pagetable_t), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, 0, 0);

	if (pagetable == (pagetable_t *)MAP_FAILED)
	{
		PyErr_SetString(PyExc_MemoryError, "FATAL: mmap() of PageTable failed");
		munmap(pagemap, pageMapSize);
		return 0;
	}

	hashmap_init(&(pagemap->hashmap), pageCount, PAGEMAP_KEY_SIZE + sizeof(shm_t));
	hashmap_ref(&(pagemap->hashmap));
	return 1;
}

PyMODINIT_FUNC init_shared(void)
{
	PyObject *module = Py_InitModule3("_shared", ModuleMethods, "HashMap using shared memory");

	if (init_memory() && PyType_Ready(&HashMapType) >= 0)
	{
		Py_AtExit(free_shms);
		PyModule_AddObject(module, "HashMap", (PyObject *)&HashMapType);
	}
}