#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)
    {
        // Release active blocks (Current chain)
        for (Arena *node = arena->Current, *previous = nullptr; node != nullptr; node = previous)
        {
            previous = node->Previous;

            JULIET_DEBUG_ONLY(DebugUnregisterArena(node);)
            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;
    }

    TempArena ArenaTempBegin(NonNullPtr<Arena> arena)
    {
        index_t position = ArenaPos(arena);
        return { arena.Get(), position };
    }

    void ArenaTempEnd(TempArena temp)
    {
        ArenaPopTo(temp.Arena, temp.Position);
    }
} // namespace Juliet
