#include <Core/Memory/MemoryArena.h> // For Arena definition
#include <Core/Memory/MemoryArenaDebug.h>

#if JULIET_DEBUG

namespace Juliet
{
    Arena*           g_ArenaGlobalHead         = nullptr;
    Arena*           g_DebugInfoArena          = nullptr;
    ArenaDebugInfo*  g_DebugInfoFreeList       = nullptr;
    ArenaAllocation* g_ArenaAllocationFreeList = nullptr;

    // --- Internal Helpers ---
    static ArenaDebugInfo* AllocDebugInfo()
    {
        if (g_DebugInfoFreeList)
        {
            ArenaDebugInfo* info = g_DebugInfoFreeList;
            g_DebugInfoFreeList  = info->Next;
            info->Next           = nullptr;
            return info;
        }

        if (!g_DebugInfoArena)
        {
            // Create a dedicated arena for debug info
            g_DebugInfoArena = ArenaAllocate(
                { .ReserveSize = Megabytes(16), .CommitSize = Kilobytes(64) } JULIET_DEBUG_ONLY(, "Debug Info Arena"));
        }

        return ArenaPushStruct<ArenaDebugInfo>(g_DebugInfoArena JULIET_DEBUG_PARAM("ArenaDebugInfo"));
    }

    static void FreeDebugInfo(ArenaDebugInfo* info)
    {
        if (info)
        {
            info->Next          = g_DebugInfoFreeList;
            g_DebugInfoFreeList = info;
        }
    }

    void DebugRegisterArena(NonNullPtr<Arena> arena)
    {
        // Add to Global List
        // Note: Not thread safe, assuming single thread for now as per context
        if (g_ArenaGlobalHead)
        {
            g_ArenaGlobalHead->GlobalPrev = arena;
        }
        arena->GlobalNext = g_ArenaGlobalHead;
        arena->GlobalPrev = nullptr;
        g_ArenaGlobalHead = arena;
    }

    void DebugUnregisterArena(NonNullPtr<Arena> arena)
    {
        // Remove from Global List
        if (arena->GlobalPrev)
        {
            arena->GlobalPrev->GlobalNext = arena->GlobalNext;
        }
        else
        {
            g_ArenaGlobalHead = arena->GlobalNext;
        }
        if (arena->GlobalNext)
        {
            arena->GlobalNext->GlobalPrev = arena->GlobalPrev;
        }
        arena->GlobalPrev = nullptr;
        arena->GlobalNext = nullptr;
    }

    void DebugArenaSetDebugName(NonNullPtr<Arena> arena, const char* name)
    {
        arena->Name = name;
    }

    bool IsDebugInfoArena(const Arena* arena)
    {
        return arena == g_DebugInfoArena;
    }

    void DebugArenaFreeBlock(Arena* block)
    {
        ArenaDebugInfo* info = block->FirstDebugInfo;
        while (info)
        {
            ArenaDebugInfo* nextInfo = info->Next;
            FreeDebugInfo(info);
            info = nextInfo;
        }
        block->FirstDebugInfo = nullptr;
    }

    void DebugArenaRemoveAllocation(Arena* block, size_t oldOffset)
    {
        ArenaDebugInfo** prevInfo = &block->FirstDebugInfo;
        ArenaDebugInfo*  info     = block->FirstDebugInfo;

        while (info)
        {
            if (info->Offset == oldOffset) // Found it
            {
                *prevInfo = info->Next;
                FreeDebugInfo(info);
                break;
            }
            prevInfo = &info->Next;
            info     = info->Next;
        }
    }

    void DebugArenaPopTo(Arena* block, size_t newPosition)
    {
        ArenaDebugInfo** prevInfo = &block->FirstDebugInfo;
        ArenaDebugInfo*  info     = block->FirstDebugInfo;
        while (info)
        {
            if (info->Offset >= newPosition)
            {
                ArenaDebugInfo* toFree = info;
                *prevInfo              = info->Next; // Unlink
                info                   = info->Next; // Advance
                FreeDebugInfo(toFree);
            }
            else
            {
                prevInfo = &info->Next;
                info     = info->Next;
            }
        }
    }

    void DebugArenaAddDebugInfo(Arena* block, size_t size, size_t offset, const char* tag)
    {
        ArenaDebugInfo* info = AllocDebugInfo();
        if (info)
        {
            info->Tag             = tag ? tag : "Untagged";
            info->Size            = size;
            info->Offset          = offset;
            info->Next            = block->FirstDebugInfo;
            block->FirstDebugInfo = info;
        }
    }

    Arena* GetGlobalArenaListHead()
    {
        return g_ArenaGlobalHead;
    }

} // namespace Juliet

#endif
