/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: AuMemory.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuMemory.hpp"
#include <mimalloc.h>
#include <Source/HWInfo/AuHWInfo.hpp>
#include <Source/Debug/MemoryCrunch.hpp>
#include <Source/Debug/ErrorStack.hpp>

namespace Aurora::Memory
{
    AuUInt gBytesCounterAllocated {};
    AuUInt gBytesCounterPeak {};

    thread_local AuInt64 tlsLastOutOfMemory {};
    thread_local MemoryLowNotification_f tlsMemoryLowNotification;
    
    static LeakFinderAlloc_f gLeakFinderAlloc;
    static LeakFinderFree_f  gLeakFinderFree;

    static void AddBytesToCounter(AuUInt uBytes)
    {
        gBytesCounterPeak = AuMax(gBytesCounterPeak, AuAtomicAdd(&gBytesCounterAllocated, uBytes));
    }

    static void RemoveBytesFromCounter(AuUInt uBytes)
    {
        AuAtomicSub(&gBytesCounterAllocated, uBytes);
    }
    
    AUKN_SYM void SetMemoryLowNotification(MemoryLowNotification_f pFunc)
    {
        tlsMemoryLowNotification = pFunc;
    }

    AUKN_SYM void SetLeakFinder(LeakFinderAlloc_f pAlloc,
                                LeakFinderFree_f pFree)
    {
        gLeakFinderAlloc = pAlloc;
        gLeakFinderFree = pFree;
    }

    AUKN_SYM void ReserveHeapMemory(AuUInt uHeapSize, bool bCommit)
    {
        mi_reserve_os_memory(uHeapSize, bCommit, true);
    }

    AUKN_SYM AuUInt GetChunkSize(const void *head)
    {
        if (AuDebug::IsPointerReserveRange((void *)head))
        {
            return AuDebug::gReserveHeap->GetChunkSize(head);
        }
        else
        {
            return ::mi_malloc_size(head);
        }
    }

    static void OnOOM(AuUInt uLength)
    {
        tlsLastOutOfMemory = AuTime::CurrentClockNS();

        AuDebug::OnOutOfMemory();

        if (auto pLowNotification = tlsMemoryLowNotification)
        {
            pLowNotification(uLength, 4);
        }
    }

    #define CHECK_WRAP_RETURN(exp, exp2, string)                                                 \
        auto pRet = exp;                                                                         \
        if (!pRet)                                                                               \
        {                                                                                        \
            OnOOM(length);                                                                       \
                                                                                                 \
            bool bCrunchFlag {};                                                                 \
            if (((bCrunchFlag = AuDebug::IsTlsMemoryCrunchActive()) ||                           \
                (gRuntimeConfig.debug.bIsApplicationClientSoftwareOnJitteryMemorySystem)) &&     \
                AuDebug::gReserveHeap)                                                           \
            {                                                                                    \
                if (auto pLowNotification = tlsMemoryLowNotification)                            \
                {                                                                                \
                    pLowNotification(length, 2);                                                 \
                }                                                                                \
                if (!(pRet = AuDebug::gReserveHeap-> AU_WHAT exp2))                              \
                {                                                                                \
                    if (bCrunchFlag || gRuntimeConfig.debug.bIsMemoryErrorFatal)                 \
                    {                                                                            \
                        SysPanic("out of contingency memory");                                   \
                    }                                                                            \
                }                                                                                \
            }                                                                                    \
            else if (gRuntimeConfig.debug.bIsMemoryErrorFatal)                                   \
            {                                                                                    \
                SysPanic(string);                                                                \
            }                                                                                    \
            if (auto pLowNotification = tlsMemoryLowNotification)                                \
            {                                                                                    \
                pLowNotification(length, !pRet ? 0 : 1);                                         \
            }                                                                                    \
        }                                                                                        \
        else                                                                                     \
        {                                                                                        \
            auto uCount = ::mi_malloc_size(pRet);                                                \
            if (gLeakFinderAlloc)                                                                \
            {                                                                                    \
                gLeakFinderAlloc(pRet, uCount);                                                  \
            }                                                                                    \
            AddBytesToCounter(uCount);                                                           \
        }                                                                                        \
        return pRet;

    AUKN_SYM void *_ZAlloc(Types::size_t length)
    {
        SysAssertDbg(length);
        CHECK_WRAP_RETURN(::mi_zalloc(length), (ZAlloc(length)), "ZAlloc out of memory");
    }

    AUKN_SYM void *_ZAlloc(Types::size_t length, Types::size_t align)
    {
        SysAssertDbg(length);
        CHECK_WRAP_RETURN(::mi_zalloc_aligned(length, align), (ZAlloc(length, align)), "ZAlloc out of memory");
    }

    AUKN_SYM void *_FAlloc(Types::size_t length)
    {
        SysAssertDbg(length);
        CHECK_WRAP_RETURN(::mi_malloc(length), (FAlloc(length)), "FAlloc out of memory");
    }

    AUKN_SYM void *_FAlloc(Types::size_t length, Types::size_t align)
    {
        SysAssertDbg(length);
        CHECK_WRAP_RETURN(::mi_malloc_aligned(length, align), (FAlloc(length, align)), "FAllocEx out of memory");
    }

    AUKN_SYM void *_ZRealloc(void *buffer, Types::size_t length, Types::size_t align)
    {
        void *pRet = nullptr;

        if (AuDebug::IsPointerReserveRange(buffer))
        {
            pRet = AuDebug::gReserveHeap->ZRealloc(buffer, length, align);
        }
        else
        {
            auto oldLen = ::mi_malloc_size(buffer);

            if (buffer && !length)
            {
                ::mi_free(buffer);
            }
            else if (!buffer && length)
            {
                return _ZAlloc(length, align);
            }
            else if ((pRet = ::mi_rezalloc_aligned(buffer, length, align)))
            {
                auto uNewSize = ::mi_malloc_size(pRet);

                RemoveBytesFromCounter(oldLen);
                AddBytesToCounter(uNewSize);
                
                if (gLeakFinderAlloc)
                {
                    if (gLeakFinderFree)
                    {
                        gLeakFinderFree(buffer, oldLen);
                    }

                    gLeakFinderAlloc(pRet, uNewSize);
                }
            }
        }

        if (!pRet && buffer && length)
        {
            OnOOM(length);

            if (gRuntimeConfig.debug.bIsMemoryErrorFatal)
            {
                SysPanic("ZAllocEx out of memory");
            }
        }

        return pRet;
    }

    AUKN_SYM void *_ZRealloc(void *buffer, Types::size_t length)
    {
        void *pRet = nullptr;

        if (AuDebug::IsPointerReserveRange(buffer))
        {
            pRet = AuDebug::gReserveHeap->ZRealloc(buffer, length);
        }
        else
        {
            auto oldLen = ::mi_malloc_size(buffer);

            if (buffer && !length)
            {
                ::mi_free(buffer);
            }
            else if (!buffer && length)
            {
                return _ZAlloc(length);
            }
            else if ((pRet = ::mi_rezalloc(buffer, length)))
            {
                auto uNewSize = ::mi_malloc_size(pRet);

                RemoveBytesFromCounter(oldLen);
                AddBytesToCounter(uNewSize);
                
                if (gLeakFinderAlloc)
                {
                    if (gLeakFinderFree)
                    {
                        gLeakFinderFree(buffer, oldLen);
                    }

                    gLeakFinderAlloc(pRet, uNewSize);
                }
            }
        }

        if (!pRet && buffer && length)
        {
            OnOOM(length);

            if (gRuntimeConfig.debug.bIsMemoryErrorFatal)
            {
                SysPanic("ZAlloc out of memory");
            }
        }

        return pRet;
    }

    AUKN_SYM void *_FRealloc(void *buffer, Types::size_t length, Types::size_t align)
    {
        void *pRet = nullptr;

        if (AuDebug::IsPointerReserveRange(buffer))
        {
            pRet = AuDebug::gReserveHeap->FRealloc(buffer, length, align);
        }
        else
        {
            auto oldLen = ::mi_malloc_size(buffer);

            if (buffer && !length)
            {
                ::mi_free(buffer);
            }
            else if (!buffer && length)
            {
                return _FAlloc(length, align);
            }
            else if ((pRet = ::mi_realloc_aligned(buffer, length, align)))
            {
                auto uNewSize = ::mi_malloc_size(pRet);

                RemoveBytesFromCounter(oldLen);
                AddBytesToCounter(uNewSize);
                
                if (gLeakFinderAlloc)
                {
                    if (gLeakFinderFree)
                    {
                        gLeakFinderFree(buffer, oldLen);
                    }

                    gLeakFinderAlloc(pRet, uNewSize);
                }
            }
        }

        if (!pRet && buffer && length)
        {
            OnOOM(length);

            if (gRuntimeConfig.debug.bIsMemoryErrorFatal)
            {
                SysPanic("FReallocEx out of memory");
            }
        }

        return pRet;
    }

    AUKN_SYM void *_FRealloc(void *buffer, Types::size_t length)
    {
        void *pRet = nullptr;

        if (AuDebug::IsPointerReserveRange(buffer))
        {
            pRet = AuDebug::gReserveHeap->FRealloc(buffer, length);
        }
        else
        {
            auto oldLen = ::mi_malloc_size(buffer);

            if (buffer && !length)
            {
                ::mi_free(buffer);
            }
            else if (!buffer && length)
            {
                return _FAlloc(length);
            }
            else if ((pRet = ::mi_realloc(buffer, length)))
            {
                auto uNewSize = ::mi_malloc_size(pRet);

                RemoveBytesFromCounter(oldLen);
                AddBytesToCounter(uNewSize);
                
                if (gLeakFinderAlloc)
                {
                    if (gLeakFinderFree)
                    {
                        gLeakFinderFree(buffer, oldLen);
                    }

                    gLeakFinderAlloc(pRet, uNewSize);
                }
            }
        }

        if (!pRet && buffer && length)
        {
            OnOOM(length);

            if (gRuntimeConfig.debug.bIsMemoryErrorFatal)
            {
                SysPanic("FRealloc out of memory");
            }
        }

        return pRet;
    }

    AUKN_SYM void _Free(void *pHead)
    {
        if (!pHead)
        {
            return;
        }

        if (AuDebug::IsPointerReserveRange(pHead))
        {
            AuDebug::gReserveHeap->Free(pHead);
        }
        else
        {
            auto uCount = ::mi_malloc_size(pHead);

            RemoveBytesFromCounter(uCount);
            ::mi_free(pHead);

            if (gLeakFinderFree)
            {
                gLeakFinderFree(pHead, uCount);
            }
        }
    }

    AUKN_SYM AuInt64 GetLastOutOfMemoryTimeNS()
    {
        return tlsLastOutOfMemory;
    }

    AUKN_SYM AuUInt GetPageSize()
    {
        return AuHwInfo::GetPageSize();
    }
}