#pragma once

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

#include "lock.h"
#include "xmalloc.h"

/* Simple K-V store based on The Practice of Programming by Kernighan and Pike */

/* Bucket count is sized to be a prime that is approximately 20% larger than the desired capacity (6k keys) */
#define MAP_BUCKET_COUNT 7907
#define MAP_HASH         jenkins_hash

struct map_node {
	struct map_node *next;
	char            *key;
	void            *value;
	uint32_t         key_len;
	uint32_t         value_len;
	uint32_t         hash;
	bool             manage_mvalue;
};

struct map_bucket {
	lock_t           lock;
	struct map_node *head;
};

struct hashmap {
	struct map_bucket buckets[MAP_BUCKET_COUNT];
};

static inline void
map_init(struct hashmap *restrict map)
{
	for (int i = 0; i < MAP_BUCKET_COUNT; i++) {
		map->buckets[i].head = NULL;
		LOCK_INIT(&map->buckets[i].lock);
	}
}

/* See https://en.wikipedia.org/wiki/Jenkins_hash_function */
static inline uint32_t
jenkins_hash(char *key, uint32_t key_len)
{
	uint32_t i    = 0;
	uint32_t hash = 0;
	while (i != key_len) {
		hash += key[i++];
		hash += hash << 10;
		hash ^= hash >> 6;
	}
	hash += hash << 3;
	hash ^= hash >> 11;
	hash += hash << 15;
	return hash;
}

static inline void *
map_get(struct hashmap *map, char *key, uint32_t key_len, uint32_t *ret_value_len)
{
	void *value = NULL;

	uint32_t hash = MAP_HASH(key, key_len);

	struct map_bucket *bucket = &map->buckets[hash % MAP_BUCKET_COUNT];

	LOCK_LOCK(&bucket->lock);
	for (struct map_node *node = bucket->head; node != NULL; node = node->next) {
		if (node->hash == hash) {
			value          = node->value;
			*ret_value_len = node->value_len;
			goto DONE;
		}
	}

	if (value == NULL) *ret_value_len = 0;

DONE:
	LOCK_UNLOCK(&bucket->lock);
	return value;
}

static inline bool
map_set(struct hashmap *map, char *key, uint32_t key_len, void *value, uint32_t value_len, bool manage_mvalue)
{
	bool did_set = false;

	uint32_t           hash   = MAP_HASH(key, key_len);
	struct map_bucket *bucket = &map->buckets[hash % MAP_BUCKET_COUNT];

	LOCK_LOCK(&bucket->lock);

	for (struct map_node *node = bucket->head; node != NULL; node = node->next) {
		if (node->hash == hash) goto DONE;
	}

	struct map_node *new_node = (struct map_node *)xmalloc(sizeof(struct map_node));
	*(new_node)               = (struct map_node){ .hash      = hash,
		                                       .key       = xmalloc(key_len),
		                                       .key_len   = key_len,
		                                       .value_len = value_len,
		                                       .next      = bucket->head };

	// Copy Key and Value
	memcpy(new_node->key, key, key_len);
	if (manage_mvalue) {
		new_node->value = xmalloc(value_len);
        memcpy(new_node->value, value, value_len); 
    } else {
        new_node->value = value; 
    }
	new_node->manage_mvalue = manage_mvalue;
	bucket->head = new_node;
	did_set      = true;

DONE:
	LOCK_UNLOCK(&bucket->lock);
	return did_set;
}

/**
 * @returns boolean if node was deleted or not
 */
static inline bool
map_delete(struct hashmap *map, char *key, uint32_t key_len)
{
	bool did_delete = false;

	uint32_t           hash   = MAP_HASH(key, key_len);
	struct map_bucket *bucket = &map->buckets[hash % MAP_BUCKET_COUNT];
	LOCK_LOCK(&bucket->lock);

	struct map_node *prev = bucket->head;
	if (prev != NULL && prev->hash == hash) {
		bucket->head = prev->next;
		free(prev->key);
		if (prev->manage_mvalue) {
        	free(prev->value);  
        }
		free(prev);
		did_delete = true;
		goto DONE;
	}

	for (struct map_node *node = prev->next; node != NULL; prev = node, node = node->next) {
		prev->next = node->next;
		free(node->key);
		if (node->manage_mvalue) {
            free(node->value);  
        }
		free(node);
		did_delete = true;
		goto DONE;
	}

DONE:
	LOCK_UNLOCK(&bucket->lock);
	return did_delete;
}

static inline void
map_upsert(struct hashmap *map, char *key, uint32_t key_len, void *value, uint32_t value_len)
{
	uint32_t           hash   = MAP_HASH(key, key_len);
	struct map_bucket *bucket = &map->buckets[hash % MAP_BUCKET_COUNT];
	LOCK_LOCK(&bucket->lock);

	for (struct map_node *node = bucket->head; node != NULL; node = node->next) {
		if (node->hash == hash) {
			node->value_len = value_len;
			node->value     = realloc(node->value, value_len);
            //node->value = value;
			assert(node->value);
			if (node->manage_mvalue)
			{
				memcpy(node->value, value, value_len);
			}			
			goto DONE;
		}
	}

	struct map_node *new_node = (struct map_node *)xmalloc(sizeof(struct map_node));

	*(new_node) = (struct map_node){ .hash      = hash,
		                         .key       = xmalloc(key_len),
		                         .key_len   = key_len,
		                         .value     = xmalloc(value_len),
		                         .value_len = value_len,
		                         .next      = bucket->head };

	assert(new_node->key);
	assert(new_node->value);

	// Copy Key and Value
	memcpy(new_node->key, key, key_len);
	memcpy(new_node->value, value, value_len);

	bucket->head = new_node;

DONE:
	LOCK_UNLOCK(&bucket->lock);
}