#include <assert.h>
#include <fx/collections/hashtable.h>
#include <fx/misc.h>
#include <fx/reflection/property.h>
#include <fx/status.h>
#include <fx/string.h>
#include <fx/vector.h>
#include <stdbool.h>
#include <stdlib.h>

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

/* these primes are used as array capacities to avoid hash collisions. */
static const size_t primes[] = {
	11,
	23,
	47,
	97,
	197,
	397,
	797,
	1597,
	3203,
	6421,
	12853,
	25717,
	51437,
};

struct table_item {
	fx_hashtable_item *i_item;
};

struct fx_hashtable_item_p {
	fx_hashtable_item *i_self;
	fx_value i_key, i_value;
};

struct fx_hashtable_p {
	struct table_item *t_items;
	/* number of items currently stored in t_items */
	size_t t_count;
	/* index in primes[] that represents the current maximum size of
	 * t_items */
	size_t t_max_index;
};

struct fx_hashtable_iterator_p {
	size_t i;
	fx_hashtable_item *item;
	fx_value item_value;

	const fx_hashtable *_h;
	struct fx_hashtable_p *_h_p;
};

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

static fx_status convert_key_to_index(
	const fx_value *v,
	struct table_item *table,
	size_t table_size,
	size_t *out_index)
{
	uint64_t hash = 0;
	fx_status status = fx_value_hash(v, &hash);
	if (!FX_OK(status)) {
		return status;
	}

	size_t index = hash % table_size, i = 1;
	bool index_found = false;
	while (table[index].i_item && i <= table_size) {
		const struct fx_hashtable_item_p *item = fx_object_get_private(
			table[index].i_item,
			FX_TYPE_HASHTABLE_ITEM);
		int cmp = fx_value_compare(&item->i_key, v);
		if (cmp == 0) {
			index_found = true;
			break;
		}

		index = (index + (i * i)) % table_size;
		i++;
	}

	if (!index_found) {
		return FX_ERR_NO_ENTRY;
	}

	assert(index < table_size);
	*out_index = index;
	return FX_SUCCESS;
}

static fx_status allocate_index_for_key(
	const fx_value *v,
	struct table_item *table,
	size_t table_size,
	size_t *out_index)
{
	uint64_t hash = 0;
	fx_status status = fx_value_hash(v, &hash);
	if (!FX_OK(status)) {
		return status;
	}

	size_t index = hash % table_size;
	bool ok = false;

	for (size_t i = 1; i < table_size; i++) {
		if (!table[index].i_item) {
			ok = true;
			break;
		}

		const struct fx_hashtable_item_p *item = fx_object_get_private(
			table[index].i_item,
			FX_TYPE_HASHTABLE_ITEM);
		int cmp = fx_value_compare(&item->i_key, v);
		if (cmp == 0) {
			ok = true;
			break;
		}

		index = (index + (i * i)) % table_size;
		i++;
	}

	if (!ok) {
		return FX_ERR_NO_SPACE;
	}

	assert(index < table_size);
	*out_index = index;
	return FX_SUCCESS;
}

static fx_status add_item_to_table(
	fx_hashtable_item *item,
	struct table_item *table,
	size_t table_capacity)
{
	size_t index = 0;
	struct fx_hashtable_item_p *item_p
		= fx_object_get_private(item, FX_TYPE_HASHTABLE_ITEM);
	fx_status status = allocate_index_for_key(
		&item_p->i_key,
		table,
		table_capacity,
		&index);
	if (!FX_OK(status)) {
		return status;
	}

	if (index >= table_capacity) {
		return FX_ERR_NO_SPACE;
	}

	struct table_item *bucket = &table[index];
	if (bucket->i_item) {
		fx_hashtable_item_unref(bucket->i_item);
	}

	bucket->i_item = item;

	return FX_SUCCESS;
}

static fx_status destroy_table(
	struct table_item *table,
	size_t length,
	bool destroy_items)
{
	if (!destroy_items) {
		free(table);
		return FX_SUCCESS;
	}

	for (size_t i = 0; i < length; i++) {
		struct table_item *item = &table[i];
		if (!item->i_item) {
			continue;
		}

		fx_hashtable_item_unref(item->i_item);
	}

	free(table);
	return FX_SUCCESS;
}

static fx_status migrate_table(
	struct table_item *src_table,
	size_t src_length,
	struct table_item *dest_table,
	size_t dest_capacity)
{
	for (size_t i = 0; i < src_length; i++) {
		struct table_item *item = &src_table[i];
		if (!item->i_item) {
			continue;
		}

		add_item_to_table(item->i_item, dest_table, dest_capacity);
	}

	return FX_SUCCESS;
}

static fx_status resize_table(
	struct fx_hashtable_p *hashtable,
	size_t new_capacity_index)
{
	size_t capacity = primes[new_capacity_index];
	struct table_item *new_table = calloc(capacity, sizeof *new_table);
	if (!new_table) {
		return FX_ERR_NO_MEMORY;
	}

	fx_status status = migrate_table(
		hashtable->t_items,
		primes[hashtable->t_max_index],
		new_table,
		capacity);
	if (!FX_OK(status)) {
		free(new_table);
		return status;
	}

	free(hashtable->t_items);
	hashtable->t_items = new_table;
	hashtable->t_max_index = new_capacity_index;
	return FX_SUCCESS;
}

static fx_status hashtable_put(
	struct fx_hashtable_p *hashtable,
	fx_value *key,
	fx_value *value)
{
	fx_status status = FX_SUCCESS;
	size_t capacity = primes[hashtable->t_max_index];
	if (hashtable->t_count + 1 > capacity) {
		status = resize_table(hashtable, hashtable->t_max_index + 1);
		if (!FX_OK(status)) {
			return status;
		}

		capacity = primes[hashtable->t_max_index];
	}

	size_t index = 0;
	do {
		status = allocate_index_for_key(
			key,
			hashtable->t_items,
			capacity,
			&index);
		if (status == FX_ERR_NO_SPACE) {
			status = resize_table(
				hashtable,
				hashtable->t_max_index + 1);
			if (!FX_OK(status)) {
				return status;
			}

			capacity = primes[hashtable->t_max_index];
			continue;
		}

		if (!FX_OK(status)) {
			return status;
		}

		break;
	} while (1);

	fx_hashtable_item *item = fx_object_create(FX_TYPE_HASHTABLE_ITEM);
	if (!item) {
		return FX_ERR_NO_MEMORY;
	}

	struct fx_hashtable_item_p *item_p
		= fx_object_get_private(item, FX_TYPE_HASHTABLE_ITEM);
	fx_value_copy(&item_p->i_key, key);
	fx_value_copy(&item_p->i_value, value);

	hashtable->t_items[index].i_item = item;
	hashtable->t_count++;
	return FX_SUCCESS;
}

static const fx_value *hashtable_get(
	const struct fx_hashtable_p *hashtable,
	const fx_value *key)
{
	size_t index = 0;
	size_t capacity = primes[hashtable->t_max_index];
	fx_status status = convert_key_to_index(
		key,
		hashtable->t_items,
		capacity,
		&index);
	if (!FX_OK(status)) {
		return NULL;
	}

	const fx_hashtable_item *item = hashtable->t_items[index].i_item;
	if (!item) {
		return NULL;
	}

	const struct fx_hashtable_item_p *item_p
		= fx_object_get_private(item, FX_TYPE_HASHTABLE_ITEM);
	return &item_p->i_value;
}

static size_t hashtable_get_count(const struct fx_hashtable_p *hashtable)
{
	return hashtable->t_count;
}

static bool hashtable_is_empty(const struct fx_hashtable_p *hashtable)
{
	return hashtable->t_count == 0;
}

static const fx_value *hashtable_item_get_key(
	const struct fx_hashtable_item_p *item)
{
	return &item->i_key;
}

static const fx_value *hashtable_item_get_value(
	const struct fx_hashtable_item_p *item)
{
	return &item->i_value;
}

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

fx_status fx_hashtable_put(
	fx_hashtable *hashtable,
	fx_value *key,
	fx_value *value)
{
	FX_CLASS_DISPATCH_STATIC(
		FX_TYPE_HASHTABLE,
		hashtable_put,
		hashtable,
		key,
		value);
}

const fx_value *fx_hashtable_get(
	const fx_hashtable *hashtable,
	const fx_value *key)
{
	FX_CLASS_DISPATCH_STATIC(
		FX_TYPE_HASHTABLE,
		hashtable_get,
		hashtable,
		key);
}

size_t fx_hashtable_get_count(const fx_hashtable *hashtable)
{
	FX_CLASS_DISPATCH_STATIC_0(
		FX_TYPE_HASHTABLE,
		hashtable_get_count,
		hashtable);
}

bool fx_hashtable_is_empty(const fx_hashtable *hashtable)
{
	FX_CLASS_DISPATCH_STATIC_0(
		FX_TYPE_HASHTABLE,
		hashtable_is_empty,
		hashtable);
}

const fx_value *fx_hashtable_item_get_key(const fx_hashtable_item *item)
{
	FX_CLASS_DISPATCH_STATIC_0(
		FX_TYPE_HASHTABLE_ITEM,
		hashtable_item_get_key,
		item);
}

const fx_value *fx_hashtable_item_get_value(const fx_hashtable_item *item)
{
	FX_CLASS_DISPATCH_STATIC_0(
		FX_TYPE_HASHTABLE_ITEM,
		hashtable_item_get_value,
		item);
}

const fx_iterator *fx_hashtable_begin(const fx_hashtable *hashtable)
{
	fx_hashtable_iterator *it_obj
		= fx_object_create(FX_TYPE_HASHTABLE_ITERATOR);
	struct fx_hashtable_iterator_p *it
		= fx_object_get_private(it_obj, FX_TYPE_HASHTABLE_ITERATOR);

	it->_h = hashtable;
	it->_h_p = fx_object_get_private(hashtable, FX_TYPE_HASHTABLE);

	it->i = 0;
	if (fx_hashtable_is_empty(hashtable)) {
		fx_iterator_set_status(it_obj, FX_ERR_NO_DATA);
		return it_obj;
	}

	struct table_item *table = it->_h_p->t_items;
	size_t capacity = primes[it->_h_p->t_max_index];

	while (it->i < capacity && table[it->i].i_item == NULL) {
		it->i++;
	}

	it->item_value = FX_VALUE_OBJECT_REF(table[it->i].i_item);
	return it_obj;
}

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

static void hashtable_init(fx_object *obj, void *priv)
{
	struct fx_hashtable_p *map = priv;
	map->t_max_index = 0;
	map->t_count = 0;

	size_t capacity = primes[map->t_max_index];
	map->t_items = calloc(capacity, sizeof *map->t_items);
}

static void hashtable_fini(fx_object *obj, void *priv)
{
	struct fx_hashtable_p *map = priv;

	size_t capacity = primes[map->t_max_index];
	size_t nr_destroyed = 0;
	for (size_t i = 0; i < capacity; i++) {
		if (map->t_items[i].i_item) {
			nr_destroyed++;
			fx_hashtable_item_unref(map->t_items[i].i_item);
		}
	}

	free(map->t_items);
}

static void hashtable_item_fini(fx_object *obj, void *priv)
{
	struct fx_hashtable_item_p *item = priv;
	fx_value_unset(&item->i_key);
	fx_value_unset(&item->i_value);
}

static void hashtable_iterator_fini(fx_object *obj, void *priv)
{
	struct fx_hashtable_iterator_p *it = priv;
	fx_value_unset(&it->item_value);
}

static fx_status to_string(
	const fx_value *obj,
	fx_stream *out,
	const char *format)
{
	struct fx_hashtable_p *hashtable
		= fx_object_get_private(obj->v_object, FX_TYPE_HASHTABLE);

	if (!hashtable->t_count) {
		fx_stream_write_cstr(out, "{}", NULL);
		return FX_SUCCESS;
	}

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

	fx_stream_push_indent(out, 1);

	size_t nr_written = 0;
	size_t capacity = primes[hashtable->t_max_index];
	for (size_t i = 0; i < capacity; i++) {
		fx_hashtable_item *item = hashtable->t_items[i].i_item;
		if (!item) {
			continue;
		}

		nr_written++;
		struct fx_hashtable_item_p *item_p
			= fx_object_get_private(item, FX_TYPE_HASHTABLE_ITEM);

		const char *cstr = NULL;
		fx_value_get_cstr(&item_p->i_key, &cstr);

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

		fx_value_to_string(&item_p->i_key, out, NULL);

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

		fx_stream_write_cstr(out, " = ", NULL);

		fx_value_get_cstr(&item_p->i_value, &cstr);
		if (cstr) {
			fx_stream_write_char(out, '"');
		}

		fx_value_to_string(&item_p->i_value, out, NULL);

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

		if (nr_written != hashtable->t_count) {
			fx_stream_write_char(out, ',');
		}

		fx_stream_write_char(out, '\n');
	}

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

static fx_status get_key(
	const fx_value *item_value,
	const fx_property *prop,
	fx_value *out)
{
	fx_hashtable_item *item = NULL;
	fx_value_get_object(item_value, &item);
	if (!item) {
		return FX_ERR_NOT_SUPPORTED;
	}

	struct fx_hashtable_item_p *item_p
		= fx_object_get_private(item, FX_TYPE_HASHTABLE_ITEM);
	fx_value_copy(out, &item_p->i_key);

	return FX_SUCCESS;
}

static fx_status get_value(
	const fx_value *item_value,
	const fx_property *prop,
	fx_value *out)
{
	fx_hashtable_item *item = NULL;
	fx_value_get_object(item_value, &item);
	if (!item) {
		return FX_ERR_NOT_SUPPORTED;
	}

	struct fx_hashtable_item_p *item_p
		= fx_object_get_private(item, FX_TYPE_HASHTABLE_ITEM);
	fx_value_copy(out, &item_p->i_value);

	return FX_SUCCESS;
}

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

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

	fx_value_unset(&it->item_value);
	struct table_item *table = it->_h_p->t_items;
	size_t capacity = primes[it->_h_p->t_max_index];

	it->i++;
	while (it->i < capacity && table[it->i].i_item == NULL) {
		it->i++;
	}

	if (it->i >= capacity || !table[it->i].i_item) {
		return FX_ERR_NO_DATA;
	}

	it->item_value = FX_VALUE_OBJECT_REF(table[it->i].i_item);
	return FX_SUCCESS;
}

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

	struct table_item *table = it->_h_p->t_items;
	size_t capacity = primes[it->_h_p->t_max_index];

	if (table[it->i].i_item) {
		fx_hashtable_item_unref(table[it->i].i_item);
		table[it->i].i_item = NULL;
	}

	it->i++;
	while (it->i < capacity && table[it->i].i_item == NULL) {
		it->i++;
	}

	if (it->i >= capacity || !table[it->i].i_item) {
		return FX_ERR_NO_DATA;
	}

	return FX_SUCCESS;
}

static const fx_value *iterator_get_value(const fx_iterator *obj)
{
	struct fx_hashtable_iterator_p *it
		= fx_object_get_private(obj, FX_TYPE_HASHTABLE_ITERATOR);
	struct table_item *table = it->_h_p->t_items;
	size_t capacity = primes[it->_h_p->t_max_index];

	if (it->i >= capacity) {
		return NULL;
	}

	if (table[it->i].i_item) {
		return &it->item_value;
	}

	return NULL;
}

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

// ---- fx_hashtable DEFINITION
FX_TYPE_CLASS_BEGIN(fx_hashtable)
	FX_TYPE_VTABLE_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = to_string;
	FX_TYPE_VTABLE_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_VTABLE_INTERFACE_BEGIN(fx_iterable, FX_TYPE_ITERABLE)
		FX_INTERFACE_ENTRY(it_begin) = fx_hashtable_begin;
	FX_TYPE_VTABLE_INTERFACE_END(fx_iterable, FX_TYPE_ITERABLE)
FX_TYPE_CLASS_END(fx_hashtable)

FX_TYPE_DEFINITION_BEGIN(fx_hashtable)
	FX_TYPE_ID(0xc73b431c, 0x5bd5, 0x45ac, 0x888b, 0x35d352db1fe4);
	FX_TYPE_NAME("fx.collections.hashtable");
	FX_TYPE_CLASS(fx_hashtable_class);
	FX_TYPE_IMPLEMENTS(FX_TYPE_ITERABLE);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_hashtable_p);
	FX_TYPE_INSTANCE_INIT(hashtable_init);
	FX_TYPE_INSTANCE_FINI(hashtable_fini);
FX_TYPE_DEFINITION_END(fx_hashtable)

// ---- fx_hashtable_iterator DEFINITION
FX_TYPE_CLASS_BEGIN(fx_hashtable_iterator)
	FX_TYPE_VTABLE_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = NULL;
	FX_TYPE_VTABLE_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_VTABLE_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_TYPE_VTABLE_INTERFACE_END(fx_iterator, FX_TYPE_ITERATOR)
FX_TYPE_CLASS_END(fx_hashtable_iterator)

FX_TYPE_DEFINITION_BEGIN(fx_hashtable_iterator)
	FX_TYPE_ID(0x42c6836d, 0xf801, 0x49f9, 0xba58, 0xf9cd6f1cb22d);
	FX_TYPE_NAME("fx.collections.hashtable.iterator");
	FX_TYPE_EXTENDS(FX_TYPE_ITERATOR);
	FX_TYPE_CLASS(fx_hashtable_iterator_class);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_hashtable_iterator_p);
	FX_TYPE_INSTANCE_FINI(hashtable_iterator_fini);
FX_TYPE_DEFINITION_END(fx_hashtable_iterator)

// ---- fx_hashtable_item DEFINITION
FX_TYPE_CLASS_BEGIN(fx_hashtable_item)
	FX_TYPE_VTABLE_INTERFACE_BEGIN(fx_object, FX_TYPE_OBJECT)
		FX_INTERFACE_ENTRY(to_string) = NULL;
	FX_TYPE_VTABLE_INTERFACE_END(fx_object, FX_TYPE_OBJECT)

	FX_TYPE_PROPERTY("key", get_key, NULL);
	FX_TYPE_PROPERTY("value", get_value, NULL);
FX_TYPE_CLASS_END(fx_hashtable_item)

FX_TYPE_DEFINITION_BEGIN(fx_hashtable_item)
	FX_TYPE_ID(0x748f78f3, 0x53df, 0x406c, 0x8777, 0x2075bd11fb97);
	FX_TYPE_NAME("fx.collections.hashtable.item");
	FX_TYPE_EXTENDS(FX_TYPE_OBJECT);
	FX_TYPE_CLASS(fx_hashtable_item_class);
	FX_TYPE_INSTANCE_PRIVATE(struct fx_hashtable_item_p);
	FX_TYPE_INSTANCE_FINI(hashtable_item_fini);
FX_TYPE_DEFINITION_END(fx_hashtable_item)