// vim:fileencoding=utf-8:foldmethod=marker

#include "table.h"
#include "../hasher/hasher.h"
#include "../../mem/allocator/mem_allocator.h"
#include "../../stream/stream.h"
#include "../../../common/aliases/aliases.h"
#include "../../../common/assert/assert.h"
#include "../../../common/misc/misc_utils.h"

wp_persist inline void hashTableValidate(const WpHashTable *table, u64 key_size, u64 value_size);

void *hashTableKeys(WpHashTable *table) {
	hashTableValidate(table, 0, 0);
	return table->keys;
}

void *hashTableValues(WpHashTable *table) {
	hashTableValidate(table, 0, 0);
	return table->values;
}

void *hashTableGetKeyFromBucket(WpHashTable *table, WpHashBucket bucket, u64 key_size) {
	hashTableValidate(table, key_size, 0);
	return arrayGet(table->keys, bucket.item_index, key_size);
}

void *hashTableGetValueFromBucket(WpHashTable *table, WpHashBucket bucket, u64 value_size) {
	hashTableValidate(table, 0, value_size);
	return table->has_values ? arrayGet(table->values, bucket.item_index, value_size) : NULL;
}

void *hashTableGetKeyFromBucketRef(WpHashTable *table, WpHashBucketRef ref, u64 key_size) {
	return hashTableGetKeyFromBucket(table, table->buckets[ref.bucket_index], key_size);
}

void *hashTableGetValueFromBucketRef(WpHashTable *table, WpHashBucketRef ref, u64 value_size) {
	return hashTableGetValueFromBucket(table, table->buckets[ref.bucket_index], value_size);
}

u64 hashTableCapacity(WpHashTable *table) {
	hashTableValidate(table, 0, 0);
	return wpArrayCapacity(table->keys);
}

u64 hashTableCount(WpHashTable *table) {
	hashTableValidate(table, 0, 0);
	return wpArrayCount(table->keys);
}

u64 hashTableKeySize(WpHashTable *table) {
	hashTableValidate(table, 0, 0);
	return wpArrayItemSize(table->keys);
}

u64 hashTableValueSize(WpHashTable *table) {
	hashTableValidate(table, 0, 0);
	return table->has_values ? wpArrayItemSize(table->values) : 0;
}

u64 hashTableCalcHash(WpHashTable *table, void *key, u64 key_size) {
	hashTableValidate(table, key_size, 0);
	wpDebugAssert(key != NULL, "`key` should not be NULL");
	WpU8Stream stream = table->encoder(key);
	return wpHasherGetHash64(&table->hasher, &stream);
}

void hashTableGrow(WpHashTable *table, const WpAllocator *allocator, u64 key_size, u64 value_size) {
	hashTableValidate(table, key_size, value_size);
	wpDebugAssert(allocator != NULL, "`allocator` should not be NULL");

	u64 capacity              = hashTableCapacity(table) * 2; 
	u64 mask                  = capacity - wpU64Const(1);
	WpArray keys              = arrayAllocCapacity(allocator, capacity, WP_ARRAY_INIT_NONE, key_size);
	WpArray values            = NULL;
	WpHashBucketArray buckets = wpArrayAllocCapacity(WpHashBucket, allocator, capacity, WP_ARRAY_INIT_FILLED);

	if (table->has_values) {
		values = arrayAllocCapacity(allocator, capacity, WP_ARRAY_INIT_NONE, value_size);
	}

	WpArray old_keys   = table->keys;
	WpArray old_values = table->values;

	table->keys        = keys;
	table->values      = values;
	table->buckets     = buckets;
	table->bucket_mask = mask;

	for (u64 i = 0; i < wpArrayCount(old_keys); ++i) {
		void *key           = arrayGet(old_keys, i, key_size);
		void *value         = table->has_values ? arrayGet(old_values, i, value_size) : NULL;
		u64 hash            = hashTableCalcHash(table, key, key_size);
		WpHashBucketRef ref = hashTableBucketRef(table, hash);
		hashTableInsert(table, key, value, ref, key_size, value_size);
	}
}

WpHashLookupResult hashTableLookup(WpHashTable *table, void *key, WpHashBucketRef ref, u64 key_size) {
	hashTableValidate(table, key_size, 0);
	wpDebugAssert(key != NULL, "`key` should not be NULL");

	while (true) {
		WpHashBucket *bucket = &table->buckets[ref.bucket_index];

		if (ref.dist_and_tag == bucket->dist_and_tag) {
			b8 key_equal = table->key_eq(key, arrayGet(table->keys, bucket->item_index, key_size));
			return (WpHashLookupResult){
				.ref.dist_and_tag = ref.dist_and_tag,
				.ref.bucket_index = ref.bucket_index,
				.found            = key_equal,
			};
		} else if (hashTableGetDist(ref.dist_and_tag) > hashTableGetDist(bucket->dist_and_tag)) {
			return (WpHashLookupResult){
				.ref.dist_and_tag = ref.dist_and_tag,
				.ref.bucket_index = ref.bucket_index,
				.found            = false,
			};
		}

		ref.dist_and_tag = hashTableDistInc(ref.dist_and_tag);
		ref.bucket_index = hashTableNextBucket(table, ref.bucket_index);
	}
}

void hashTableUpdate(WpHashTable *table, void *value, WpHashBucketRef ref, u64 value_size) {
	hashTableValidate(table, 0, value_size);
	if (!(table->has_values)) {
		wpRuntimeAssert(0, "Attempting to update a hash table that has no values");
	}
	wpDebugAssert(value != NULL, "`value` should not be NULL");

	WpHashBucket *bucket = &table->buckets[ref.bucket_index];
	arraySet(table->values, bucket->item_index, value, value_size);
}

void hashTableInsert(WpHashTable *table, void *key, void *value, WpHashBucketRef ref, u64 key_size,
                     u64 value_size) {
	hashTableValidate(table, key_size, value_size);
	wpDebugAssert(key != NULL, "`key` should not be NULL");

	u64 index = wpArrayCount(table->keys);
	while (true) {
		WpHashBucket *bucket = &table->buckets[ref.bucket_index];

		if (bucket->dist_and_tag == 0) {
			bucket->dist_and_tag = ref.dist_and_tag;
			bucket->item_index   = index;
			break;
		} else if (hashTableGetDist(ref.dist_and_tag) > hashTableGetDist(bucket->dist_and_tag)) {
			wpMiscUtilsSwap(u64, &ref.dist_and_tag, &bucket->dist_and_tag);
			wpMiscUtilsSwap(u64, &index, &bucket->item_index);
		}

		ref.dist_and_tag = hashTableDistInc(ref.dist_and_tag);
		ref.bucket_index = hashTableNextBucket(table, ref.bucket_index);
	}

	arrayAppendCapped(table->keys, key, key_size);
	if (table->has_values) {
		arrayAppendCapped(table->values, value, value_size);
	}
	table->load_factor = (f32)wpArrayCount(table->keys) / (f32)wpArrayCapacity(table->keys);
}

void hashTableRemove(WpHashTable *table, WpHashBucketRef ref, u64 key_size, u64 value_size) {
	hashTableValidate(table, key_size, value_size);

	WpHashBucket *bucket = &table->buckets[ref.bucket_index];

	u64 last_item_index = wpArrayCount(table->keys) - 1;

	if (bucket->item_index < last_item_index) {
		// Update the bucket for the last item before swapping
		u64 hash = hashTableCalcHash(table, arrayGet(table->keys, last_item_index, key_size), key_size);

		WpHashLookupResult result = hashTableLookup(table, arrayGet(table->keys, last_item_index, key_size),
		                                            hashTableBucketRef(table, hash), key_size);
		wpRuntimeAssert(result.found, "Hash table item should exist but it doesn't");

		table->buckets[result.ref.bucket_index].item_index = bucket->item_index;

		arrayRemoveSwapEnd(table->keys, bucket->item_index, key_size);
		if (table->has_values) {
			arrayRemoveSwapEnd(table->values, bucket->item_index, value_size);
		}
	}

	while (bucket->dist_and_tag > 0) {
		u64 next_index     = hashTableNextBucket(table, ref.bucket_index);
		WpHashBucket *next = &table->buckets[next_index];

		if (hashTableGetDist(next->dist_and_tag) > WP_HASH_DIST_INC) {
			bucket->dist_and_tag = hashTableDistDec(next->dist_and_tag);
			bucket->item_index   = next->item_index;

			ref.bucket_index = next_index;
			bucket           = &table->buckets[ref.bucket_index];
		} else {
			bucket->dist_and_tag = 0;
			bucket->item_index   = 0;
		}
	}
}

void hashTableClear(WpHashTable *table, u64 key_size, u64 value_size) {
	hashTableValidate(table, key_size, value_size);
	arrayClear(table->keys, key_size);
	if (table->has_values) {
		arrayClear(table->values, value_size);
	}
	wpArrayZero(WpHashBucket, table->buckets);
}

WpHashBucketRef hashTableBucketRef(WpHashTable *table, u64 hash) {
	hashTableValidate(table, 0, 0);
	return (WpHashBucketRef){
		.dist_and_tag = WP_HASH_DIST_INC | (hash & WP_HASH_TAG_MASK),
		.bucket_index = hash & table->bucket_mask,
	};
}

u64 hashTableGetDist(u64 dist_and_tag) {
	return dist_and_tag & (~WP_HASH_TAG_MASK);
}

u64 hashTableNextBucket(WpHashTable *table, u64 bucket_index) {
	return (++bucket_index) & table->bucket_mask;
}

u64 hashTableDistInc(u64 dist_and_tag) {
	return dist_and_tag + WP_HASH_DIST_INC;
}

u64 hashTableDistDec(u64 dist_and_tag) {
	return dist_and_tag - WP_HASH_DIST_INC;
}

wp_persist inline void hashTableValidate(const WpHashTable *table, u64 key_size, u64 value_size) {
	wpDebugAssert(table != NULL, "`table` should not be NULL");
	wpRuntimeAssert(WP_HASH_TABLE_MAGIC == table->magic, "`table` is not a valid wapp hash table");
	if (key_size > 0) {
		wpRuntimeAssert(key_size == wpArrayItemSize(table->keys), "Invalid key type provided");
	}
	if (value_size > 0) {
		wpRuntimeAssert(value_size == wpArrayItemSize(table->values), "Invalid value type provided");
	}
}