#include "memory.hpp"

const std::size_t MINIMUM_ALIGNMENT = 1;
const std::size_t MAXIMUM_ALIGNMENT = 64;

struct MemBlock {
	char *data;
	std::size_t size;
};

struct MemBucket {
	int64_t size;
	int64_t head;
	int64_t lostBytes;
	int64_t allocs;
	int64_t lastEmptyBlockIndex;
	int64_t maxBlockCount;
	char *blocks;
	char *ptr;
	bool transient;
};

MemBucket buckets[128];
int64_t bucketHead = 0;

int64_t InitNewBucket(int64_t sz, bool transient = false) {
	int64_t blockCount = sz * 0.01;
	auto &bkt = buckets[bucketHead];
	bkt.size = sz;
	bkt.head = 0;
	bkt.lostBytes = 0;
	bkt.allocs = 0;
	bkt.lastEmptyBlockIndex = -1;
	bkt.maxBlockCount = blockCount < 10 ? 10 : blockCount;
	bkt.blocks = reinterpret_cast<char *>(std::malloc(sz));
	assert(bkt.blocks != nullptr && "failed to allocate memory");
	bkt.ptr = bkt.blocks + (sizeof(MemBlock) * bkt.maxBlockCount);
	size_t misalignment = reinterpret_cast<uint64_t>(bkt.ptr) % MAXIMUM_ALIGNMENT;
	if (misalignment != 0) {
		size_t moreBlocks = misalignment / sizeof(MemBlock);
		bkt.maxBlockCount += moreBlocks;
		bkt.ptr += (MAXIMUM_ALIGNMENT - misalignment);
	}
	bkt.transient = transient;
	auto *memBlock = reinterpret_cast<MemBlock *>(bkt.blocks);
	memBlock->data = bkt.ptr;
	memBlock->size = sz - (sizeof(MemBlock) * bkt.maxBlockCount);
	return bucketHead++;
}

void Pke_ResetBucket(MemBucket *bkt) {
	bkt->head = 0;
	bkt->lostBytes = 0;
	bkt->allocs = 0;
	bkt->lastEmptyBlockIndex = -1;
}

void DestroyBucket(MemBucket *bkt) {
	std::free(bkt->blocks);
	bkt->size = 0;
	bkt->head = 0;
	bkt->lostBytes = 0;
	bkt->allocs = 0;
	bkt->lastEmptyBlockIndex = -1;
	bkt->maxBlockCount = 0;
	bkt->blocks = CAFE_BABE(char);
	bkt->ptr = CAFE_BABE(char);
	bkt->transient = false;
}

void *Pke_New(std::size_t sz, std::size_t alignment, MemBucket *bkt) {
	MemBlock *blocks = reinterpret_cast<MemBlock *>(bkt->blocks);
	std::size_t calculatedAlignment = alignment < MINIMUM_ALIGNMENT ? MINIMUM_ALIGNMENT : alignment;
	std::size_t misalignment = 0;
	long index = -1;
	MemBlock *block = nullptr;
	void *data = nullptr;
	for (int64_t i = 0; i <= bkt->lastEmptyBlockIndex; ++i) {
		auto &blk = blocks[i];
		misalignment = reinterpret_cast<std::size_t>(blk.data) % calculatedAlignment;
		misalignment = (calculatedAlignment - misalignment) % calculatedAlignment;
		if (blk.size >= sz + misalignment) {
			index = i;
			block = &blk;
			break;
		}
	}
	if (block != nullptr) {
		MemBlock *nextBlock = &blocks[index + 1];
		bool touchesValidNeighbor = index != bkt->lastEmptyBlockIndex && nextBlock->data == block->data + block->size; 
		data = block->data + misalignment;
		if (misalignment) {
			size_t afterSize = block->size - misalignment - sz;
			block->size = misalignment;
			if (afterSize) {
				if (index == bkt->maxBlockCount) {
					bkt->lostBytes += afterSize;
				} else if (touchesValidNeighbor) {
					nextBlock->data -= afterSize;
					nextBlock->size += afterSize;
				} else {
					if (index != bkt->lastEmptyBlockIndex) {
						long moveCount = bkt->lastEmptyBlockIndex - index;
						if (moveCount > 0) {
							char *srcPos = bkt->blocks + (sizeof(MemBlock) * (index + 1));
							char *dstPos = bkt->blocks + (sizeof(MemBlock) * (index + 2));
							memmove(dstPos, srcPos, sizeof(MemBlock) * moveCount);
						}
					}
					bkt->lastEmptyBlockIndex += 1;
					nextBlock->data = block->data + misalignment + sz;
					nextBlock->size = afterSize;
				}
			}
		} else {
			size_t afterSize = block->size - sz;
			if (afterSize && !touchesValidNeighbor) {
					block->data += sz;
					block->size -= sz;
			} else {
				if (afterSize && touchesValidNeighbor) {
					nextBlock->data -= afterSize;
					nextBlock->size += afterSize;
				}
				long moveCount = bkt->lastEmptyBlockIndex - index;
				bkt->lastEmptyBlockIndex -= 1;
				if (moveCount > 0) {
					char *srcPos = bkt->blocks + (sizeof(MemBlock) * (index + 1));
					char *dstPos = bkt->blocks + (sizeof(MemBlock) * (index + 0));
					memmove(dstPos, srcPos, sizeof(MemBlock) * moveCount);
				}
			}
		}
	} else {
		assert(bkt->head + sz <= bkt->size && "memory bucket specified, but full");
		misalignment = reinterpret_cast<std::size_t>(bkt->ptr + bkt->head) % calculatedAlignment;
		misalignment = (calculatedAlignment - misalignment) % calculatedAlignment;
		if (misalignment != 0) {
			if (bkt->lastEmptyBlockIndex == bkt->maxBlockCount) {
				bkt->lostBytes += misalignment;
			} else {
				bkt->lastEmptyBlockIndex += 1;
				blocks[bkt->lastEmptyBlockIndex].data = bkt->ptr + bkt->head;
				blocks[bkt->lastEmptyBlockIndex].size = misalignment;
			}
			bkt->head = bkt->head + misalignment;
		}
		data = bkt->ptr + bkt->head;
		bkt->head = bkt->head + sz;
	}
	bkt->allocs++;
	assert(data >= bkt->ptr && "allocated data is before bucket data");
	assert(data <= bkt->ptr + bkt->size && "allocated data is after bucket data");
	return data;
}

void *Pke_New(std::size_t sz, std::size_t alignment) {
	MemBucket *bkt = nullptr;
	for (long i = 0; i < bucketHead; ++i) {
		if (buckets[i].transient == false && buckets[i].size - buckets[i].head > sz + MAXIMUM_ALIGNMENT) {
			bkt = &buckets[i];
		}
	}
	if (bkt == nullptr) {
		bkt = &buckets[InitNewBucket(DEFAULT_BUCKET_SIZE)];
	}
	return Pke_New(sz, alignment, bkt);
}

void inline Pke_CollapseEmptyBlocksToHead(MemBucket *bkt) {
	MemBlock *blocks = reinterpret_cast<MemBlock *>(bkt->blocks);
	while (bkt->lastEmptyBlockIndex > -1) {
		MemBlock *lastBlock = &blocks[bkt->lastEmptyBlockIndex];
		if (lastBlock->data + lastBlock->size != bkt->ptr + bkt->head) {
			return;
		}
		bkt->head -= lastBlock->size;
		lastBlock->data = 0;
		lastBlock->size = 0;
		bkt->lastEmptyBlockIndex -= 1;
	}
}

void inline Pke_CollapseBlocks(MemBucket *bkt) {
	/* JCB 2023-11-21
	 * After writing this, I realized that if the MemBlocks are sorted,
	 *  one free should cause at most 1 memmove shift, and only if both
	 *  the block before and after are identified. This method does not
	 *  know when to stop and is therefore less efficient, especially
	 *  considering it would find 1 instance to memmove at most if called
	 *  on every free.
	 * Free has been refactored to handle the before+after block scenarios,
	 *  so this method is no longer necessary, unless someday in the future
	 *  we add a mechanism to toggle whether or not to resize blocks on free
	 * Example: if we don't enforce sorting, then this would be a good GC
	 *  after a sort. However, in that case it might be better to rebuild
	 *  the blocks into a new array and memcpy over the original, instead
	 *  of performing memmove an unspecified number of times.
	 */
	assert(!"for reference use only");
	MemBlock *blocks = reinterpret_cast<MemBlock *>(bkt->blocks);
	long skipDistance = 1;
	long lastStartingIndex = 0;
	for (long i = 0; i <= bkt->lastEmptyBlockIndex - 1; ++i) {
		lastStartingIndex = i;
		MemBlock &block = blocks[i];
		MemBlock &nextBlock = blocks[i + skipDistance];
		if (block.data + block.size == nextBlock.data) {
			block.size += nextBlock.size;
			nextBlock.size = 0;
			skipDistance += 1;
			i -= 1;
		} else {
			if (skipDistance > 1) {
				char *srcPos = bkt->blocks + (sizeof(MemBlock) * (i + skipDistance));
				char *dstPos = bkt->blocks + (sizeof(MemBlock) * (i + 1));
				memmove(dstPos, srcPos, sizeof(MemBlock) * (skipDistance - 1));
				bkt->lastEmptyBlockIndex -= (skipDistance - 1);
			}
			i += skipDistance - 1;
			skipDistance = 1;
		}
	}
	if (skipDistance > 1) {
		char *srcPos = bkt->blocks + (sizeof(MemBlock) * (lastStartingIndex + skipDistance));
		char *dstPos = bkt->blocks + (sizeof(MemBlock) * (lastStartingIndex + 1));
		memmove(dstPos, srcPos, sizeof(MemBlock) * (skipDistance - 1));
		bkt->lastEmptyBlockIndex -= (skipDistance - 1);
	}
}

void Pke_Delete(const void *ptr, std::size_t sz, MemBucket *bkt) {
	assert(ptr >= bkt->ptr && ptr < bkt->ptr + bkt->size && "pointer not in memory bucket range");
	bkt->allocs--;
	if (bkt->allocs == 0) {
		bkt->head = 0;
		bkt->lastEmptyBlockIndex = -1;
		return;
	}
	if (ptr == bkt->ptr + bkt->head - sz) {
		bkt->head -= sz;
		Pke_CollapseEmptyBlocksToHead(bkt);
		return;
	}
	MemBlock *blocks = reinterpret_cast<MemBlock *>(bkt->blocks);
	size_t prevBlockIndex = 0xFFFFFFFFFFFFFFFF;
	void *prevPointer = reinterpret_cast<void *>(prevBlockIndex);
	long beforeIndex = -1;
	long afterIndex = -1;
	for (int64_t i = 0; i <= bkt->lastEmptyBlockIndex; ++i) {
		const auto &blk = blocks[i];
		if (blk.data < ptr && prevPointer < blk.data) {
			prevBlockIndex = i;
			prevPointer = blk.data;
		}
		if (reinterpret_cast<const char *>(ptr) == blk.data + blk.size ) {
			beforeIndex = i;
			continue;
		}
		if (blk.data == reinterpret_cast<const char *>(ptr) + sz) {
			afterIndex = i;
			break;
		}
		// we found beforeIndex on the previous loop, but then didn't find afterIndex on this loop
		if (beforeIndex != -1) {
			break;
		}
	}
	if (beforeIndex != -1 && afterIndex != -1) {
		MemBlock &beforeBlk = blocks[beforeIndex];
		MemBlock &afterBlk = blocks[afterIndex];
		beforeBlk.size += sz + afterBlk.size;
		afterBlk.size = 0;
		char *srcPos = bkt->blocks + (sizeof(MemBlock) * (afterIndex + 1));
		char *dstPos = bkt->blocks + (sizeof(MemBlock) * (afterIndex));
		memmove(dstPos, srcPos, sizeof(MemBlock) * (bkt->lastEmptyBlockIndex - afterIndex));
		bkt->lastEmptyBlockIndex -= 1;
	} else if (beforeIndex != -1) {
		MemBlock &beforeBlk = blocks[beforeIndex];
		beforeBlk.size += sz;
	} else if (afterIndex != -1) {
		MemBlock &afterBlk = blocks[afterIndex];
		afterBlk.data -= sz;
		afterBlk.size += sz;
	} else {
		if (bkt->lastEmptyBlockIndex == bkt->maxBlockCount) {
			bkt->lostBytes += sz;
		} else {
			MemBlock *targetBlock = nullptr;
			if (prevBlockIndex < bkt->lastEmptyBlockIndex) {
				long moveCount = bkt->lastEmptyBlockIndex - prevBlockIndex;
				assert(moveCount > 0);
				char *srcPos = bkt->blocks + (sizeof(MemBlock) * (prevBlockIndex + 1));
				char *dstPos = bkt->blocks + (sizeof(MemBlock) * (prevBlockIndex + 2));
				memmove(dstPos, srcPos, sizeof(MemBlock) * moveCount);
				targetBlock = &blocks[prevBlockIndex + 1];
			} else {
				targetBlock = &blocks[bkt->lastEmptyBlockIndex + 1];
			}
			bkt->lastEmptyBlockIndex += 1;
			targetBlock->data = reinterpret_cast<char *>(const_cast<void *>(ptr));
			targetBlock->size = sz;
		}
	}
	Pke_CollapseEmptyBlocksToHead(bkt);
}

void Pke_Delete(const void *ptr, std::size_t sz) {
	MemBucket *bkt = nullptr;
	for (long i = 0; i < bucketHead; ++i) {
		bkt = &buckets[i];
		if (ptr >= bkt->ptr && ptr < bkt->ptr + bkt->size) break;
	}
	assert(bkt != nullptr && "failed to determine correct memory bucket");
	Pke_Delete(ptr, sz, bkt);
}

MemBucket *Pke_BeginTransientBucket(int64_t sz) {
	return &buckets[InitNewBucket(sz, true)];
}

void Pke_EndTransientBucket(MemBucket *bkt) {
	int64_t foundIndex = -1;
	for (int64_t i = 0; i < bucketHead; ++i) {
		if (&buckets[i] == bkt) {
			foundIndex = i;
			DestroyBucket(&buckets[i]);
			break;
		}
	}
	if (foundIndex == bucketHead) {
		bucketHead--;
	}
}

void Pke_MemoryFlush() {
	for (long i = bucketHead - 2; i > -1; --i) {
		if (buckets[i].head != 0) break;
		if (buckets[i+1].head != 0) break;
		if (buckets[i].transient == true) break;
		if (buckets[i+1].transient == true) break;
		bucketHead--;
		DestroyBucket(&buckets[i + 1]);
	}
}

PkeHandle Buckets_NewHandle(std::size_t bucketBytes, std::size_t alignment, PkeHandleItemIndex_T bucketItemCount, PkeHandleBucketIndex_T &bucketIncrementer, PkeHandleBucketIndex_T &bucketCounter, PkeHandleItemIndex_T &itemCounter, void*& buckets, bool &moved) {
	moved = false;
	PkeHandle returnValue {
		.bucketIndex = bucketCounter,
		.itemIndex = itemCounter,
	};

	itemCounter += 1;
	if (itemCounter >= bucketItemCount) {
		itemCounter = 0ULL;
		bucketCounter += 1;
	}
	if (bucketCounter >= bucketIncrementer) {
		std::size_t calculatedAlignment = alignment < MINIMUM_ALIGNMENT ? MINIMUM_ALIGNMENT : alignment;
		moved = true;
		int64_t newIncrement = bucketIncrementer * 1.5;
		char * newBuckets = reinterpret_cast<char *>(Pke_New(bucketBytes * newIncrement, calculatedAlignment));
		std::memcpy(newBuckets, buckets, bucketBytes * bucketIncrementer);
		Pke_Delete(buckets, bucketBytes * bucketIncrementer);
		buckets = newBuckets;
		bucketIncrementer = newIncrement;
	}

	return returnValue;
}

void Pke_DebugPrint() {
	printf("Memory Manager printout:\nBucket count: %li\n", bucketHead);
	for (long i = 0; i < bucketHead; ++i) {
		printf("- bucket #%li\n", i);
		printf("\tsize:                %li\n", buckets[i].size);
		printf("\thead:                %li\n", buckets[i].head);
		printf("\tlostBytes:           %li\n", buckets[i].lostBytes);
		printf("\tallocs:              %li\n", buckets[i].allocs);
		printf("\tlastEmptyBlockIndex: %li\n", buckets[i].lastEmptyBlockIndex);
		printf("\tmaxBlockCount:       %li\n", buckets[i].maxBlockCount);
		printf("\tblocks:              %p\n", buckets[i].blocks);
		printf("\tptr:                 %p\n", buckets[i].ptr);
		printf("\ttransient:           %i\n", buckets[i].transient);
	}
}