#include <Core/Memory/Allocator.h>
#include <Core/Memory/MemoryArena.h>
#include <Core/Memory/Utils.h>

#include <cstring>

namespace Juliet
{
    namespace UnitTest
    {
        extern void TestMemoryArena();
    }

    void MemoryArenaCreate(MemoryArena* arena, void* backingMemory, size_t size)
    {
        Assert(arena);
        Assert(backingMemory);
        arena->Data   = static_cast<uint8*>(backingMemory);
        arena->Size   = size;
        arena->Offset = 0;
#if JULIET_DEBUG
        arena->AllocationCount = 0;
#endif
    }

    void* ArenaPush(MemoryArena* arena, size_t size, size_t alignment, [[maybe_unused]] String tag)
    {
        Assert(arena);

        // Alignment must be power of 2
        Assert((alignment & (alignment - 1)) == 0);

        size_t currentPtr = reinterpret_cast<size_t>(arena->Data + arena->Offset);
        size_t offset     = (currentPtr + (alignment - 1)) & ~(alignment - 1);
        size_t newOffset  = offset - reinterpret_cast<size_t>(arena->Data) + size;

        if (newOffset > arena->Size)
        {
            Assert(false, "Memory Arena overflow");
            return nullptr;
        }

        void* result  = arena->Data + (offset - reinterpret_cast<size_t>(arena->Data));
        arena->Offset = newOffset;

#if JULIET_DEBUG
        if (arena->AllocationCount < MemoryArena::kMaxAllocations)
        {
            arena->Allocations[arena->AllocationCount] = { offset - reinterpret_cast<size_t>(arena->Data), size,
                                                           IsValid(tag) ? tag : WrapString("Unknown") };
            arena->AllocationCount++;
        }
#endif

        return result;
    }

    void* ArenaRealloc(MemoryArena* arena, void* oldPtr, size_t oldSize, size_t newSize, size_t alignment, String tag)
    {
        Assert(arena);
        // Alignment must be power of 2
        Assert((alignment & (alignment - 1)) == 0);

        if (oldPtr == nullptr)
        {
            return ArenaPush(arena, newSize, alignment, tag);
        }

        if (newSize == 0)
        {
            return nullptr;
        }

        // Check if the old allocation is at the top of the stack
        // We need to verify if (oldPtr + oldSize) == (Data + Offset)
        // Note: usage of reinterpret_cast is to be careful with pointer arithmetic on void*

        uint8* oldPtrBytes = static_cast<uint8*>(oldPtr);
        uint8* arenaEnd    = arena->Data + arena->Offset;

        if (oldPtrBytes + oldSize == arenaEnd)
        {
            // It is the last allocation! We can reuse the space.
            // We just need to check if we can expand it (if growing)

            // Re-calculate the offset start for this block (to ensure nothing weird with padding)
            // Ideally oldPtr was aligned.

            // Current Offset corresponds to oldPtrBytes + oldSize.
            // We want to move Offset to oldPtrBytes + newSize.

            size_t oldPtrOffset = static_cast<size_t>(oldPtrBytes - arena->Data);
            size_t newOffset    = oldPtrOffset + newSize;

            if (newOffset > arena->Size)
            {
                // Cannot expand in place, not enough space.
                // Fallthrough to Alloc + Copy
            }
            else
            {
                // In-place Resize success
                arena->Offset = newOffset;
#if JULIET_DEBUG
                // Update the last allocation size
                if (arena->AllocationCount > 0)
                {
                    arena->Allocations[arena->AllocationCount - 1].Size = newSize;
                }
#endif
                return oldPtr;
            }
        }

        // Fallback: Alloc + Copy
        void* newPtr = ArenaPush(arena, newSize, alignment, tag);

        if (newPtr)
        {
            size_t copySize = oldSize < newSize ? oldSize : newSize;
            MemCopy(newPtr, oldPtr, copySize);
        }
        return newPtr;
    }

    void ArenaReset(MemoryArena* arena)
    {
        Assert(arena);
        arena->Offset = 0;
#if JULIET_DEBUG
        arena->AllocationCount = 0;
#endif
    }

    size_t ArenaGetMarker(MemoryArena* arena)
    {
        Assert(arena);
        return arena->Offset;
    }

    void ArenaResetToMarker(MemoryArena* arena, size_t marker)
    {
        Assert(arena);
        Assert(marker <= arena->Offset);
        arena->Offset = marker;

#if JULIET_DEBUG
        while (arena->AllocationCount > 0)
        {
            if (arena->Allocations[arena->AllocationCount - 1].Offset >= marker)
            {
                arena->AllocationCount--;
            }
            else
            {
                break;
            }
        }
#endif
    }

    // --- Global Arenas & Management ---
    namespace
    {
        MemoryArena g_ScratchArena;
        MemoryArena g_EngineArena;
        MemoryArena g_GameArena;

        void* g_ScratchBacking = nullptr;
        void* g_EngineBacking  = nullptr;
        void* g_GameBacking    = nullptr;

        constexpr size_t kScratchSize = 64 * 1024 * 1024;  // 64MB
        constexpr size_t kEngineSize  = 256 * 1024 * 1024; // 256MB
        constexpr size_t kGameSize    = 512 * 1024 * 1024; // 512MB
    } // namespace

    MemoryArena* GetScratchArena()
    {
        return &g_ScratchArena;
    }

    MemoryArena* GetEngineArena()
    {
        return &g_EngineArena;
    }

    MemoryArena* GetGameArena()
    {
        return &g_GameArena;
    }

    void ScratchArenaReset()
    {
        ArenaReset(&g_ScratchArena);
    }

    void MemoryArenasInit()
    {
        // TODO: Use the VirtualAlloc API for this on windows
        g_ScratchBacking = Malloc(kScratchSize);
        MemSet(g_ScratchBacking, 0, kScratchSize);
        g_EngineBacking = Malloc(kEngineSize);
        MemSet(g_EngineBacking, 0, kEngineSize);
        g_GameBacking = Malloc(kGameSize);
        MemSet(g_GameBacking, 0, kGameSize);

        MemoryArenaCreate(&g_ScratchArena, g_ScratchBacking, kScratchSize);
        MemoryArenaCreate(&g_EngineArena, g_EngineBacking, kEngineSize);
        MemoryArenaCreate(&g_GameArena, g_GameBacking, kGameSize);

#if JULIET_DEBUG
        UnitTest::TestMemoryArena();
#endif
    }

    void MemoryArenasShutdown()
    {
        SafeFree(g_ScratchBacking);
        SafeFree(g_EngineBacking);
        SafeFree(g_GameBacking);
    }
} // namespace Juliet
