#include <blue/core/misc.h>
#include <blue/core/status.h>
#include <blue/ds/hashmap.h>
#include <blue/ds/string.h>
#include <stdbool.h>
#include <stdlib.h>

#define HASH_OFFSET_BASIS 0xcbf29ce484222325
#define HASH_PRIME        0x100000001b3

/*** PRIVATE DATA *************************************************************/

struct b_hashmap_bucket_item {
	struct b_queue_entry bi_entry;
	struct b_hashmap_key bi_key;
	struct b_hashmap_value bi_value;
};

struct b_hashmap_bucket {
	struct b_btree_node bk_node;
	uint64_t bk_hash;
	struct b_queue bk_items;
};

struct b_hashmap_p {
	size_t h_count;
	struct b_btree h_buckets;
	b_hashmap_key_destructor h_key_dtor;
	b_hashmap_value_destructor h_value_dtor;
};

struct b_hashmap_iterator_p {
	size_t i;
	b_hashmap_item item;

	b_hashmap *_h;
	struct b_hashmap_p *_h_p;
	b_btree_node *_cbn;
	b_queue_entry *_cqe;
};

/*** PRIVATE FUNCTIONS ********************************************************/

static B_BTREE_DEFINE_SIMPLE_GET(
	struct b_hashmap_bucket, uint64_t, bk_node, bk_hash, get_bucket);
static B_BTREE_DEFINE_SIMPLE_INSERT(
	struct b_hashmap_bucket, bk_node, bk_hash, put_bucket);

static uint64_t hash_data(const void *p, size_t size)
{
	const unsigned char *s = p;
	uint64_t hash = HASH_OFFSET_BASIS;

	for (size_t i = 0; s[i]; i++) {
		hash *= HASH_PRIME;
		hash ^= s[i];
	}

	return hash;
}

static uint64_t hash_key(const struct b_hashmap_key *key)
{
	if (key->key_flags & B_HASHMAP_KEY_F_INTVALUE) {
		return hash_data(&key->key_data, sizeof key->key_data);
	} else {
		return hash_data(key->key_data, key->key_size);
	}
}

static bool compare_key(
	const struct b_hashmap_key *a, const struct b_hashmap_key *b)
{
	const void *a_data = NULL, *b_data = NULL;
	size_t a_len = 0, b_len = 0;

	if (a->key_flags & B_HASHMAP_KEY_F_INTVALUE) {
		a_data = &a->key_data;
		a_len = sizeof a->key_data;
	} else {
		a_data = a->key_data;
		a_len = a->key_size;
	}

	if (b->key_flags & B_HASHMAP_KEY_F_INTVALUE) {
		b_data = &b->key_data;
		b_len = sizeof b->key_data;
	} else {
		b_data = b->key_data;
		b_len = b->key_size;
	}

	if (a_len != b_len) {
		return false;
	}

	size_t cmp_len = a_len;
	return memcmp(a_data, b_data, cmp_len) == 0;
}

static bool get_next_node(
	struct b_btree_node *cur_node, struct b_queue_entry *cur_entry,
	struct b_btree_node **out_next_node, struct b_queue_entry **out_next_entry)
{
	struct b_hashmap_bucket *cur_bucket
		= b_unbox(struct b_hashmap_bucket, cur_node, bk_node);
	if (!cur_bucket) {
		return false;
	}

	struct b_hashmap_bucket_item *cur_item
		= b_unbox(struct b_hashmap_bucket_item, cur_entry, bi_entry);
	if (!cur_item) {
		return false;
	}

	struct b_btree_node *next_node = cur_node;
	struct b_queue_entry *next_entry = b_queue_next(cur_entry);
	if (!next_entry) {
		next_node = b_btree_next(cur_node);
		if (!next_node) {
			return false;
		}

		struct b_hashmap_bucket *next_bucket
			= b_unbox(struct b_hashmap_bucket, next_node, bk_node);
		if (!next_bucket) {
			return false;
		}

		next_entry = b_queue_first(&next_bucket->bk_items);
		if (!next_entry) {
			return false;
		}
	}

	struct b_hashmap_bucket_item *next_item
		= b_unbox(struct b_hashmap_bucket_item, next_entry, bi_entry);
	if (!next_item) {
		return false;
	}

	*out_next_node = next_node;
	*out_next_entry = next_entry;

	return true;
}

static struct b_hashmap_bucket *create_bucket(void)
{
	struct b_hashmap_bucket *bucket = malloc(sizeof *bucket);
	if (!bucket) {
		return NULL;
	}

	memset(bucket, 0x0, sizeof *bucket);
	return bucket;
}

static struct b_hashmap_bucket_item *create_bucket_item(void)
{
	struct b_hashmap_bucket_item *item = malloc(sizeof *item);
	if (!item) {
		return NULL;
	}

	memset(item, 0x0, sizeof *item);
	return item;
}

static b_status hashmap_put(
	struct b_hashmap_p *hashmap, const b_hashmap_key *key,
	const b_hashmap_value *value)
{
	uint64_t hash = hash_key(key);
	struct b_hashmap_bucket *bucket = get_bucket(&hashmap->h_buckets, hash);

	if (!bucket) {
		bucket = create_bucket();
		if (!bucket) {
			return B_ERR_NO_MEMORY;
		}

		bucket->bk_hash = hash;
		put_bucket(&hashmap->h_buckets, bucket);
	}

	struct b_queue_entry *entry = b_queue_first(&bucket->bk_items);
	while (entry) {
		struct b_hashmap_bucket_item *item
			= b_unbox(struct b_hashmap_bucket_item, entry, bi_entry);

		if (compare_key(&item->bi_key, key)) {
			memcpy(&item->bi_value, value, sizeof *value);
			return B_SUCCESS;
		}

		entry = b_queue_next(entry);
	}

	struct b_hashmap_bucket_item *item = create_bucket_item();
	if (!item) {
		return B_ERR_NO_MEMORY;
	}

	memcpy(&item->bi_key, key, sizeof *key);
	memcpy(&item->bi_value, value, sizeof *value);

	b_queue_push_back(&bucket->bk_items, &item->bi_entry);
	hashmap->h_count++;

	return B_SUCCESS;
}

static const struct b_hashmap_value *hashmap_get(
	const struct b_hashmap_p *hashmap, const struct b_hashmap_key *key)
{
	uint64_t hash = hash_key(key);

	struct b_hashmap_bucket *bucket = get_bucket(&hashmap->h_buckets, hash);
	if (!bucket) {
		return NULL;
	}

	struct b_queue_entry *entry = b_queue_first(&bucket->bk_items);
	while (entry) {
		struct b_hashmap_bucket_item *item
			= b_unbox(struct b_hashmap_bucket_item, entry, bi_entry);

		if (compare_key(&item->bi_key, key)) {
			return &item->bi_value;
		}

		entry = b_queue_next(entry);
	}

	return NULL;
}

static bool hashmap_has_key(
	const struct b_hashmap_p *hashmap, const b_hashmap_key *key)
{
	uint64_t hash = hash_key(key);
	struct b_hashmap_bucket *bucket = get_bucket(&hashmap->h_buckets, hash);
	if (!bucket) {
		return false;
	}

	struct b_queue_entry *entry = b_queue_first(&bucket->bk_items);
	while (entry) {
		struct b_hashmap_bucket_item *item
			= b_unbox(struct b_hashmap_bucket_item, entry, bi_entry);

		if (compare_key(&item->bi_key, key)) {
			return true;
		}

		entry = b_queue_next(entry);
	}

	return false;
}

static size_t hashmap_get_size(const struct b_hashmap_p *hashmap)
{
	return hashmap->h_count;
}

static bool hashmap_is_empty(const struct b_hashmap_p *hashmap)
{
	b_btree_node *first_node = b_btree_first(&hashmap->h_buckets);
	struct b_hashmap_bucket *first_bucket
		= b_unbox(struct b_hashmap_bucket, first_node, bk_node);
	if (!first_bucket) {
		return true;
	}

	b_queue_entry *first_entry = b_queue_first(&first_bucket->bk_items);
	struct b_hashmap_bucket_item *first_item
		= b_unbox(struct b_hashmap_bucket_item, first_entry, bi_entry);
	if (!first_item) {
		return true;
	}

	return false;
}

static b_status delete_item(
	struct b_hashmap_p *hashmap, struct b_hashmap_bucket *bucket,
	struct b_hashmap_bucket_item *item)
{
	b_queue_delete(&bucket->bk_items, &item->bi_entry);

	if (hashmap->h_key_dtor) {
		hashmap->h_key_dtor((void *)item->bi_key.key_data);
	}

	if (hashmap->h_value_dtor) {
		hashmap->h_value_dtor((void *)item->bi_value.value_data);
	}

	free(item);

	if (b_queue_empty(&bucket->bk_items)) {
		b_btree_delete(&hashmap->h_buckets, &bucket->bk_node);
		free(bucket);
	}

	hashmap->h_count--;
	return B_SUCCESS;
}

/*** PUBLIC FUNCTIONS *********************************************************/

b_hashmap *b_hashmap_create(
	b_hashmap_key_destructor key_dtor, b_hashmap_value_destructor value_dtor)
{
	b_hashmap *hashmap = b_object_create(B_TYPE_HASHMAP);
	if (!hashmap) {
		return NULL;
	}

	return hashmap;
}

b_hashmap *b_hashmap_create_with_items(const b_hashmap_item *items)
{
	b_hashmap *hashmap = b_hashmap_create(NULL, NULL);
	if (!hashmap) {
		return NULL;
	}

	struct b_hashmap_p *p = b_object_get_private(hashmap, B_TYPE_HASHMAP);

	for (size_t i = 0; items[i].key.key_data && items[i].key.key_size; i++) {
		hashmap_put(p, &items[i].key, &items[i].value);
	}

	return hashmap;
}

b_status b_hashmap_put(
	b_hashmap *hashmap, const b_hashmap_key *key, const b_hashmap_value *value)
{
	B_CLASS_DISPATCH_STATIC(B_TYPE_HASHMAP, hashmap_put, hashmap, key, value);
}

const struct b_hashmap_value *b_hashmap_get(
	const b_hashmap *hashmap, const struct b_hashmap_key *key)
{
	B_CLASS_DISPATCH_STATIC(B_TYPE_HASHMAP, hashmap_get, hashmap, key);
}

bool b_hashmap_has_key(const b_hashmap *hashmap, const b_hashmap_key *key)
{
	B_CLASS_DISPATCH_STATIC(B_TYPE_HASHMAP, hashmap_has_key, hashmap, key);
}

size_t b_hashmap_get_size(const b_hashmap *hashmap)
{
	B_CLASS_DISPATCH_STATIC_0(B_TYPE_HASHMAP, hashmap_get_size, hashmap);
}

bool b_hashmap_is_empty(const b_hashmap *hashmap)
{
	B_CLASS_DISPATCH_STATIC_0(B_TYPE_HASHMAP, hashmap_is_empty, hashmap);
}

/*** VIRTUAL FUNCTIONS ********************************************************/

static void hashmap_init(b_object *obj, void *priv)
{
	struct b_hashmap_p *map = priv;
}

static void hashmap_fini(b_object *obj, void *priv)
{
	struct b_hashmap_p *map = priv;

	struct b_btree_node *node = b_btree_first(&map->h_buckets);
	while (node) {
		struct b_hashmap_bucket *b
			= b_unbox(struct b_hashmap_bucket, node, bk_node);
		struct b_btree_node *next_node = b_btree_next(node);
		b_btree_delete(&map->h_buckets, node);

		struct b_queue_entry *entry = b_queue_first(&b->bk_items);
		while (entry) {
			struct b_hashmap_bucket_item *item = b_unbox(
				struct b_hashmap_bucket_item, entry, bi_entry);
			struct b_queue_entry *next_entry = b_queue_next(entry);
			b_queue_delete(&b->bk_items, entry);

			if (map->h_key_dtor) {
				map->h_key_dtor((void *)item->bi_key.key_data);
			}

			if (map->h_value_dtor) {
				map->h_value_dtor((void *)item->bi_value.value_data);
			}

			free(item);
			entry = next_entry;
		}

		free(b);
		node = next_node;
	}
}

/*** ITERATOR FUNCTIONS *******************************************************/

static b_iterator *iterable_begin(b_hashmap *hashmap)
{
	b_hashmap_iterator *it_obj = b_object_create(B_TYPE_HASHMAP_ITERATOR);
	struct b_hashmap_iterator_p *it
		= b_object_get_private(it_obj, B_TYPE_HASHMAP_ITERATOR);

	it->_h = hashmap;
	it->_h_p = b_object_get_private(hashmap, B_TYPE_HASHMAP);

	it->i = 0;
	if (b_hashmap_is_empty(hashmap)) {
		memset(&it->item, 0x0, sizeof it->item);
		b_iterator_set_status(it_obj, B_ERR_NO_DATA);
		return it_obj;
	}

	struct b_btree_node *first_node = b_btree_first(&it->_h_p->h_buckets);
	struct b_hashmap_bucket *first_bucket
		= b_unbox(struct b_hashmap_bucket, first_node, bk_node);
	if (!first_bucket) {
		memset(&it->item, 0x0, sizeof it->item);
		b_iterator_set_status(it_obj, B_ERR_NO_DATA);
		return it_obj;
	}

	struct b_queue_entry *first_entry = b_queue_first(&first_bucket->bk_items);
	struct b_hashmap_bucket_item *first_item
		= b_unbox(struct b_hashmap_bucket_item, first_entry, bi_entry);
	if (!first_item) {
		memset(&it->item, 0x0, sizeof it->item);
		b_iterator_set_status(it_obj, B_ERR_NO_DATA);
		return it_obj;
	}

	memcpy(&it->item.key, &first_item->bi_key, sizeof it->item.key);
	memcpy(&it->item.value, &first_item->bi_value, sizeof it->item.value);

	it->_cbn = first_node;
	it->_cqe = first_entry;

	return 0;
}

static const b_iterator *iterable_cbegin(const b_hashmap *hashmap)
{
	return iterable_begin((b_hashmap *)hashmap);
}

static enum b_status iterator_move_next(const b_iterator *obj)
{
	struct b_hashmap_iterator_p *it
		= b_object_get_private(obj, B_TYPE_HASHMAP_ITERATOR);

	struct b_btree_node *next_node;
	struct b_queue_entry *next_entry;
	if (!get_next_node(it->_cbn, it->_cqe, &next_node, &next_entry)) {
		memset(&it->item, 0x0, sizeof it->item);
		return B_ERR_NO_DATA;
	}

	struct b_hashmap_bucket_item *next_item
		= b_unbox(struct b_hashmap_bucket_item, next_entry, bi_entry);

	if (!next_item) {
		memset(&it->item, 0x0, sizeof it->item);
		return B_ERR_NO_DATA;
	}

	it->i++;
	memcpy(&it->item.key, &next_item->bi_key, sizeof it->item.key);
	memcpy(&it->item.value, &next_item->bi_value, sizeof it->item.value);

	it->_cbn = next_node;
	it->_cqe = next_entry;

	return true;
}

static enum b_status iterator_erase(b_iterator *obj)
{
	struct b_hashmap_iterator_p *it
		= b_object_get_private(obj, B_TYPE_HASHMAP_ITERATOR);

	if ((it->item.key.key_data || it->item.value.value_data)
	    && !(it->_cbn && it->_cqe)) {
		return B_ERR_BAD_STATE;
	}

	if (!it->item.key.key_data || !it->_cqe) {
		return B_ERR_NO_DATA;
	}

	struct b_btree_node *next_node;
	struct b_queue_entry *next_entry;
	if (!get_next_node(it->_cbn, it->_cqe, &next_node, &next_entry)) {
		memset(&it->item, 0x0, sizeof it->item);
		return B_ERR_NO_DATA;
	}

	struct b_hashmap_bucket *cur_bucket
		= b_unbox(struct b_hashmap_bucket, it->_cbn, bk_node);
	struct b_hashmap_bucket_item *cur_item
		= b_unbox(struct b_hashmap_bucket_item, it->_cqe, bi_entry);

	struct b_hashmap_bucket_item *next_item
		= b_unbox(struct b_hashmap_bucket_item, next_entry, bi_entry);

	b_status status = delete_item(it->_h_p, cur_bucket, cur_item);
	if (B_ERR(status)) {
		return status;
	}

	if (next_item) {
		memcpy(&it->item.key, &next_item->bi_key, sizeof it->item.key);
		memcpy(&it->item.value, &next_item->bi_value,
		       sizeof it->item.value);

		it->_cbn = next_node;
		it->_cqe = next_entry;
	} else {
		memset(&it->item, 0x0, sizeof it->item);

		it->_cbn = NULL;
		it->_cqe = NULL;
	}

	return B_SUCCESS;
}

static b_iterator_value iterator_get_value(b_iterator *obj)
{
	struct b_hashmap_iterator_p *it
		= b_object_get_private(obj, B_TYPE_HASHMAP_ITERATOR);
	return B_ITERATOR_VALUE_PTR(&it->item);
}

static const b_iterator_value iterator_get_cvalue(const b_iterator *obj)
{
	const struct b_hashmap_iterator_p *it
		= b_object_get_private(obj, B_TYPE_HASHMAP_ITERATOR);
	return B_ITERATOR_VALUE_CPTR(&it->item);
}

/*** CLASS DEFINITION *********************************************************/

// ---- b_hashmap DEFINITION
B_TYPE_CLASS_DEFINITION_BEGIN(b_hashmap)
	B_TYPE_CLASS_INTERFACE_BEGIN(b_object, B_TYPE_OBJECT)
		B_INTERFACE_ENTRY(to_string) = NULL;
	B_TYPE_CLASS_INTERFACE_END(b_object, B_TYPE_OBJECT)

	B_TYPE_CLASS_INTERFACE_BEGIN(b_iterable, B_TYPE_ITERABLE)
		B_INTERFACE_ENTRY(it_begin) = iterable_begin;
		B_INTERFACE_ENTRY(it_cbegin) = iterable_cbegin;
	B_TYPE_CLASS_INTERFACE_END(b_iterable, B_TYPE_ITERABLE)
B_TYPE_CLASS_DEFINITION_END(b_hashmap)

B_TYPE_DEFINITION_BEGIN(b_hashmap)
	B_TYPE_ID(0x7bf5bcd1, 0x1ff3, 0x4e43, 0xbed8, 0x7c74f28348bf);
	B_TYPE_CLASS(b_hashmap_class);
	B_TYPE_IMPLEMENTS(B_TYPE_ITERABLE);
	B_TYPE_INSTANCE_PRIVATE(struct b_hashmap_p);
	B_TYPE_INSTANCE_INIT(hashmap_init);
	B_TYPE_INSTANCE_FINI(hashmap_fini);
B_TYPE_DEFINITION_END(b_hashmap)

// ---- b_hashmap_iterator DEFINITION
B_TYPE_CLASS_DEFINITION_BEGIN(b_hashmap_iterator)
	B_TYPE_CLASS_INTERFACE_BEGIN(b_object, B_TYPE_OBJECT)
		B_INTERFACE_ENTRY(to_string) = NULL;
	B_TYPE_CLASS_INTERFACE_END(b_object, B_TYPE_OBJECT)

	B_TYPE_CLASS_INTERFACE_BEGIN(b_iterator, B_TYPE_ITERATOR)
		B_INTERFACE_ENTRY(it_move_next) = iterator_move_next;
		B_INTERFACE_ENTRY(it_erase) = iterator_erase;
		B_INTERFACE_ENTRY(it_get_value) = iterator_get_value;
		B_INTERFACE_ENTRY(it_get_cvalue) = iterator_get_cvalue;
	B_TYPE_CLASS_INTERFACE_END(b_iterator, B_TYPE_ITERATOR)
B_TYPE_CLASS_DEFINITION_END(b_hashmap_iterator)

B_TYPE_DEFINITION_BEGIN(b_hashmap_iterator)
	B_TYPE_ID(0xd9658456, 0xdd80, 0x419a, 0xb23a, 0xb513013e6431);
	B_TYPE_EXTENDS(B_TYPE_ITERATOR);
	B_TYPE_CLASS(b_hashmap_iterator_class);
	B_TYPE_INSTANCE_PRIVATE(struct b_hashmap_iterator_p);
B_TYPE_DEFINITION_END(b_hashmap_iterator)