#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));
	bkt.ptr = bkt.blocks + (sizeof(MemBlock) * bkt.maxBlockCount);
	bkt.transient = transient;
	auto *memBlock = reinterpret_cast<MemBlock *>(bkt.blocks);
	memBlock->data = bkt.ptr;
	memBlock->size = sz - (sizeof(MemBlock) * bkt.maxBlockCount);
	return bucketHead++;
}

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;
	int64_t availBlockIndex = -1;
	void *data = nullptr;
	for (int64_t i = 0; i <= bkt->lastEmptyBlockIndex; ++i) {
		auto blk = blocks[i];
		std::size_t misalignment = reinterpret_cast<std::size_t>(blk.data) % calculatedAlignment;
		if (misalignment != 0)
			continue;
		if (blk.size > sz) {
			availBlockIndex = i;
			break;
		}
	}
	if (availBlockIndex != -1) {
		data = blocks[availBlockIndex].data;
		blocks[availBlockIndex].data += sz;
		blocks[availBlockIndex].size -= sz;
		assert(blocks[availBlockIndex].size < bkt->size);
		if (blocks[availBlockIndex].size == 0) {
			if (availBlockIndex != bkt->lastEmptyBlockIndex) {
				blocks[availBlockIndex].data = blocks[bkt->lastEmptyBlockIndex].data;
				blocks[availBlockIndex].size = blocks[bkt->lastEmptyBlockIndex].size;
			}
			bkt->lastEmptyBlockIndex -= 1;
		}
	} else {
		assert(bkt->head + sz <= bkt->size && "memory bucket specified, but full");
		std::size_t misalignment = reinterpret_cast<std::size_t>(bkt->ptr + bkt->head) % calculatedAlignment;
		std::size_t padding = calculatedAlignment - misalignment;
		if (misalignment != 0) {
			bkt->lastEmptyBlockIndex += 1;
			blocks[bkt->lastEmptyBlockIndex].data = bkt->ptr + bkt->head;
			blocks[bkt->lastEmptyBlockIndex].size = padding;
		}
		data = bkt->ptr + bkt->head + padding;
		bkt->head += sz + padding;
	}
	bkt->allocs++;
	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 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;
	} else {
		MemBlock *blocks = reinterpret_cast<MemBlock *>(bkt->blocks);
		for (int64_t i = 0; i < bkt->lastEmptyBlockIndex; ++i) {
			auto &blk = blocks[i];
			if (blk.data == reinterpret_cast<const char *>(ptr) + sz) {
				blk.data -= sz;
				blk.size += sz;
				return;
			}
			if (reinterpret_cast<const char *>(ptr) == blk.data + blk.size ) {
				blk.size += sz;
				return;
			}
		}
		if (bkt->lastEmptyBlockIndex != bkt->maxBlockCount) {
			bkt->lastEmptyBlockIndex += 1;
			blocks[bkt->lastEmptyBlockIndex].data = reinterpret_cast<char *>(const_cast<void *>(ptr));
			blocks[bkt->lastEmptyBlockIndex].size = sz;
		} else {
			bkt->lostBytes += sz;
		}
	}
}

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]);
	}
}

uint64_t Buckets_NewHandle(std::size_t bucketBytes, std::size_t alignment, uint64_t bucketItemCount, uint64_t &bucketIncrementer, uint64_t &bucketCounter, uint64_t &itemCounter, void*& buckets) {
	uint64_t newHandle{itemCounter | bucketCounter};

	std::size_t calculatedAlignment = alignment < MINIMUM_ALIGNMENT ? MINIMUM_ALIGNMENT : alignment;

	itemCounter += uint64_t{1ULL << 32};
	if (itemCounter > uint64_t{(bucketItemCount - 1) << 32}) {
		itemCounter = 0ULL;
		bucketCounter += 1;
	}
	if (bucketCounter > bucketIncrementer) {
		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 newHandle;
}

void Pke_DebugPrint() {
	printf("Memory Manager printout:\nBucket count: %li\n", bucketHead + 1);
	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);
	}
}