#include <fx/core/status.h>
#include <fx/core/stream.h>
#include <fx/ds/dict.h>
#include <fx/ds/string.h>
#include <stdbool.h>
#include <stdlib.h>

#define HASH_OFFSET_BASIS 0xcbf29ce484222325
#define HASH_PRIME        0x100000001b3

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

struct fx_dict_bucket_item {
	fx_queue_entry bi_entry;
	fx_string *bi_str;
	fx_object *bi_value;
};

struct fx_dict_bucket {
	fx_bst_node bk_node;
	uint64_t bk_hash;
	fx_queue bk_items;
};

struct fx_dict_p {
	fx_bst d_buckets;
};

struct fx_dict_iterator_p {
	size_t i;
	fx_dict_item item;

	fx_dict *_d;
	struct fx_dict_p *_d_p;
	fx_bst_node *_cbn;
	fx_queue_entry *_cqe;
};

/*** MISC FUNCTIONS ***********************************************************/

static FX_BST_DEFINE_SIMPLE_GET(
	struct fx_dict_bucket, uint64_t, bk_node, bk_hash, get_bucket);
static FX_BST_DEFINE_SIMPLE_INSERT(
	struct fx_dict_bucket, bk_node, bk_hash, put_bucket);

uint64_t fx_cstr_hash(const char *s)
{
	uint64_t hash = HASH_OFFSET_BASIS;

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

	return hash;
}

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

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

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

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

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

static fx_status dict_put(struct fx_dict_p *dict, const char *key, fx_object *value)
{
	uint64_t hash = fx_cstr_hash(key);
	struct fx_dict_bucket *bucket = get_bucket(&dict->d_buckets, hash);
	if (!bucket) {
		bucket = create_bucket();
		if (!bucket) {
			return FX_ERR_NO_MEMORY;
		}

		bucket->bk_hash = hash;
		put_bucket(&dict->d_buckets, bucket);
	}

	struct fx_dict_bucket_item *item = create_bucket_item();
	if (!item) {
		return FX_ERR_NO_MEMORY;
	}

	item->bi_str = fx_string_create_from_cstr(key);
	item->bi_value = fx_object_ref(value);

	fx_queue_push_back(&bucket->bk_items, &item->bi_entry);

	return FX_SUCCESS;
}

static fx_status dict_put_sk(
	struct fx_dict_p *dict, const fx_string *key, fx_object *value)
{
	uint64_t hash = fx_string_hash(key);
	struct fx_dict_bucket *bucket = get_bucket(&dict->d_buckets, hash);
	if (!bucket) {
		bucket = create_bucket();
		if (!bucket) {
			return FX_ERR_NO_MEMORY;
		}

		bucket->bk_hash = hash;
		put_bucket(&dict->d_buckets, bucket);
	}

	struct fx_dict_bucket_item *item = create_bucket_item();
	if (!item) {
		return FX_ERR_NO_MEMORY;
	}

	item->bi_str = fx_string_duplicate(key);
	item->bi_value = fx_object_ref(value);

	fx_queue_push_back(&bucket->bk_items, &item->bi_entry);

	return FX_SUCCESS;
}

static fx_object *dict_at(const struct fx_dict_p *dict, const char *key)
{
	uint64_t hash = fx_cstr_hash(key);
	struct fx_dict_bucket *bucket = get_bucket(&dict->d_buckets, hash);
	if (!bucket) {
		return NULL;
	}

	struct fx_queue_entry *entry = fx_queue_first(&bucket->bk_items);
	while (entry) {
		struct fx_dict_bucket_item *item
			= fx_unbox(struct fx_dict_bucket_item, entry, bi_entry);

		if (!strcmp(fx_string_ptr(item->bi_str), key)) {
			return item->bi_value;
		}

		entry = fx_queue_next(entry);
	}

	return NULL;
}

static fx_object *dict_at_sk(const struct fx_dict_p *dict, const fx_string *key)
{
	uint64_t hash = fx_string_hash(key);
	struct fx_dict_bucket *bucket = get_bucket(&dict->d_buckets, hash);
	if (!bucket) {
		return NULL;
	}

	struct fx_queue_entry *entry = fx_queue_first(&bucket->bk_items);
	while (entry) {
		struct fx_dict_bucket_item *item
			= fx_unbox(struct fx_dict_bucket_item, entry, bi_entry);

		if (fx_string_compare(item->bi_str, key)) {
			return item->bi_value;
		}

		entry = fx_queue_next(entry);
	}

	return NULL;
}

static fx_object *dict_get(struct fx_dict_p *dict, const char *key)
{
	fx_object *value = dict_at(dict, key);
	if (value) {
		fx_object_ref(value);
	}

	return value;
}

static fx_object *dict_get_sk(struct fx_dict_p *dict, const fx_string *key)
{
	fx_object *value = dict_at_sk(dict, key);
	if (value) {
		fx_object_ref(value);
	}

	return value;
}

static bool dict_has_key(const struct fx_dict_p *dict, const char *key)
{
	return dict_at(dict, key) != NULL;
}

static bool dict_has_skey(const struct fx_dict_p *dict, const fx_string *key)
{
	return dict_at_sk(dict, key) != NULL;
}

static size_t dict_get_size(const struct fx_dict_p *dict)
{
	size_t count = 0;
#if 0
	fx_bst_iterator it1;
	fx_queue_iterator it2;

	fx_bst_foreach (&it1, &dict->d_buckets) {
		struct fx_dict_bucket *bucket
			= fx_unbox(struct fx_dict_bucket, it1.node, bk_node);

		fx_queue_foreach (&it2, &bucket->bk_items) {
			count++;
		}
	}
#endif

	return count;
}

static bool dict_is_empty(const struct fx_dict_p *dict)
{
	fx_bst_node *first_node = fx_bst_first(&dict->d_buckets);
	struct fx_dict_bucket *first_bucket
		= fx_unbox(struct fx_dict_bucket, first_node, bk_node);
	if (!first_bucket) {
		return true;
	}

	fx_queue_entry *first_entry = fx_queue_first(&first_bucket->bk_items);
	struct fx_dict_bucket_item *first_item
		= fx_unbox(struct fx_dict_bucket_item, first_entry, bi_entry);
	if (!first_item) {
		return true;
	}

	return false;
}

static bool get_next_node(
	struct fx_bst_node *cur_node, struct fx_queue_entry *cur_entry,
	struct fx_bst_node **out_next_node, struct fx_queue_entry **out_next_entry)
{
	struct fx_dict_bucket *cur_bucket
		= fx_unbox(struct fx_dict_bucket, cur_node, bk_node);
	if (!cur_bucket) {
		return false;
	}

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

	struct fx_bst_node *next_node = cur_node;
	struct fx_queue_entry *next_entry = fx_queue_next(cur_entry);
	if (!next_entry) {
		next_node = fx_bst_next(cur_node);
		if (!next_node) {
			return false;
		}

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

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

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

	*out_next_node = next_node;
	*out_next_entry = next_entry;

	return true;
}

static fx_status delete_item(
	struct fx_dict_p *dict, struct fx_dict_bucket *bucket,
	struct fx_dict_bucket_item *item)
{
	fx_queue_delete(&bucket->bk_items, &item->bi_entry);
	free(item);

	if (fx_queue_empty(&bucket->bk_items)) {
		fx_bst_delete(&dict->d_buckets, &bucket->bk_node);
		free(bucket);
	}

	return FX_SUCCESS;
}

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

#if 0
fx_dict *fx_dict_create_with_items(const fx_dict_item *items)
{
	fx_dict *dict = fx_dict_create();
	if (!dict) {
		return NULL;
	}

	struct fx_dict_p *p = fx_object_get_private(dict, FX_TYPE_DICT);

	for (size_t i = 0; items[i].key; i++) {
		dict_put(p, items[i].key, items[i].value);
	}

	return dict;
}
#endif

fx_status fx_dict_put(fx_dict *dict, const char *key, fx_object *value)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_put, dict, key, value);
}

fx_status fx_dict_put_sk(fx_dict *dict, const fx_string *key, fx_object *value)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_put_sk, dict, key, value);
}

fx_object *fx_dict_at(const fx_dict *dict, const char *key)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_at, dict, key);
}

fx_object *fx_dict_at_sk(const fx_dict *dict, const fx_string *key)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_at_sk, dict, key);
}

fx_object *fx_dict_get(fx_dict *dict, const char *key)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_get, dict, key);
}

fx_object *fx_dict_get_sk(fx_dict *dict, const fx_string *key)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_get_sk, dict, key);
}

bool fx_dict_has_key(const fx_dict *dict, const char *key)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_has_key, dict, key);
}

bool fx_dict_has_skey(const fx_dict *dict, const fx_string *key)
{
	FX_CLASS_DISPATCH_STATIC(FX_TYPE_DICT, dict_has_skey, dict, key);
}

size_t fx_dict_get_size(const fx_dict *dict)
{
	FX_CLASS_DISPATCH_STATIC_0(FX_TYPE_DICT, dict_get_size, dict);
}

bool fx_dict_is_empty(const fx_dict *dict)
{
	FX_CLASS_DISPATCH_STATIC_0(FX_TYPE_DICT, dict_is_empty, dict);
}

/*** PUBLIC ALIAS FUNCTIONS ***************************************************/
/*** VIRTUAL FUNCTIONS ********************************************************/

static void dict_init(fx_object *obj, void *priv)
{
	struct fx_dict_p *dict = priv;
}

static void dict_fini(fx_object *obj, void *priv)
{
	struct fx_dict_p *dict = priv;

	struct fx_bst_node *node = fx_bst_first(&dict->d_buckets);
	while (node) {
		struct fx_dict_bucket *bucket
			= fx_unbox(struct fx_dict_bucket, node, bk_node);
		struct fx_bst_node *next_node = fx_bst_next(node);
		fx_bst_delete(&dict->d_buckets, node);

		struct fx_queue_entry *entry = fx_queue_first(&bucket->bk_items);
		while (entry) {
			struct fx_dict_bucket_item *item = fx_unbox(
				struct fx_dict_bucket_item, entry, bi_entry);
			struct fx_queue_entry *next_entry = fx_queue_next(entry);
			fx_queue_delete(&bucket->bk_items, entry);

			free(item->bi_str);
			fx_object_unref(item->bi_value);
			free(item);

			entry = next_entry;
		}

		free(bucket);
		node = next_node;
	}
}

static void dict_to_string(const fx_object *obj, fx_stream *out)
{
	struct fx_dict_p *dict = fx_object_get_private(obj, FX_TYPE_DICT);

	if (dict_is_empty(dict)) {
		fx_stream_write_string(out, "{}", NULL);
		return;
	}

	fx_stream_write_string(out, "{\n", NULL);

	fx_stream_push_indent(out, 1);
	size_t len = dict_get_size(dict);
	size_t i = 0;

	struct fx_bst_node *node = fx_bst_first(&dict->d_buckets);
	while (node) {
		struct fx_dict_bucket *bucket
			= fx_unbox(struct fx_dict_bucket, node, bk_node);

		struct fx_queue_entry *entry = fx_queue_first(&bucket->bk_items);
		while (entry) {
			struct fx_dict_bucket_item *item = fx_unbox(
				struct fx_dict_bucket_item, entry, bi_entry);

			fx_object_to_string(item->bi_str, out);
			fx_stream_write_string(out, ": ", NULL);

			bool is_string
				= fx_object_is_type(item->bi_value, FX_TYPE_STRING);

			if (is_string) {
				fx_stream_write_char(out, '"');
			}

			fx_object_to_string(item->bi_value, out);

			if (is_string) {
				fx_stream_write_char(out, '"');
			}

			if (i < len - 1) {
				fx_stream_write_string(out, ",", NULL);
			}

			fx_stream_write_char(out, '\n');
			entry = fx_queue_next(entry);
			i++;
		}
		node = fx_bst_next(node);
	}

	fx_stream_pop_indent(out);
	fx_stream_write_char(out, '}');
}

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

static fx_iterator *iterable_begin(fx_dict *dict)
{
	fx_iterator *it_obj = fx_object_create(FX_TYPE_DICT_ITERATOR);
	if (!it_obj) {
		return NULL;
	}

	struct fx_dict_iterator_p *it
		= fx_object_get_private(it_obj, FX_TYPE_DICT_ITERATOR);

	it->i = 0;
	it->_d = dict;
	it->_d_p = fx_object_get_private(dict, FX_TYPE_DICT);

	if (dict_is_empty(it->_d_p)) {
		it->item.key = NULL;
		it->item.value = NULL;
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
		return it_obj;
	}

	fx_bst_node *first_node = fx_bst_first(&it->_d_p->d_buckets);
	struct fx_dict_bucket *first_bucket
		= fx_unbox(struct fx_dict_bucket, first_node, bk_node);
	if (!first_bucket) {
		it->item.key = NULL;
		it->item.value = NULL;
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
		return it_obj;
	}

	fx_queue_entry *first_entry = fx_queue_first(&first_bucket->bk_items);
	struct fx_dict_bucket_item *first_item
		= fx_unbox(struct fx_dict_bucket_item, first_entry, bi_entry);
	if (!first_item) {
		it->item.key = NULL;
		it->item.value = NULL;
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
		return it_obj;
	}

	it->item.key = first_item->bi_str;
	it->item.value = first_item->bi_value;

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

	return it_obj;
}

static const fx_iterator *iterable_cbegin(const fx_dict *dict)
{
	fx_iterator *it_obj = fx_object_create(FX_TYPE_DICT_ITERATOR);
	if (!it_obj) {
		return NULL;
	}

	struct fx_dict_iterator_p *it
		= fx_object_get_private(it_obj, FX_TYPE_DICT_ITERATOR);

	it->i = 0;
	it->_d = (fx_dict *)dict;
	it->_d_p = fx_object_get_private(dict, FX_TYPE_DICT);

	if (dict_is_empty(it->_d_p)) {
		it->item.key = NULL;
		it->item.value = NULL;
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
		return it_obj;
	}

	fx_bst_node *first_node = fx_bst_first(&it->_d_p->d_buckets);
	struct fx_dict_bucket *first_bucket
		= fx_unbox(struct fx_dict_bucket, first_node, bk_node);
	if (!first_bucket) {
		it->item.key = NULL;
		it->item.value = NULL;
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
		return it_obj;
	}

	fx_queue_entry *first_entry = fx_queue_first(&first_bucket->bk_items);
	struct fx_dict_bucket_item *first_item
		= fx_unbox(struct fx_dict_bucket_item, first_entry, bi_entry);
	if (!first_item) {
		it->item.key = NULL;
		it->item.value = NULL;
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
		return it_obj;
	}

	it->item.key = first_item->bi_str;
	it->item.value = first_item->bi_value;

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

	return it_obj;
}

static enum fx_status iterator_move_next(const fx_iterator *obj)
{
	struct fx_dict_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_DICT_ITERATOR);

	struct fx_bst_node *next_node;
	struct fx_queue_entry *next_entry;
	if (!get_next_node(it->_cbn, it->_cqe, &next_node, &next_entry)) {
		it->item.key = NULL;
		it->item.value = NULL;
		return FX_ERR_NO_DATA;
	}

	struct fx_dict_bucket_item *next_item
		= fx_unbox(struct fx_dict_bucket_item, next_entry, bi_entry);

	if (!next_item) {
		it->item.key = NULL;
		it->item.value = NULL;
		return FX_ERR_NO_DATA;
	}

	it->i++;
	it->item.key = next_item->bi_str;
	it->item.value = next_item->bi_value;

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

	return FX_SUCCESS;
}

static enum fx_status iterator_erase(fx_iterator *obj)
{
	struct fx_dict_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_DICT_ITERATOR);

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

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

	struct fx_bst_node *next_node;
	struct fx_queue_entry *next_entry;
	if (!get_next_node(it->_cbn, it->_cqe, &next_node, &next_entry)) {
		it->item.key = NULL;
		it->item.value = NULL;
		return FX_ERR_NO_DATA;
	}

	struct fx_dict_bucket *cur_bucket
		= fx_unbox(struct fx_dict_bucket, it->_cbn, bk_node);
	struct fx_dict_bucket_item *cur_item
		= fx_unbox(struct fx_dict_bucket_item, it->_cqe, bi_entry);

	struct fx_dict_bucket_item *next_item
		= fx_unbox(struct fx_dict_bucket_item, next_entry, bi_entry);

	fx_status status = delete_item(it->_d_p, cur_bucket, cur_item);
	if (FX_ERR(status)) {
		return status;
	}

	if (next_item) {
		it->item.key = next_item->bi_str;
		it->item.value = next_item->bi_value;

		it->_cbn = next_node;
		it->_cqe = next_entry;
	} else {
		it->item.key = NULL;
		it->item.value = NULL;

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

	return FX_SUCCESS;
}

static fx_iterator_value iterator_get_value(fx_iterator *obj)
{
	struct fx_dict_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_DICT_ITERATOR);

	return FX_ITERATOR_VALUE_PTR(&it->item);
}

static const fx_iterator_value iterator_get_cvalue(const fx_iterator *obj)
{
	struct fx_dict_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_DICT_ITERATOR);

	return FX_ITERATOR_VALUE_CPTR(&it->item);
}

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

// ---- fx_dict DEFINITION
FX_TYPE_CLASS_DEFINITION_BEGIN(fx_dict)
	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = dict_to_string;
	FX_TYPE_CLASS_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_iterable, FX_TYPE_ITERABLE)
		FX_INTERFACE_ENTRY(it_begin) = iterable_begin;
		FX_INTERFACE_ENTRY(it_cbegin) = iterable_cbegin;
	FX_TYPE_CLASS_INTERFACE_END(fx_iterable, FX_TYPE_ITERABLE)
FX_TYPE_CLASS_DEFINITION_END(fx_dict)

FX_TYPE_DEFINITION_BEGIN(fx_dict)
	FX_TYPE_ID(0xd2af61d9, 0xd0be, 0x4960, 0xbe3f, 0x509749814c10);
	FX_TYPE_CLASS(fx_dict_class);
	FX_TYPE_IMPLEMENTS(FX_TYPE_ITERABLE);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_dict_p);
	FX_TYPE_INSTANCE_INIT(dict_init);
	FX_TYPE_INSTANCE_FINI(dict_fini);
FX_TYPE_DEFINITION_END(fx_dict)

// ---- fx_dict_iterator DEFINITION
FX_TYPE_CLASS_DEFINITION_BEGIN(fx_dict_iterator)
	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = NULL;
	FX_TYPE_CLASS_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_CLASS_INTERFACE_BEGIN(fx_iterator, FX_TYPE_ITERATOR)
		FX_INTERFACE_ENTRY(it_move_next) = iterator_move_next;
		FX_INTERFACE_ENTRY(it_erase) = iterator_erase;
		FX_INTERFACE_ENTRY(it_get_value) = iterator_get_value;
		FX_INTERFACE_ENTRY(it_get_cvalue) = iterator_get_cvalue;
	FX_TYPE_CLASS_INTERFACE_END(fx_iterator, FX_TYPE_ITERATOR)
FX_TYPE_CLASS_DEFINITION_END(fx_dict_iterator)

FX_TYPE_DEFINITION_BEGIN(fx_dict_iterator)
	FX_TYPE_ID(0x9ea96701, 0x1713, 0x4a3e, 0xbf63, 0xdc856b456f3b);
	FX_TYPE_EXTENDS(FX_TYPE_ITERATOR);
	FX_TYPE_CLASS(fx_dict_iterator_class);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_dict_iterator_p);
FX_TYPE_DEFINITION_END(fx_dict_iterator)