#include <Core/Common/CoreTypes.h>
#include <Core/Common/CoreUtils.h>
#include <Core/Memory/Allocator.h>
#include <Core/Memory/MemoryArena.h>

#if JULIET_DEBUG

namespace Juliet::UnitTest
{
    struct TestStruct
    {
        uint8  First;
        uint16 Second;
    };
    // Need access to internal Pool functions? They are in the header now!
    // MemoryPool is declared in header.
    void TestMemoryArena()
    {
        printf("Running Paged Memory Arena Tests...\n");

        // New Arena!
        ArenaParams param{ .ReserveSize = Megabytes(64llu), .CommitSize = Kilobytes(64llu) };
        Arena*      testArena = ArenaAllocate(param JULIET_DEBUG_ONLY(, "Test Arena"));

        size_t pos = ArenaPos(testArena);
        Assert(pos == k_ArenaHeaderSize);

        TestStruct* myStruct = ArenaPushStruct<TestStruct>(testArena);
        Assert(myStruct->First == 0);
        Assert(myStruct->Second == 0);
        pos = ArenaPos(testArena);
        Assert(pos == k_ArenaHeaderSize + sizeof(TestStruct));

        TestStruct* myStruct2 = ArenaPushStruct<TestStruct>(testArena);
        Assert(myStruct2->First == 0);
        Assert(myStruct2->Second == 0);

        myStruct->First   = 5;
        myStruct->Second  = 125;
        myStruct2->First  = 5;
        myStruct2->Second = 125;

        ArenaClear(testArena);
        pos = ArenaPos(testArena);
        Assert(pos == k_ArenaHeaderSize);

        myStruct = ArenaPushStruct<TestStruct>(testArena);
        Assert(myStruct->First == 0);
        Assert(myStruct->Second == 0);
        pos = ArenaPos(testArena);
        Assert(pos == k_ArenaHeaderSize + sizeof(TestStruct));

        TestStruct* myStruct3 = ArenaPushStruct<TestStruct>(testArena);
        Assert(myStruct3->First == 0);
        Assert(myStruct3->Second == 0);
        pos = ArenaPos(testArena);
        Assert(pos == k_ArenaHeaderSize + sizeof(TestStruct) * 2);

        ArenaClear(testArena);

        // TEST COMMIT
        // Push more than the default commit size
        uint8* myArray = ArenaPushArray<uint8>(testArena, 100'000);
        *myArray       = 5;
        Assert(*myArray == 5);
        bool isCommitSizeSame = testArena->Committed == (2 * param.CommitSize);
        Assert(isCommitSizeSame);
        ArenaClear(testArena);

        // Push again above default commit size
        myArray = ArenaPushArray<uint8>(testArena, 100'000);
        // Shoud be zero again
        Assert(*myArray == 0);
        // Should not have commited more
        isCommitSizeSame = testArena->Current->Committed == (2 * param.CommitSize);
        Assert(isCommitSizeSame);
        ArenaClear(testArena);

        // TEST RESERVE
        size_t posPreBigArray          = ArenaPos(testArena);
        size_t reservedSizePreBigArray = testArena->Current->Reserved;

        size_t arraySize = Megabytes(100llu);
        myArray          = ArenaPushArray<uint8>(testArena, arraySize);

        // Align and support the header size because we need to allocate a new block + header
        size_t reserve_size = AlignPow2(arraySize + k_ArenaHeaderSize, AlignOf(uint8));
        Assert(reserve_size == testArena->Current->Position);
        Assert(reservedSizePreBigArray == testArena->Current->BasePosition);

        // TODO Get page size from os kernel call (dwPageSize)
        constexpr uint64 k_PageSizeUnitTest      = Kilobytes(4);
        size_t           reserve_sizeAlignToPage = AlignPow2(reserve_size, k_PageSizeUnitTest);

        bool isReserveSizeSame = testArena->Current->Reserved == reserve_sizeAlignToPage;
        Assert(isReserveSizeSame);
        Assert(testArena->Current->Reserved == testArena->Current->Committed);

        size_t basePos = ArenaPos(testArena);
        Assert(basePos == reserve_size + testArena->Current->BasePosition);

        myStruct = ArenaPushStruct<TestStruct>(testArena);

        // Should not delete the current block because we only delete the struct
        ArenaPopTo(testArena, basePos);
        Assert(testArena->Current->Previous != nullptr);

        // Should still not delete the current block. We only remove the array
        ArenaPop(testArena, arraySize);
        Assert(testArena->Current->Previous != nullptr);
        // We should only have the header remaining
        Assert(testArena->Current->Position == k_ArenaHeaderSize);

        // This one should delete the block
        ArenaPopTo(testArena, posPreBigArray);
        Assert(testArena->Current->Previous == nullptr);

        ArenaRelease(testArena);

        {
            // Test reallocate
            Arena* arena           = ArenaAllocate({ .AllowRealloc = true } JULIET_DEBUG_ONLY(, "Test Realloc"));
            char*  charArray       = ArenaPushArray<char>(arena, 128);
            char*  secondCharArray = ArenaPushArray<char>(arena, 128);
            char*  thirdCharArray  = ArenaPushArray<char>(arena, 128);

            secondCharArray = static_cast<char*>(ArenaReallocate(arena, secondCharArray, 128, 256, alignof(char), true JULIET_DEBUG_ONLY(, "ReallocChar")));

            char* fourthCharArray = ArenaPushArray<char>(arena, 128);

            Assert(charArray);
            Assert(secondCharArray);
            Assert(thirdCharArray);
            Assert(fourthCharArray);

            ArenaRelease(arena);
        }

        {
            // Test Reallocate Shrink (Buffer Overflow Bug Check)
            Arena* arena = ArenaAllocate({ .AllowRealloc = true } JULIET_DEBUG_ONLY(, "Test Shrink"));
            
            size_t largeSize = 100;
            char* large = ArenaPushArray<char>(arena, largeSize);
            MemSet(large, 'A', largeSize);

            size_t smallSize = 50;
            char* smallData = static_cast<char*>(ArenaReallocate(arena, large, largeSize, smallSize, alignof(char), true JULIET_DEBUG_ONLY(, "ResizeSmall")));

            for (size_t i = 0; i < smallSize; ++i)
            {
                Assert(smallData[i] == 'A');
            }

            // Allocate next block (should be immediately after 'small')
            // Don't zero it, if overflow happened it will contain 'A's
            char* next = static_cast<char*>(ArenaPush(arena, 50, alignof(char), false JULIET_DEBUG_ONLY(, "OverflowCheck")));

            bool corrupted = false;
            for(size_t i = 0; i < 50; ++i)
            {
                if (next[i] == 'A')
                {
                    corrupted = true;
                    break;
                }
            }
            
            Assert(!corrupted);

            ArenaRelease(arena);
        }

        printf("All Paged MemoryArena tests passed.\n");
    }
} // namespace Juliet::UnitTest
#endif
