#pragma once

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


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

typedef pthread_rwlock_t  rwlock_t;
#define LOCK_INIT_RW(lock) pthread_rwlock_init(lock, NULL)
#define LOCK_RDLOCK_RW(lock) pthread_rwlock_rdlock(lock)
#define LOCK_WRLOCK_RW(lock) pthread_rwlock_wrlock(lock)
#define LOCK_UNLOCK_RW(lock) pthread_rwlock_unlock(lock)

/* 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 {
	rwlock_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_RW(&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_RDLOCK_RW(&bucket->lock);
	for (struct map_node *node = bucket->head; node != NULL; node = node->next) {
		if (node->hash == hash && memcmp(node->key, key, key_len) == 0) {
			
			value          = node->value;
			*ret_value_len = node->value_len;
			goto DONE;
		}
	}

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

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

/**
 * @manage_mvalue manage _ mvalue determines whether the hash table or the caller manages value memory, and true is managed by the hash.
 */
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_WRLOCK_RW(&bucket->lock);

	for (struct map_node *node = bucket->head; node != NULL; node = node->next) {
		if (node->hash == hash && memcmp(node->key, key, key_len) == 0) 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_RW(&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_WRLOCK_RW(&bucket->lock);

	struct map_node *prev = NULL;
    struct map_node *node = bucket->head;
	while (node != NULL) {
        if (node->hash == hash && memcmp(node->key, key, key_len) == 0) {
            if (prev == NULL) { 
                bucket->head = node->next;
            } else {
                prev->next   = node->next;
            }
            free(node->key);
			node->key = NULL; 
			if (node->manage_mvalue) {
        		free(node->value);  
				node->value = NULL;
        	} 
			free(node);
			node = NULL;     
            did_delete = true;
            break;
        }
        prev = node;
        node = node->next;
    }

	LOCK_UNLOCK_RW(&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_WRLOCK_RW(&bucket->lock);

	for (struct map_node *node = bucket->head; node != NULL; node = node->next) {
		if (node->hash == hash && memcmp(node->key, key, key_len) == 0) {
			node->value_len = value_len;
			node->value     = realloc(node->value, value_len);
			assert(node->value);
			if (node->manage_mvalue)
			{
				memcpy(node->value, value, value_len);
			}else  node->value = value;			
			goto DONE;
		}
	}
	panic("map_upsert: key not found");
	/*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_RW(&bucket->lock);
} 

/*static inline void 
map_destroy(struct hashmap *map) {
    for (int i = 0; i < MAP_BUCKET_COUNT; i++) {
        LOCK_LOCK_RW(&map->buckets[i].lock);
        struct map_node *current = map->buckets[i].head;
        struct map_node *next;

        while (current != NULL) {
            next = current->next;
            free(current->key); 
            if (current->manage_mvalue && current->value != NULL) {
                free(current->value); 
            }
            free(current); 
            current = next;
        }

        LOCK_UNLOCK_RW(&map->buckets[i].lock);

        pthread_mutex_destroy(&map->buckets[i].lock);

        // 确保头指针设置为 NULL
        map->buckets[i].head = NULL;
    }
    // 如果 hashmap 是动态分配的，这里还需要释放 hashmap 结构本身
    // free(map);
}*/