#include "memory.hpp"

#include <gmock/gmock.h>
#include <gtest/gtest.h>

struct TMemBlock {
	char *data;
	size_t size;
};

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

class MemoryTest : public ::testing::Test {
protected:
	void SetUp() override {
		MemBkt_Test = Pke_BeginTransientBucket(1024 * 1024);
		bkt = reinterpret_cast<TMemBucket *>(MemBkt_Test);
	};
	void TearDown() override {
		Pke_EndTransientBucket(MemBkt_Test);
	};
	MemBucket *MemBkt_Test = nullptr;
	TMemBucket *bkt = nullptr;
};

/*
 * Test that memory can be allocated at the head without misalignment
 */
TEST_F(MemoryTest, EnsureAllocation001) {
	// arrange
	size_t sz = 2;
	size_t alignment = 1;
	
	// act
	Pke_New(sz, alignment, MemBkt_Test);
	
	// assert
	EXPECT_EQ(-1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(sz, bkt->head);
};

/*
 * Test that memory can be allocated at the head with misalignment
 */
TEST_F(MemoryTest, EnsureAllocation002) {
	// arrange

	// act
	Pke_New(2, 1, MemBkt_Test);
	Pke_New(8, 16, MemBkt_Test);

	// assert
	EXPECT_EQ(0, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(24, bkt->head);

	EXPECT_EQ(bkt->ptr + 2, bkt->blocks[0].data);
	EXPECT_EQ(14, bkt->blocks[0].size);
};

/*
 * Test that memory can be allocated into a blank block without misalignment
 */
TEST_F(MemoryTest, EnsureAllocation003) {
	// arrange

	// act
	Pke_New(1, 1, MemBkt_Test);
	Pke_New(1, 64, MemBkt_Test);
	Pke_New(9, 1, MemBkt_Test);

	// assert
	EXPECT_EQ(0, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(65, bkt->head);

	EXPECT_EQ(bkt->ptr + 10, bkt->blocks[0].data);
	EXPECT_EQ(54, bkt->blocks[0].size);
};

/*
 * Test that memory can be allocated into a blank block with misalignment
 */
TEST_F(MemoryTest, EnsureAllocation004) {
	// arrange

	// act
	Pke_New(1, 1, MemBkt_Test);
	Pke_New(1, 64, MemBkt_Test);
	Pke_New(9, 16, MemBkt_Test);

	// assert
	EXPECT_EQ(1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(65, bkt->head);

	EXPECT_EQ(bkt->ptr + 1, bkt->blocks[0].data);
	EXPECT_EQ(15, bkt->blocks[0].size);

	EXPECT_EQ(bkt->ptr + 25, bkt->blocks[1].data);
	EXPECT_EQ(39, bkt->blocks[1].size);
};

/*
 * Test that memory can be freed predictably
 */
TEST_F(MemoryTest, EnsureFree001) {
	// arrange
	size_t sz = 2;
	size_t alignment = 1;
	
	// act
	void *ptr1 = Pke_New(sz, alignment, MemBkt_Test);
	Pke_Delete(ptr1, sz, MemBkt_Test);

	// assert
	EXPECT_EQ(-1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(0, bkt->head);
};

/* allocation events
 * 1st = alignment 01, size 02, move head ( =02 )
 * 2nd = alignment 64, size 16, move head ( =80 )   memBlock { +02, 62 }
 * 3rd = alignment 16, size 16, memBlock { +2, 14 } memBlock { +32, 32 }
 * 4th = alignment 08, size 04, memBlock { +2, 06 } memBlock { +12, 04 } memBlock { +32, 32 }
 */

/*
 * Test that memory can be allocated and freed predictably
 */
class MemoryTestScenario100 : public MemoryTest {
	void SetUp() override {
		MemoryTest::SetUp();
		for (long i = 0; i < cnt; ++i) {
			ptrs[i] = Pke_New(sz[i], alignment[i], MemBkt_Test);
		}
	}
protected:
	void TestExpectedState() {
		EXPECT_EQ(2, bkt->lastEmptyBlockIndex);
		EXPECT_EQ(80, bkt->head);
		EXPECT_EQ(bkt->ptr + 00, ptrs[0]);
		EXPECT_EQ(bkt->ptr + 64, ptrs[1]);
		EXPECT_EQ(bkt->ptr + 16, ptrs[2]);
		EXPECT_EQ(bkt->ptr +  8, ptrs[3]);
	}
	const long cnt = 4;
	const size_t sz[4] = {2, 16, 16, 4};
	const size_t alignment[4] = {1, 64, 16, 8};
	void * ptrs[4];
};

TEST_F(MemoryTestScenario100, EnsureAllocFree0000) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange
	
	// act
	
	// assert
	EXPECT_EQ(bkt->ptr + 2, bkt->blocks[0].data);
	EXPECT_EQ(6, bkt->blocks[0].size);

	EXPECT_EQ(bkt->ptr + 12, bkt->blocks[1].data);
	EXPECT_EQ(4, bkt->blocks[1].size);

	EXPECT_EQ(bkt->ptr + 32, bkt->blocks[2].data);
	EXPECT_EQ(32, bkt->blocks[2].size);
}

TEST_F(MemoryTestScenario100, EnsureAllocFree1000) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange
	size_t expectedHead = 80;
	long freeIndex = 0;

	// act
	/* free events
	 * 1st = ptr +00, size 02, memBlock { +00, 08 }
	 */
	Pke_Delete(ptrs[freeIndex], sz[freeIndex], MemBkt_Test);

	// assert
	EXPECT_EQ(2, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(expectedHead, bkt->head);

	EXPECT_EQ(bkt->ptr + 0, bkt->blocks[0].data);
	EXPECT_EQ(8, bkt->blocks[0].size);

	EXPECT_EQ(bkt->ptr + 12, bkt->blocks[1].data);
	EXPECT_EQ(4, bkt->blocks[1].size);

	EXPECT_EQ(bkt->ptr + 32, bkt->blocks[2].data);
	EXPECT_EQ(32, bkt->blocks[2].size);
};

TEST_F(MemoryTestScenario100, EnsureAllocFree0100) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange
	size_t expectedHead = 32;
	long freeIndex = 1;

	// act
	/* free events
	 * 1st = ptr +64, size 16, move head ( =32 ) memBlock { +02, 06 } memBlock { +12, 4 }
	 */
	Pke_Delete(ptrs[freeIndex], sz[freeIndex], MemBkt_Test);

	// assert
	EXPECT_EQ(1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(expectedHead, bkt->head);

	EXPECT_EQ(bkt->ptr + 2, bkt->blocks[0].data);
	EXPECT_EQ(6, bkt->blocks[0].size);

	EXPECT_EQ(bkt->ptr + 12, bkt->blocks[1].data);
	EXPECT_EQ(4, bkt->blocks[1].size);
};

TEST_F(MemoryTestScenario100, EnsureAllocFree0010) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange
	size_t expectedHead = 80;
	long freeIndex = 2;
	
	// act
	/* free events
	 * 1st = ptr +16, size 16, memBlock { +02, 06 } memBlock { +12, 04 } memBlock { +16, 48 }
	 *                         memBlock { +02, 06 } memBlock { +12, 52 }
	 */
	Pke_Delete(ptrs[freeIndex], sz[freeIndex], MemBkt_Test);
	
	// assert
	EXPECT_EQ(1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(expectedHead, bkt->head);

	EXPECT_EQ(bkt->ptr + 2, bkt->blocks[0].data);
	EXPECT_EQ(6, bkt->blocks[0].size);

	EXPECT_EQ(bkt->ptr + 12, bkt->blocks[1].data);
	EXPECT_EQ(52, bkt->blocks[1].size);
}

TEST_F(MemoryTestScenario100, EnsureAllocFree0001) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange
	size_t expectedHead = 80;
	long freeIndex = 3;
	
	// act
	/* free events
	 * 1st = ptr +08, size 4, memBlock { +02, 06 } memBlock { +08, 08 } memBlock { +32, 32 }
	 *                        memBlock { +02, 14 } memBlock { +32, 32 }
	 */
	Pke_Delete(ptrs[freeIndex], sz[freeIndex], MemBkt_Test);
	
	// assert
	EXPECT_EQ(1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(expectedHead, bkt->head);

	EXPECT_EQ(bkt->ptr + 2, bkt->blocks[0].data);
	EXPECT_EQ(14, bkt->blocks[0].size);

	EXPECT_EQ(bkt->ptr + 32, bkt->blocks[1].data);
	EXPECT_EQ(32, bkt->blocks[1].size);
}

class MemoryTestScenario200 : public MemoryTest {
	void SetUp() override {
		MemoryTest::SetUp();
		for (long i = 0; i < cnt; ++i) {
			ptrs[i] = Pke_New(sz[i], alignment[i], MemBkt_Test);
		}
	}
protected:
	void TestExpectedState() {
		EXPECT_EQ(1, bkt->lastEmptyBlockIndex);
		EXPECT_EQ(65, bkt->head);
		EXPECT_EQ(bkt->ptr + 00, ptrs[0]);
		EXPECT_EQ(bkt->ptr + 64, ptrs[1]);
		EXPECT_EQ(bkt->ptr + 32, ptrs[2]);

		EXPECT_EQ(bkt->ptr + 19, bkt->blocks[0].data);
		EXPECT_EQ(13, bkt->blocks[0].size);

		EXPECT_EQ(bkt->ptr + 33, bkt->blocks[1].data);
		EXPECT_EQ(31, bkt->blocks[1].size);
	}
	const long cnt = 3;
	const size_t sz[3] = {19, 1, 1};
	const size_t alignment[3] = {1, 64, 32};
	void * ptrs[3];
};

/* memory
 * [00-18] (19) ptr0
 * [19-31] (13) MemBlock
 * [ 32  ] ( 1) ptr2
 * [33-63] (31) MemBlock
 * [ 64  ] ( 1) ptr1
 * [ 65  ] ( 0) HEAD
 */

/*
 * Test that memory can be freed touching after only
 */
TEST_F(MemoryTestScenario200, EnsureFreeTouchingAfter) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange

	// act

	// frees [00-18] which gets absorbed into [19-31]
	// [00-31] MemBlock
	Pke_Delete(ptrs[0], 19, MemBkt_Test);

	// assert
	EXPECT_EQ(1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(65, bkt->head);
	EXPECT_EQ(bkt->ptr + 64, ptrs[1]);
	EXPECT_EQ(bkt->ptr + 32, ptrs[2]);
	EXPECT_EQ(bkt->ptr +  0, bkt->blocks[0].data);
	EXPECT_EQ(32, bkt->blocks[0].size);
	EXPECT_EQ(bkt->ptr + 33, bkt->blocks[1].data);
	EXPECT_EQ(31, bkt->blocks[1].size);
};

/*
 * Test that memory can be freed touching before only
 */
TEST_F(MemoryTestScenario200, EnsureFreeTouchingBefore) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange

	// act 1

	// fill everything, then allocate [65], moving HEAD
  // [00-18] (19) ptr0
  // [19-31] (13) ptr3
  // [ 32  ] ( 1) ptr2
  // [33-63] (31) ptr4
  // [ 64  ] ( 1) ptr1
  // [ 65  ] ( 1) ptr5
  // [ 66  ] ( 0) HEAD
	void *ptr3 = Pke_New(13, 1, MemBkt_Test);
	void *ptr4 = Pke_New(31, 1, MemBkt_Test);
	void *ptr5 = Pke_New( 1, 1, MemBkt_Test);

	// assert 1
	EXPECT_EQ(-1, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(66, bkt->head);
	EXPECT_EQ(bkt->ptr + 00, ptrs[0]);
	EXPECT_EQ(bkt->ptr + 19, ptr3);
	EXPECT_EQ(bkt->ptr + 32, ptrs[2]);
	EXPECT_EQ(bkt->ptr + 33, ptr4);
	EXPECT_EQ(bkt->ptr + 64, ptrs[1]);
	EXPECT_EQ(bkt->ptr + 65, ptr5);

	// act 2

	// free [19-31]
  // [00-18] (19) ptr0
  // [19-31] (13) MemBlock
  // [ 32  ] ( 1) ptr2
  // [33-63] (31) ptr4
  // [ 64  ] ( 1) ptr1
  // [ 65  ] ( 1) ptr5
  // [ 66  ] ( 0) HEAD
	Pke_Delete(ptr3, 13, MemBkt_Test);
	ptr3 = nullptr;

	// assert 2
	EXPECT_EQ( 0, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(66, bkt->head);
	EXPECT_EQ(bkt->ptr + 00, ptrs[0]);
	EXPECT_EQ(nullptr,       ptr3);
	EXPECT_EQ(bkt->ptr + 32, ptrs[2]);
	EXPECT_EQ(bkt->ptr + 33, ptr4);
	EXPECT_EQ(bkt->ptr + 64, ptrs[1]);
	EXPECT_EQ(bkt->ptr + 65, ptr5);
	EXPECT_EQ(bkt->ptr + 19, bkt->blocks[0].data);
	EXPECT_EQ(13, bkt->blocks[0].size);

	// act 3

	// free [32] which gets absorbed into 19-32
  // [19-32] (14) MemBlock
  // [33-63] (31) ptr4
	Pke_Delete(ptrs[2], 1, MemBkt_Test);
	ptrs[2] = nullptr;

	// assert 3
	EXPECT_EQ( 0, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(66, bkt->head);
	EXPECT_EQ(bkt->ptr + 00, ptrs[0]);
	EXPECT_EQ(nullptr,       ptr3);
	EXPECT_EQ(nullptr,       ptrs[2]);
	EXPECT_EQ(bkt->ptr + 33, ptr4);
	EXPECT_EQ(bkt->ptr + 64, ptrs[1]);
	EXPECT_EQ(bkt->ptr + 65, ptr5);
	EXPECT_EQ(bkt->ptr + 19, bkt->blocks[0].data);
	EXPECT_EQ(14, bkt->blocks[0].size);
};

/*
 * Test that memory can be freed touching before and after
 */
TEST_F(MemoryTestScenario200, EnsureFreeTouchingBoth) {
	EXPECT_NO_FATAL_FAILURE(this->TestExpectedState());

	// arrange

	// act

	// frees [32] which gets absorbed into [19-31]
  // [00-18] (19) ptr0
  // [19-63] (45) MemBlock
  // [ 64  ] ( 1) ptr1
	Pke_Delete(ptrs[2], 1, MemBkt_Test);
	ptrs[2] = nullptr;

	// assert
	EXPECT_EQ(0, bkt->lastEmptyBlockIndex);
	EXPECT_EQ(65, bkt->head);
	EXPECT_EQ(bkt->ptr +  0, ptrs[0]);
	EXPECT_EQ(bkt->ptr + 64, ptrs[1]);
	EXPECT_EQ(nullptr      , ptrs[2]);
	EXPECT_EQ(bkt->ptr + 19, bkt->blocks[0].data);
	EXPECT_EQ(45, bkt->blocks[0].size);
};