#ifndef MAP_H
#define MAP_H

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

#define NGC_BUCKET_IDX(hash, pow2, bucket) \
	((hash) & (((pow2) << (bucket)) - 1))

#ifndef NGC_CONTAINER_OF
#define NGC_CONTAINER_OF(ptr, type, member) \
	(type *)((char *)(ptr) - offsetof(type, member))
#endif

/* trait would be something like is_hashable + is_comparable, but that's not yet
 * supported so just go with this for now */
typedef map[any K, any D]() {
	struct <>_tuple {
		<K> key;
		<D> data;
	};

	typedef struct <>_tuple *<>_iter;

	struct <>_node {
		size_t hash;
		struct <>_bucket *bucket;
		struct <>_tuple t;
	};

	struct <>_bucket {
		size_t size;
		struct <>_bucket *next;
		struct <>_node nodes[];
	};

	typedef struct <>_root {
		size_t len;
		size_t count;
		size_t pow2;
		struct <>_bucket **buckets;
	} <>;

	<> <>_create(size_t count)
	{
		size_t pow2 = 1;
		while (pow2 < count)
			pow2 <<= 1;

		return (<>) {
			.len = 0,
			.count = 0,
			.pow2 = pow2,
			.buckets = NULL
		};
	}

	void <>_destroy(<> *root)
	{
		for (size_t i = 0; i < root->count; ++i)
			free(root->buckets[i]);

		free(root->buckets);
	}

	<D> *<>_insert(<>* root, <K> key, <D> data)
	{
		size_t hash = <K>_hash(&key);
		/* look through buckets in order */
		for (size_t b = 0; b < root->count; ++b) {
			struct <>_bucket *bucket = root->buckets[b];
			size_t idx = NGC_BUCKET_IDX(hash, root->pow2, b);

			struct <>_node *node = &bucket->nodes[idx];
			/* free to use this slot */
			if (!node->bucket) {
				node->bucket = bucket;
				node->hash = hash;
				node->t.data = data;
				node->t.key = key;
				root->len++;
				return &node->t.data;
			}

			/* there already exists a node like this */
			if (node->hash == hash && <K>_cmp(&node->t.key, &key) == 0)
				return &node->t.data;
		}

		/* no bucket available, create new one */
		size_t size = root->pow2 << root->count;
		size_t bytes = sizeof(struct <>_bucket) + sizeof(struct <>_node) * size;
		struct <>_bucket *bucket = calloc(1, bytes);
		if (!bucket)
			return NULL;

		bucket->size = size;

		size_t buckets_bytes = sizeof(struct <>_bucket *) * (root->count + 1);
		struct <>_bucket **new_buckets = realloc(root->buckets, buckets_bytes);
		if (!new_buckets) {
			free(bucket);
			return NULL;
		}

		root->buckets = new_buckets;
		if (root->count != 0)
			root->buckets[root->count - 1]->next = bucket;

		root->buckets[root->count] = bucket;

		/* populate node */
		size_t idx = NGC_BUCKET_IDX(hash, root->pow2, root->count);
		struct <>_node *node = &bucket->nodes[idx];
		node->bucket = bucket;
		node->hash = hash;
		node->t.key = key;
		node->t.data = data;
		root->count++;
		root->len++;
		return &node->t.data;
	}

	<D> *<>_find(<> *root, <K> key)
	{
		if (root->len == 0)
			return NULL;

		size_t hash = <K>_hash(&key);
		for (size_t b = 0; b < root->count; ++b) {
			struct <>_bucket *bucket = root->buckets[b];
			size_t idx = NGC_BUCKET_IDX(hash, root->pow2, b);

			struct <>_node *node = &bucket->nodes[idx];
			if (node->hash != hash)
				continue;

			if (<K>_cmp(&node->t.key, &key) != 0)
				continue;

			return &node->t.data;
		}

		return NULL;
	}

	void <>_remove_found(<> *root, <D> *data)
	{
		struct <>_tuple *tuple = NGC_CONTAINER_OF(data, struct <>_tuple, data);
		struct <>_node *node = NGC_CONTAINER_OF(tuple, struct <>_node, t);
		node->bucket = NULL;
		root->len--;
	}

	void <>_remove(<> *root, <K> key)
	{
		<D> *found = <>_find(root, key);
		if (!found)
			return;

		<>_remove_found(root, found);
	}

	struct <>_tuple *<>_find_next(struct <>_bucket *bucket, struct <>_tuple *tuple)
	{
		struct <>_node *node = NGC_CONTAINER_OF(tuple, struct <>_node, t);
		size_t idx = node - bucket->nodes;
		for (; idx < bucket->size; ++idx) {
			struct <>_node *candidate = &bucket->nodes[idx];
			if (candidate->bucket)
				return &candidate->t;
		}

		struct <>_bucket *next = bucket->next;
		if (!next)
			return NULL;

		return <>_find_next(next, &next->nodes[0].t);
	}

	struct <>_tuple *<>_begin(<> *root)
	{
		if (root->len == 0)
			return NULL;

		struct <>_bucket *bucket = root->buckets[0];
		return <>_find_next(bucket, &bucket->nodes[0].t);
	}

	struct <>_tuple *<>_next(struct <>_tuple *t)
	{
		struct <>_node *node = NGC_CONTAINER_OF(t, struct <>_node, t);
		return <>_find_next(node->bucket, &(node + 1)->t);
	}

	bool <>_end(<> *root, struct <>_tuple *t)
	{
		(void)root;
		return t == NULL;
	}

	size_t <>_len(<> *root)
	{
		return root->len;
	}
};

#endif /* MAP_H */