#ifndef PK_PKBKTARR_H
#define PK_PKBKTARR_H

#include "./pkmem.h" /* deleteme */
#include "./pkmacros.h" /* deleteme */
#include "pkiter.h" /*deleteme*/

#ifndef PK_BKT_ARR_ALL_UNUSED_VAL
	#define PK_BKT_ARR_ALL_UNUSED_VAL 0xFFFFFFFFFFFFFFFF
#endif
#define PK_BKT_ARR_HANDLE_B_MAX 0xFFFFFF
#define PK_BKT_ARR_HANDLE_I_MAX 64

typedef bool (pk_bkt_arr_compare_fn)(void *user_data, const void *user_obj_data, const void *arr_obj_data);
typedef void (pk_bkt_arr_iterate_fn)(void *user_data, void *arr_obj_data);

struct pk_bkt_arr_handle {
	unsigned int b : 24;
	unsigned int i : 8;
};

#if ! defined(__cplusplus)
#define pk_bkt_arr_handle_MAX ((struct pk_bkt_arr_handle){ .b = PK_BKT_ARR_HANDLE_B_MAX, .i = PK_BKT_ARR_HANDLE_I_MAX })
#else
#define pk_bkt_arr_handle_MAX (pk_bkt_arr_handle{ .b = PK_BKT_ARR_HANDLE_B_MAX, .i = PK_BKT_ARR_HANDLE_I_MAX })
constexpr struct pk_bkt_arr_handle pk_bkt_arr_handle_MAX_constexpr = pk_bkt_arr_handle_MAX;
inline constexpr bool
operator==(const pk_bkt_arr_handle &lhs, const pk_bkt_arr_handle &rhs)
{
	return lhs.b == rhs.b && lhs.i == rhs.i;
}
#endif

struct pk_bkt_arr {
	struct pk_membucket *bkt_buckets;
	struct pk_membucket *bkt_data;
	unsigned long long *idx_unused;
	void **bucketed_data;
	struct pk_bkt_arr_handle head_l;
	struct pk_bkt_arr_handle head_r;
	struct pk_bkt_arr_handle limits;
	unsigned int reserved_buckets;
	unsigned long stride;
	unsigned long alignment;
};

enum PK_BKT_ARR_HANDLE_VALIDATION : uint8_t {
	PK_BKT_ARR_HANDLE_VALIDATION_VALID                      = 0,
	PK_BKT_ARR_HANDLE_VALIDATION_BUCKET_INDEX_TOO_HIGH      = 1 << 0,
	PK_BKT_ARR_HANDLE_VALIDATION_ITEM_INDEX_TOO_HIGH        = 2 << 1,
};

enum PK_BKT_ARR_HANDLE_VALIDATION pk_bkt_arr_handle_validate(struct pk_bkt_arr *bkt_arr, struct pk_bkt_arr_handle handle);

void pk_bkt_arr_init(struct pk_bkt_arr *bkt_arr, unsigned long stride, unsigned long alignment, struct pk_bkt_arr_handle limits, struct pk_membucket *bkt_buckets, struct pk_membucket *bkt_data);
void pk_bkt_arr_clear(struct pk_bkt_arr *bkt_arr);
void pk_bkt_arr_reserve(struct pk_bkt_arr *bkt_arr, size_t count);
struct pk_bkt_arr_handle pk_bkt_arr_find_first_handle(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_compare_fn fn, void *user_data, const void *user_obj_data);
void pk_bkt_arr_iterate(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_iterate_fn fn, void *user_data);
void pk_bkt_arr_teardown(struct pk_bkt_arr *bkt_arr);
struct pk_bkt_arr_handle pk_bkt_arr_new_handle(struct pk_bkt_arr *bkt_arr);
void pk_bkt_arr_free_handle(struct pk_bkt_arr *bkt_arr, struct pk_bkt_arr_handle handle);
int pk_bkt_arr_handle_compare(struct pk_bkt_arr_handle lhs, struct pk_bkt_arr_handle rhs);
struct pk_bkt_arr_handle pk_bkt_arr_handle_increment(struct pk_bkt_arr *arr, struct pk_bkt_arr_handle h);
struct pk_bkt_arr_handle pk_bkt_arr_handle_decrement(struct pk_bkt_arr *arr, struct pk_bkt_arr_handle h);
bool pk_bkt_arr_iter_begin(struct pk_bkt_arr *arr, struct pk_iter *it);
bool pk_bkt_arr_iter_end(struct pk_bkt_arr *arr, struct pk_iter *it);
bool pk_bkt_arr_iter_increment(struct pk_bkt_arr *arr, struct pk_iter *it);
bool pk_bkt_arr_iter_decrement(struct pk_bkt_arr *arr, struct pk_iter *it);

#if defined (__cplusplus)
#include "pktmpln.h" /*deleteme*/
#include <assert.h>
template<typename T>
struct pk_bkt_arr_t : public pk_bkt_arr {
	pk_bkt_arr_t() = default;
	pk_bkt_arr_t(struct pk_bkt_arr_handle limits, struct pk_membucket *bkt_buckets, struct pk_membucket *bkt_data);
	~pk_bkt_arr_t() = default;
	T &operator[](struct pk_bkt_arr_handle);
	using FN_Iter = pk_tmpln_1<void, T*, void*>;
	using FN_Find = pk_tmpln_2<bool, const T*, const T*, const void*, const void*>;
};
template<typename T>
pk_bkt_arr_t<T>::pk_bkt_arr_t(struct pk_bkt_arr_handle limits, struct pk_membucket *bkt_buckets, struct pk_membucket *bkt_data) {
	pk_bkt_arr_init(this, sizeof(T), alignof(T), limits, bkt_buckets, bkt_data);
}
template<typename T>
T &pk_bkt_arr_t<T>::operator[](struct pk_bkt_arr_handle handle) {
	assert(this->idx_unused != nullptr);
	assert(this->bucketed_data != nullptr);
	assert(handle.b <= this->limits.b);
	assert(handle.i <= this->limits.i);
	assert(handle.b != this->head_r.b || handle.i <= this->head_r.i);
	T** two_star_programmer = reinterpret_cast<T**>(this->bucketed_data);
	return two_star_programmer[handle.b][handle.i];
}
#endif

#endif /* PK_PKBKTARR_H */
#ifdef PK_IMPL_BKTARR

#include <assert.h>
#include <string.h>

enum PK_BKT_ARR_HANDLE_VALIDATION pk_bkt_arr_handle_validate(struct pk_bkt_arr *bkt_arr, struct pk_bkt_arr_handle handle) {
	assert(bkt_arr != NULL);
	uint8_t ret = 0;
	if (handle.b >= bkt_arr->reserved_buckets || handle.b >= bkt_arr->limits.b) {
		ret |= PK_BKT_ARR_HANDLE_VALIDATION_BUCKET_INDEX_TOO_HIGH;
	}
	if (handle.i >= bkt_arr->limits.i) {
		ret |= PK_BKT_ARR_HANDLE_VALIDATION_ITEM_INDEX_TOO_HIGH;
	}
	if (handle.b == bkt_arr->head_r.b && handle.i > bkt_arr->head_r.i) {
		ret |= PK_BKT_ARR_HANDLE_VALIDATION_ITEM_INDEX_TOO_HIGH;
	}
	return (enum PK_BKT_ARR_HANDLE_VALIDATION)ret;
}

void pk_bkt_arr_init(struct pk_bkt_arr *bkt_arr, unsigned long stride, unsigned long alignment, struct pk_bkt_arr_handle limits, struct pk_membucket *bkt_buckets, struct pk_membucket *bkt_data)
{
	assert(limits.b <= PK_BKT_ARR_HANDLE_B_MAX);
	assert(limits.i <= PK_BKT_ARR_HANDLE_I_MAX);
	assert(bkt_buckets != nullptr);
	assert(bkt_data != nullptr);
	assert(bkt_arr != nullptr);
	memset(bkt_arr, 0, sizeof(struct pk_bkt_arr));
	bkt_arr->bkt_buckets = bkt_buckets;
	bkt_arr->bkt_data = bkt_data;
	bkt_arr->head_l.b = 0ul;
	bkt_arr->head_l.i = 0ul;
	bkt_arr->head_r.b = 0ul;
	bkt_arr->head_r.i = 0ul;
	bkt_arr->limits = limits;
	bkt_arr->reserved_buckets = 1;
	bkt_arr->stride = stride;
	bkt_arr->alignment = alignment;
	bkt_arr->idx_unused = (unsigned long long *)pk_new_bkt(sizeof(unsigned long long), alignof(unsigned long long), bkt_buckets);
	bkt_arr->idx_unused[0] = PK_BKT_ARR_ALL_UNUSED_VAL;
	bkt_arr->bucketed_data = (void **)pk_new_bkt(sizeof(void *), alignof(void *), bkt_buckets);
	bkt_arr->bucketed_data[0] = pk_new_bkt(stride * limits.i, alignment, bkt_data);
}

void pk_bkt_arr_clear(struct pk_bkt_arr *bkt_arr) {
	unsigned int b;
	bkt_arr->head_l.b = 0;
	bkt_arr->head_l.i = 0;
	bkt_arr->head_r.b = 0;
	bkt_arr->head_r.i = 0;
	for (b = 0; b < bkt_arr->reserved_buckets; ++b) {
		bkt_arr->idx_unused[b] = PK_BKT_ARR_ALL_UNUSED_VAL;
	}
}

void pk_bkt_arr_reserve(struct pk_bkt_arr *bkt_arr, size_t count) {
	size_t bucket_count = count / bkt_arr->limits.i;
	if (bkt_arr->reserved_buckets >= bucket_count) return;
	unsigned long long *new_idx_unused = (unsigned long long *)pk_new_bkt(sizeof(unsigned long long) * bucket_count, alignof(uint64_t), bkt_arr->bkt_buckets);
	void **new_bucketed_data = (void **)pk_new_bkt(sizeof(void *) * bucket_count, alignof(void *), bkt_arr->bkt_buckets);
	if (bkt_arr->reserved_buckets > 0) {
		memcpy(new_idx_unused, bkt_arr->idx_unused, sizeof(unsigned long long) * bkt_arr->reserved_buckets);
		memcpy(new_bucketed_data, bkt_arr->bucketed_data, sizeof(void *) * bkt_arr->reserved_buckets);
		pk_delete_bkt(bkt_arr->bucketed_data, sizeof(void *) * bkt_arr->reserved_buckets, bkt_arr->bkt_buckets);
		pk_delete_bkt(bkt_arr->idx_unused, sizeof(unsigned long long) * bkt_arr->reserved_buckets, bkt_arr->bkt_buckets);
	}
	for (size_t i = bkt_arr->reserved_buckets; i < bucket_count; ++i) {
		new_idx_unused[i] = PK_BKT_ARR_ALL_UNUSED_VAL;
		new_bucketed_data[i] = pk_new_bkt(bkt_arr->stride * bkt_arr->limits.i, bkt_arr->alignment, bkt_arr->bkt_data);
	}
	bkt_arr->idx_unused = new_idx_unused;
	bkt_arr->bucketed_data = new_bucketed_data;
	bkt_arr->reserved_buckets = bucket_count;
}

struct pk_bkt_arr_handle pk_bkt_arr_find_first_handle(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_compare_fn fn, void *user_data, const void *user_obj_data) {
	assert(bkt_arr != NULL);
	assert(fn != NULL);
	struct pk_bkt_arr_handle ret;
	unsigned int b, i, ii;
	ret.b = PK_BKT_ARR_HANDLE_B_MAX;
	ret.i = PK_BKT_ARR_HANDLE_I_MAX;
	for (b = 0; b < bkt_arr->reserved_buckets; ++b) {
		char *arr = ((char**)(bkt_arr->bucketed_data))[b];
		ii = b == bkt_arr->reserved_buckets-1 ? bkt_arr->head_r.i : bkt_arr->limits.i;
		for (i = 0; i < ii; ++i) {
			if (PK_HAS_FLAG(bkt_arr->idx_unused[b], 1ull << i)) {
				continue;
			}
			if (fn(user_data, user_obj_data, arr+(bkt_arr->stride * i))) {
				ret.b = b;
				ret.i = i;
				return ret;
			}
		}
	}
	return ret;
}

void pk_bkt_arr_iterate(struct pk_bkt_arr *bkt_arr, pk_bkt_arr_iterate_fn fn, void *user_data) {
	assert(bkt_arr != NULL);
	assert(fn != NULL);
	unsigned int b, i, ii;
	for (b = 0; b < bkt_arr->reserved_buckets; ++b) {
		char *arr = ((char**)(bkt_arr->bucketed_data))[b];
		ii = b == bkt_arr->head_r.b ? bkt_arr->head_r.i : bkt_arr->limits.i;
		for (i = 0; i < ii; ++i) {
			if (PK_HAS_FLAG(bkt_arr->idx_unused[b], 1ull << i)) {
				continue;
			}
			fn(user_data, arr+(bkt_arr->stride * i));
		}
	}
}

void pk_bkt_arr_teardown(struct pk_bkt_arr *bkt_arr)
{
	int b;
	size_t sz = bkt_arr->limits.i * bkt_arr->stride;
	if (bkt_arr->idx_unused == nullptr && bkt_arr->bucketed_data == nullptr) return;
	for (b = bkt_arr->reserved_buckets - 1; b > -1; --b) {
		pk_delete_bkt(bkt_arr->bucketed_data[b], sz, bkt_arr->bkt_data);
	}
	pk_delete_bkt((void *)bkt_arr->idx_unused, sizeof(unsigned long long) * (bkt_arr->reserved_buckets), bkt_arr->bkt_buckets);
	pk_delete_bkt((void *)bkt_arr->bucketed_data, sizeof(void *) * (bkt_arr->reserved_buckets), bkt_arr->bkt_buckets);
	memset(bkt_arr, 0, sizeof(struct pk_bkt_arr));
	bkt_arr->bkt_buckets = NULL;
	bkt_arr->bkt_data = NULL;
	bkt_arr->idx_unused = NULL;
	bkt_arr->bucketed_data = NULL;
}

struct pk_bkt_arr_handle pk_bkt_arr_new_handle(struct pk_bkt_arr *bkt_arr)
{
	struct pk_bkt_arr_handle ret;
	unsigned int b, i, ii;
	assert(bkt_arr != nullptr);
	// if we have an existing open slot
	if (pk_bkt_arr_handle_compare(bkt_arr->head_l, bkt_arr->head_r) != 0) {
		ret = bkt_arr->head_l;
		for (b = bkt_arr->head_l.b; b < bkt_arr->reserved_buckets; ++b) {
			if (bkt_arr->idx_unused[b] == 0ull) continue;
			// I feel like you could do a binary search here, but for 64 elements is it worth it?
			i = bkt_arr->head_l.b == b ? bkt_arr->head_l.i + 1 : 0;
			ii = bkt_arr->head_r.b == b ? bkt_arr->head_r.i : PK_MIN(64, bkt_arr->limits.i);
			for (; i < ii; ++i) {
				if (bkt_arr->idx_unused[b] & (1ull << i)) {
					bkt_arr->head_l.b = b;
					bkt_arr->head_l.i = i;
					goto done;
				}
			}
		}
		bkt_arr->head_l = bkt_arr->head_r;
		goto done;
	}
	if (pk_bkt_arr_handle_compare(pk_bkt_arr_handle_increment(bkt_arr, bkt_arr->head_l), bkt_arr->head_l) == 0 && bkt_arr->reserved_buckets == bkt_arr->limits.b && bkt_arr->idx_unused[bkt_arr->head_r.b] == 0) {
		PK_LOGV_ERR("[pk_bkt_arr_new_handle] Exceeded bucket limits!: b:%u i:%u\n", bkt_arr->limits.b, bkt_arr->limits.i);
		exit(1);
	}
	if (bkt_arr->head_r.b == bkt_arr->reserved_buckets && bkt_arr->head_r.i == 0) {
		bkt_arr->reserved_buckets += 1;

		unsigned long long *new_idx_unused = (unsigned long long *)pk_new_bkt(sizeof(unsigned long long) * bkt_arr->reserved_buckets, alignof(unsigned long long), bkt_arr->bkt_buckets);
		void **new_data_ptrs = (void **)pk_new_bkt(sizeof(void *) * bkt_arr->reserved_buckets, alignof(void *), bkt_arr->bkt_buckets);

		for (b = 0; b < bkt_arr->reserved_buckets - 1; ++b) {
			new_idx_unused[b] = bkt_arr->idx_unused[b];
			new_data_ptrs[b] = bkt_arr->bucketed_data[b];
		}
		new_idx_unused[bkt_arr->reserved_buckets - 1] = PK_BKT_ARR_ALL_UNUSED_VAL;
		new_data_ptrs[bkt_arr->reserved_buckets - 1] = pk_new_bkt(bkt_arr->stride * bkt_arr->limits.i, bkt_arr->alignment, bkt_arr->bkt_data);

		pk_delete_bkt((void *)bkt_arr->idx_unused, sizeof(unsigned long long) * (bkt_arr->reserved_buckets - 1), bkt_arr->bkt_buckets);
		pk_delete_bkt((void *)bkt_arr->bucketed_data, sizeof(void *) * (bkt_arr->reserved_buckets - 1), bkt_arr->bkt_buckets);
		bkt_arr->idx_unused = new_idx_unused;
		bkt_arr->bucketed_data = new_data_ptrs;
	}
	ret = bkt_arr->head_r;
	bkt_arr->head_r = pk_bkt_arr_handle_increment(bkt_arr, bkt_arr->head_r);
	bkt_arr->head_l = pk_bkt_arr_handle_increment(bkt_arr, bkt_arr->head_l);
done:
	bkt_arr->idx_unused[ret.b] &= ~(1ull << ret.i);
	return ret;
}

void pk_bkt_arr_free_handle(struct pk_bkt_arr *bkt_arr, struct pk_bkt_arr_handle handle)
{
	assert(bkt_arr != nullptr);
	assert(pk_bkt_arr_handle_validate(bkt_arr, handle) == PK_BKT_ARR_HANDLE_VALIDATION_VALID);
	bkt_arr->idx_unused[handle.b] |= (1ull << handle.i);
	if (handle.b < bkt_arr->head_l.b || (handle.b == bkt_arr->head_l.b && handle.i < bkt_arr->head_l.i)) {
		bkt_arr->head_l = handle;
		return;
	}
}

int pk_bkt_arr_handle_compare(struct pk_bkt_arr_handle lhs, struct pk_bkt_arr_handle rhs)
{
	if (lhs.b == rhs.b && lhs.i == rhs.i) return 0;
	if (lhs.b == rhs.b) return (int)rhs.i - (int)lhs.i;
	return (int)rhs.b - (int)lhs.b;
}

struct pk_bkt_arr_handle pk_bkt_arr_handle_increment(struct pk_bkt_arr *arr, struct pk_bkt_arr_handle h)
{
	h.i += 1;
	if (arr->limits.i == h.i) {
		if (h.b + 1 < arr->limits.b) {
			h.b += 1;
			h.i = 0;
		} else {
			h.i -= 1;
		}
	}
	return h;
}

struct pk_bkt_arr_handle pk_bkt_arr_handle_decrement(struct pk_bkt_arr *arr, struct pk_bkt_arr_handle h)
{
	if (h.i == 0) {
		if (h.b != 0) {
			h.b -= 1;
			h.i = arr->limits.i;
		} else {
			return h;
		}
	}
	h.i -= 1;
	return h;
}

bool pk_bkt_arr_iter_begin(struct pk_bkt_arr *arr, struct pk_iter *it) {
	it->data = nullptr;
	it->id.bkt.b = 0;
	it->id.bkt.i = 0;
	if (arr->head_l.b == 0 && arr->head_l.i == 0 && (arr->head_l.b != arr->head_r.b || arr->head_l.i != arr->head_r.i))  {
		return pk_bkt_arr_iter_increment(arr, it);
	}
	if ((arr->idx_unused[it->id.bkt.b] & (1ull << it->id.bkt.i)) != 0) return false;
	it->data = (char*)(arr->bucketed_data[it->id.bkt.b]) + (arr->stride * it->id.bkt.i);
	return true;
}

bool pk_bkt_arr_iter_end(struct pk_bkt_arr *arr, struct pk_iter *it) {
	it->data = nullptr;
	it->id.bkt.b = 0;
	it->id.bkt.i = 0;
	if (arr->head_r.b == 0 && arr->head_r.i == 0) return false;
	do {
		struct pk_bkt_arr_handle handle = arr->head_r;
		for (;;) {
			if ((arr->idx_unused[handle.b] & (1ull << handle.i)) == 0) break;
			if (handle.b == 0 && handle.i == 0) return false;
			handle = pk_bkt_arr_handle_decrement(arr, handle);
		}
		it->id.bkt.b = handle.b;
		it->id.bkt.i = handle.i;
		break;
	} while (true);
	if (arr->bucketed_data != nullptr && arr->bucketed_data[it->id.bkt.b] != nullptr) {
		it->data = (char*)(arr->bucketed_data[it->id.bkt.b]) + (arr->stride * it->id.bkt.i);
		return true;
	}
	return false;
}

bool pk_bkt_arr_iter_increment(struct pk_bkt_arr *arr, struct pk_iter *it) {
	struct pk_bkt_arr_handle handle = {
		.b = it->id.bkt.b,
		.i = it->id.bkt.i,
	};
	if (it->id.bkt.b == arr->limits.b-1 && it->id.bkt.i == arr->limits.i-1) return false;
	for (;;) {
		handle = pk_bkt_arr_handle_increment(arr, handle);
		if (handle.b >= arr->reserved_buckets) return false;
		if ((arr->idx_unused[handle.b] & (1ull << handle.i)) == 0) break;
	}
	it->id.bkt.b = handle.b;
	it->id.bkt.i = handle.i;
	if ((arr->idx_unused[it->id.bkt.b] & (1ull << it->id.bkt.i)) != 0) return false;
	it->data = (char*)(arr->bucketed_data[it->id.bkt.b]) + (arr->stride * it->id.bkt.i);
	return true;
}

bool pk_bkt_arr_iter_decrement(struct pk_bkt_arr *arr, struct pk_iter *it) {
	struct pk_bkt_arr_handle handle = {
		.b = it->id.bkt.b,
		.i = it->id.bkt.i,
	};
	for (;;) {
		handle = pk_bkt_arr_handle_decrement(arr, handle);
		if ((arr->idx_unused[handle.b] & (1ull << handle.i)) == 0) break;
		if (handle.b == 0 && handle.i == 0) break;
	}
	if (it->id.bkt.b == handle.b && it->id.bkt.i == handle.i) return false;
	it->id.bkt.b = handle.b;
	it->id.bkt.i = handle.i;
	if ((arr->idx_unused[it->id.bkt.b] & (1ull << it->id.bkt.i)) != 0) return false;
	it->data = ((char*)(arr->bucketed_data[it->id.bkt.b])) + (arr->stride * it->id.bkt.i);
	return true;
}

#endif /* PK_IMPL_BKTARR */