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

#include <algorithm> // For std::max
#include <Core/Common/NonNullPtr.h>
#include <Core/HAL/OS/OS.h>
#include <Core/Logging/LogTypes.h>
#include <cstring>

namespace Juliet
{
    // TODO Get page size from os kernel call (dwPageSize)
    namespace
    {
        constexpr uint64 k_PageSize = Kilobytes(4);
    } // namespace

    struct ArenaFreeNode
    {
        ArenaFreeNode* Next;
        ArenaFreeNode* Previous;
        index_t        Position;
        size_t         Size;
    };

    namespace
    {
        static_assert(sizeof(ArenaFreeNode) == 32, "ArenaFreeNode should be 32 bytes and not more");
        constexpr size_t k_ArenaFreeNodeSize = sizeof(ArenaFreeNode);
    } // namespace

    // https://github.com/EpicGamesExt/raddebugger/blob/master/src/base/base_arena.c

    Arena* ArenaAllocate(const ArenaParams& params JULIET_DEBUG_ONLY(, const char* name), const std::source_location& loc)
    {
        Log(LogLevel::Message, LogCategory::Core, "Allocating from %s : %ul", loc.file_name(), loc.line());

        Byte* baseMem = nullptr;

        uint64 reserve_size = AlignPow2(params.ReserveSize, k_PageSize);
        uint64 commit_size  = AlignPow2(params.CommitSize, k_PageSize);

        // TODO: handle large pages
        baseMem = Memory::OS_Reserve(reserve_size);
        Memory::OS_Commit(baseMem, commit_size);

        Arena* arena   = reinterpret_cast<Arena*>(baseMem);
        arena->Current = arena;

        arena->CommitSize  = commit_size;
        arena->ReserveSize = reserve_size;

        arena->Committed = commit_size;
        arena->Reserved  = reserve_size;

        arena->BasePosition = 0;
        arena->Position     = k_ArenaHeaderSize;

        arena->FreeNodes    = nullptr;
        arena->AllowRealloc = params.AllowRealloc;

#if JULIET_DEBUG
        arena->CanReserveMore = params.CanReserveMore;
        arena->FirstDebugInfo = nullptr;

        DebugArenaSetDebugName(arena, name);
        DebugRegisterArena(arena);
#endif

        return arena;
    }

    void ArenaRelease(NonNullPtr<Arena> arena)
    {
        JULIET_DEBUG_ONLY(DebugUnregisterArena(arena);)

        for (Arena *node = arena->Current, *previous = nullptr; node != nullptr; node = previous)
        {
            previous = node->Previous;

            JULIET_DEBUG_ONLY(DebugArenaFreeBlock(node);)

            Memory::OS_Release(node, node->Reserved);
        }
    }

    void* ArenaPush(NonNullPtr<Arena> arena, size_t size, size_t align, bool shouldBeZeroed JULIET_DEBUG_ONLY(, const char* tag))
    {
        // Assert(IsPowerOfTwo(align));
        Arena* current          = arena->Current;
        size_t positionPrePush  = AlignPow2(current->Position, align);
        size_t positionPostPush = positionPrePush + size;

        if (arena->AllowRealloc)
        {
            size = Max(size, k_ArenaFreeNodeSize);

            for (ArenaFreeNode* freeNode = current->FreeNodes; freeNode != nullptr; freeNode = freeNode->Next)
            {
                if (size <= freeNode->Size)
                {
                    index_t position      = freeNode->Position;
                    size_t  remainingSize = freeNode->Size - size;
                    if (remainingSize < k_ArenaFreeNodeSize)
                    {
                        ArenaFreeNode* previous = freeNode->Previous;
                        ArenaFreeNode* next     = freeNode->Next;
                        if (previous)
                        {
                            previous->Next = next;
                        }
                        if (next)
                        {
                            next->Previous = previous;
                        }

#if JULIET_DEBUG
                        if (remainingSize > 0)
                        {
                            ++current->LostNodeCount;
                        }
#endif
                    }
                    else
                    {
                        freeNode->Position += size;
                    }

                    auto* result = reinterpret_cast<Byte*>(current) + position;
                    if (shouldBeZeroed)
                    {
                        MemoryZero(result, size);
                    }

                    return result;
                }
            }
        }

        // If allowed and needed, add a new block and chain it to the arena.
        if (current->Reserved < positionPostPush /* flags : chaining allowed */)
        {
            Assert(arena->CanReserveMore);

            Arena* newBlock = nullptr;
            {
                Arena* prev_block;
                for (newBlock = arena->FreeBlockLast, prev_block = nullptr; newBlock != nullptr;
                     prev_block = newBlock, newBlock = newBlock->Previous)
                {
                    if (newBlock->Reserved >= AlignPow2(newBlock->Position, align) + size)
                    {
                        if (prev_block)
                        {
                            prev_block->Previous = newBlock->Previous;
                        }
                        else
                        {
                            arena->FreeBlockLast = newBlock->Previous;
                        }
                        break;
                    }
                }
            }

            if (newBlock == nullptr)
            {
                size_t reserveSize = current->ReserveSize;
                size_t commitSize  = current->CommitSize;

                if (size + k_ArenaHeaderSize > reserveSize)
                {
                    reserveSize = AlignPow2(size + k_ArenaHeaderSize, align);
                    commitSize  = AlignPow2(size + k_ArenaHeaderSize, align);
                }

                newBlock = ArenaAllocate({ .ReserveSize = reserveSize, .CommitSize = commitSize } JULIET_DEBUG_ONLY(, arena->Name));
            }

            newBlock->BasePosition = current->BasePosition + current->Reserved;
            SingleLinkedListPushPrevious(arena->Current, newBlock);

            current          = newBlock;
            positionPrePush  = AlignPow2(current->Position, align);
            positionPostPush = positionPrePush + size;
        }

        size_t sizeToZero = 0;
        if (shouldBeZeroed)
        {
            sizeToZero = Min(current->Committed, positionPostPush) - positionPrePush;
        }

        // If needed commit new pages
        if (current->Committed < positionPostPush)
        {
            size_t commitPostAligned = positionPostPush + current->CommitSize - 1;
            commitPostAligned -= commitPostAligned % current->CommitSize;
            size_t commitPostClamped = ClampTop(commitPostAligned, current->Reserved);
            size_t commitSize        = commitPostClamped - current->Committed;
            Byte*  commitPtr         = reinterpret_cast<Byte*>(current) + current->Committed;

            // TODO commit_large
            Memory::OS_Commit(commitPtr, commitSize);

            current->Committed = commitPostClamped;
        }

        // Push on the current block
        void* result = nullptr;
        if (current->Committed >= positionPostPush)
        {
            result            = reinterpret_cast<Byte*>(current) + positionPrePush;
            current->Position = positionPostPush;
            if (sizeToZero != 0)
            {
                MemoryZero(result, sizeToZero);
            }
        }
        
        if (result == nullptr) [[unlikely]]
        {
            Log(LogLevel::Error, LogCategory::Core, "Fatal Allocation Failure - Unexpected allocation failure");
            Assert(false, "Fatal Allocation Failure - Unexpected allocation failure");
        }

        JULIET_DEBUG_ONLY(if (!IsDebugInfoArena(arena)) { DebugArenaAddDebugInfo(current, size, positionPrePush, tag); })

        return result;
    }

    void* ArenaReallocate(NonNullPtr<Arena> arena, void* oldPtr, size_t oldSize, size_t newSize, size_t align,
                          bool shouldBeZeroed JULIET_DEBUG_ONLY(, const char* tag))
    {
        void* result = ArenaPush(arena, newSize, align, shouldBeZeroed JULIET_DEBUG_ONLY(, tag));

        // Find the correct block to release
        Arena* block = nullptr;
        for (block = arena->Current; block != nullptr; block = block->Previous)
        {
            if ((reinterpret_cast<Byte*>(block) < static_cast<Byte*>(oldPtr)) &&
                (static_cast<Byte*>(oldPtr) <= (reinterpret_cast<Byte*>(block) + block->Reserved)))
            {
                break;
            }
        }
        Assert(block != nullptr);

        // Copy old to new
        MemCopy(result, oldPtr, std::min(oldSize, newSize));

        // Zero the old memory
        MemoryZero(oldPtr, oldSize);

        if (oldSize >= sizeof(ArenaFreeNode))
        {
            ArenaFreeNode* freeNode = static_cast<ArenaFreeNode*>(oldPtr);
            ptrdiff_t      posPtr   = static_cast<Byte*>(oldPtr) - reinterpret_cast<Byte*>(block);
            index_t        position = static_cast<index_t>(posPtr);
            freeNode->Position      = position;
            freeNode->Size          = oldSize;

            // Insert at head of the free list
            freeNode->Next     = block->FreeNodes;
            freeNode->Previous = nullptr;
            if (block->FreeNodes)
            {
                block->FreeNodes->Previous = freeNode;
            }
            block->FreeNodes = freeNode;
        
#if JULIET_DEBUG
            // Remove the debug info for the old allocation (since it's now free)
            size_t oldOffset = static_cast<size_t>(static_cast<Byte*>(oldPtr) - reinterpret_cast<Byte*>(block));
            DebugArenaRemoveAllocation(block, oldOffset);
#endif
        }
        return result;
    }

    void ArenaPopTo(NonNullPtr<Arena> arena, size_t position)
    {
        size_t clampedPosition = ClampBottom(k_ArenaHeaderSize, position);
        Arena* current         = arena->Current;

        for (Arena* previous = nullptr; current->BasePosition >= clampedPosition; current = previous)
        {
            previous          = current->Previous;
            current->Position = k_ArenaHeaderSize;

            JULIET_DEBUG_ONLY(DebugArenaFreeBlock(current);)

            SingleLinkedListPushPrevious(arena->FreeBlockLast, current);
        }

        arena->Current     = current;
        size_t newPosition = clampedPosition - current->BasePosition;
        Assert(newPosition <= current->Position);
        current->Position = newPosition;

        JULIET_DEBUG_ONLY(DebugArenaPopTo(current, newPosition);)
    }

    void ArenaPop(NonNullPtr<Arena> arena, size_t amount)
    {
        size_t oldPosition = ArenaPos(arena);
        size_t newPosition = oldPosition;
        if (amount < oldPosition)
        {
            newPosition = oldPosition - amount;
        }
        ArenaPopTo(arena, newPosition);
    }

    void ArenaClear(NonNullPtr<Arena> arena)
    {
        ArenaPopTo(arena, 0);
    }

    size_t ArenaPos(NonNullPtr<Arena> arena)
    {
        Arena* current  = arena->Current;
        size_t position = current->BasePosition + current->Position;
        return position;
    }

    namespace UnitTest
    {
        extern void TestMemoryArena();
    }

    // --- MemoryPool Implementation ---


    // Simple First-Fit Allocator
    MemoryBlock* MemoryPool::AllocateBlock(size_t minCapacity)
    {
        // Require space for Header + Data
        size_t totalUnalignedSize = sizeof(MemoryBlock) + minCapacity;
        size_t requiredSize       = (totalUnalignedSize + 15) & ~static_cast<size_t>(15);

        MemoryBlock** prevPtr = &FreeList;
        MemoryBlock*  curr    = FreeList;

        while (curr)
        {
            if (curr->TotalSize >= requiredSize)
            {
                // Match
                // Check if we can split this block?
                if (curr->TotalSize >= requiredSize + sizeof(MemoryBlock) + 16)
                {
                    // Split
                    size_t remainingSize = curr->TotalSize - requiredSize;

                    MemoryBlock* nextBlock = reinterpret_cast<MemoryBlock*>((uint8*)curr + requiredSize);
                    nextBlock->Magic       = MemoryBlock::kMagic;
                    nextBlock->TotalSize   = remainingSize;
                    nextBlock->Used        = 0;
                    nextBlock->Next        = curr->Next;

                    // Update FreeList to point to the new remaining block instead of curr
                    *prevPtr = nextBlock;

                    // Update curr to be the allocated chunk
                    curr->TotalSize = requiredSize;
                }
                else
                {
                    // Take the whole block
                    *prevPtr = curr->Next;
                }

                curr->Next  = nullptr;
                curr->Used  = 0;
                curr->Magic = MemoryBlock::kMagic;
#if JULIET_DEBUG
                curr->FirstAllocation = nullptr;
#endif

#if JULIET_DEBUG
                if (curr->TotalSize > sizeof(MemoryBlock))
                {
                    MemSet(curr->GetData(), 0xCD, curr->TotalSize - sizeof(MemoryBlock));
                }
#endif
                return curr;
            }

            prevPtr = &curr->Next;
            curr    = curr->Next;
        }

        // Out of Memory in Pool
        Assert(false, "MemoryPool exhausted!");
        return nullptr;
    }

    void MemoryPool::FreeBlock(MemoryBlock* block)
    {
        if (!block)
        {
            return;
        }

        Assert(block->Magic == MemoryBlock::kMagic);

        // Poison Header and Data in Debug
#if JULIET_DEBUG
        DebugFreeArenaAllocations(block);
        // 0xDD = Dead Data
        MemSet(block->GetData(), 0xDD, block->TotalSize - sizeof(MemoryBlock));
        block->Magic = 0xDEADBEEF;
#endif

        // Insert at Head of FreeList (Simplest, no coalescing yet)
        block->Next = FreeList;
        FreeList    = block;
    }

    // --- MemoryArena Implementation ---

    void MemoryArenaCreate(MemoryArena* arena, MemoryPool* pool)
    {
        Assert(arena);
        Assert(pool);
        arena->BackingPool  = pool;
        arena->CurrentBlock = nullptr;
        arena->FirstBlock   = nullptr;
    }

    // Overload for backward compatibility / tests if needed, but we should switch to using Pools.
    // NOTE: The previous signature was (Arena*, void* backing, size_t).
    // We are changing the API.

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

        // Default Block Size (e.g., 64KB or 1MB? Let's use 16KB for granular tests,
        // or larger for prod. Let's make it dynamic or standard constant.
        constexpr size_t kDefaultBlockSize = 64 * 1024; // 64KB pages

        // Alignment check
        Assert((alignment & (alignment - 1)) == 0);

        if (!arena->CurrentBlock)
        {
            // Initial Allocation
            size_t allocSize    = std::max(size, kDefaultBlockSize);
            arena->CurrentBlock = arena->BackingPool->AllocateBlock(allocSize);
            arena->FirstBlock   = arena->CurrentBlock;
        }

        // Try allocation in CurrentBlock
        MemoryBlock* blk             = arena->CurrentBlock;
        size_t       currentAddr     = reinterpret_cast<size_t>(blk->GetData()) + blk->Used;
        size_t       alignmentOffset = 0;

        size_t mask = alignment - 1;
        if (currentAddr & mask)
        {
            alignmentOffset = alignment - (currentAddr & mask);
        }

        if (blk->Used + alignmentOffset + size > blk->TotalSize - sizeof(MemoryBlock))
        {
            // Overflow! Request new block.
            // Strict minimum: what we need now.
            // Better: Max(Default, size) to avoid repeating large allocs for tiny overflow?
            size_t       allocSize = std::max(size, kDefaultBlockSize);
            MemoryBlock* newBlock  = arena->BackingPool->AllocateBlock(allocSize);

            // Link
            blk->Next           = newBlock;
            arena->CurrentBlock = newBlock;
            blk                 = newBlock;

            // Recalc for new block (Used should be 0)
            currentAddr     = reinterpret_cast<size_t>(blk->GetData());
            alignmentOffset = 0;
            // newBlock check
            if (currentAddr & mask)
            {
                alignmentOffset = alignment - (currentAddr & mask);
            }
        }

        // Commit
        blk->Used += alignmentOffset;
        void* ptr = blk->GetData() + blk->Used;

        JULIET_DEBUG_ONLY(DebugArenaAddAllocation(blk, size, blk->Used, tag);)

        blk->Used += size;

        return ptr;
    }

    void* ArenaRealloc(MemoryArena* arena, void* oldPtr, size_t oldSize, size_t newSize, size_t alignment, String tag)
    {
        Assert(arena);
        Assert(oldPtr);
        Assert(newSize != 0);

        // Optimized Case: Expanding the LAST allocation in the Current Block
        MemoryBlock* blk      = arena->CurrentBlock;
        uint8*       oldBytes = static_cast<uint8*>(oldPtr);

        // Is oldPtr inside current block?
        if (oldBytes >= blk->GetData() && oldBytes < blk->GetData() + blk->TotalSize - sizeof(MemoryBlock))
        {
            // Is it the last one?
            if (oldBytes + oldSize == blk->GetData() + blk->Used)
            {
                // Can we expand?
                if (blk->Used + (newSize - oldSize) <= blk->TotalSize - sizeof(MemoryBlock))
                {
                    // Yes, expand in place
                    blk->Used += (newSize - oldSize);
#if JULIET_DEBUG
                    {
                        ArenaAllocation* t = blk->FirstAllocation;
                        while (t && t->Next)
                            t = t->Next;
                        if (t) t->Size += (newSize - oldSize);
                    }
#endif
                    return oldPtr;
                }
            }
        }

        // Fallback: Copy
        void* newPtr = ArenaPush(arena, newSize, alignment, tag);
        MemCopy(newPtr, oldPtr, std::min(oldSize, newSize));
        return newPtr;
    }

    bool ArenaPop(MemoryArena* arena, void* ptr, size_t size)
    {
        Assert(arena);
        Assert(ptr);
        Assert(size);

        MemoryBlock* blk = arena->CurrentBlock;
        Assert(blk);

        uint8* ptrBytes   = static_cast<uint8*>(ptr);
        uint8* currentTop = blk->GetData() + blk->Used;

        // Check if this pointer is exactly at the top of the stack (LIFO)
        if (ptrBytes + size == currentTop)
        {
            // Yes, we can just rewind the Used pointer
            blk->Used -= size;
            JULIET_DEBUG_ONLY(DebugArenaRemoveLastAllocation(blk);)
            return true;
        }

        return false;
    }

    void ArenaReset(MemoryArena* arena)
    {
        Assert(arena);
        Assert(arena->FirstBlock);

        // Keep FirstBlock, Free the rest.
        MemoryBlock* curr = arena->FirstBlock->Next;
        while (curr)
        {
            MemoryBlock* next = curr->Next;
            arena->BackingPool->FreeBlock(curr);
            curr = next;
        }

        arena->FirstBlock->Next = nullptr;
        arena->FirstBlock->Used = 0;
        arena->CurrentBlock     = arena->FirstBlock;

#if JULIET_DEBUG
        // Poison First Block
        DebugFreeArenaAllocations(arena->FirstBlock);
        MemSet(arena->FirstBlock->GetData(), 0xCD, arena->FirstBlock->TotalSize - sizeof(MemoryBlock));
#endif
    }

    ArenaMarker ArenaGetMarker(MemoryArena* arena)
    {
        Assert(arena);
        return { arena->CurrentBlock, arena->CurrentBlock ? arena->CurrentBlock->Used : 0 };
    }

    void ArenaResetToMarker(MemoryArena* arena, ArenaMarker marker)
    {
        Assert(arena);
        if (!marker.Block)
        {
            // If marker block is null, it might mean "start" or "empty".
            // But if the arena has blocks, this is suspicious.
            // If the arena was empty when marker was taken, this is valid.
            ArenaReset(arena);
            return;
        }

        // Free blocks *after* the marker block
        MemoryBlock* curr = marker.Block->Next;
        while (curr)
        {
            MemoryBlock* next = curr->Next;
            arena->BackingPool->FreeBlock(curr);
            curr = next;
        }

        marker.Block->Next  = nullptr;
        marker.Block->Used  = marker.Offset;
        arena->CurrentBlock = marker.Block;
    }

    // --- Global Arenas ---

    namespace
    {
        MemoryPool g_ScratchMemory;
        MemoryPool g_GameMemory;

        MemoryArena g_ScratchArena;
        MemoryArena g_GameArena;

        // Backing Buffers
        void* g_ScratchBuffer = nullptr;
        void* g_EngineBuffer  = nullptr;
        void* g_GameBuffer    = nullptr;

        constexpr size_t kScratchSize = Megabytes(64);
        constexpr size_t kGameSize    = Megabytes(512);

        void InitPool(MemoryPool* pool, void* buffer, size_t size)
        {
            pool->BaseAddress = buffer;
            pool->TotalSize   = size;

            // Create one giant initial block
            Assert(size > sizeof(MemoryBlock));
            MemoryBlock* block = static_cast<MemoryBlock*>(buffer);
            block->Magic       = MemoryBlock::kMagic;
            block->Next        = nullptr;
            block->TotalSize   = size;
            block->Used        = 0;

            pool->FreeList = block;
        }
    } // namespace

    MemoryArena* GetScratchArena()
    {
        return &g_ScratchArena;
    }

    MemoryArena* GetGameArena()
    {
        return &g_GameArena;
    }

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

    void MemoryArenasInit()
    {
        g_ScratchBuffer = Malloc(kScratchSize);
        g_GameBuffer    = Malloc(kGameSize);

        InitPool(&g_ScratchMemory, g_ScratchBuffer, kScratchSize);
        InitPool(&g_GameMemory, g_GameBuffer, kGameSize);

        MemoryArenaCreate(&g_ScratchArena, &g_ScratchMemory);
        MemoryArenaCreate(&g_GameArena, &g_GameMemory);

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

    void MemoryArenasShutdown()
    {
        // Technically we should free blocks?
        // But since we own the giant buffers, we can just free them.
        SafeFree(g_ScratchBuffer);
        SafeFree(g_EngineBuffer);
        SafeFree(g_GameBuffer);
    }

} // namespace Juliet
