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

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

#if defined(AURORA_IS_LINUX_DERIVED)
    #include <sys/mman.h>
#endif

#include <Source/Debug/MemoryCrunch.hpp>

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

    thread_local AuInt64 tlsLastOutOfMemory {};

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

    static void RemoveBytesFromCounter(AuUInt uBytes)
    {
        AuAtomicSub(&gBytesCounterAllocated, uBytes);
    }

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

    #define CHECK_WRAP_RETURN(exp, exp2, string)                                                 \
        auto pRet = exp;                                                                         \
        if (!pRet)                                                                               \
        {                                                                                        \
            tlsLastOutOfMemory = AuTime::CurrentClockNS();                                       \
                                                                                                 \
            bool bCrunchFlag {};                                                                 \
            if (((bCrunchFlag = AuDebug::IsTlsMemoryCrunchActive()) ||                           \
                (gRuntimeConfig.debug.bIsApplicationClientSoftwareOnJitteryMemorySystem)) &&     \
                AuDebug::gReserveHeap)                                                           \
            {                                                                                    \
                if (!(pRet = AuDebug::gReserveHeap-> AU_WHAT exp2))                              \
                {                                                                                \
                    if (bCrunchFlag || gRuntimeConfig.debug.bIsMemoryErrorFatal)                 \
                    {                                                                            \
                        SysPanic("out of contingency memory");                                   \
                    }                                                                            \
                }                                                                                \
            }                                                                                    \
            else if (gRuntimeConfig.debug.bIsMemoryErrorFatal)                                   \
            {                                                                                    \
                SysPanic(string);                                                                \
            }                                                                                    \
        }                                                                                        \
        else                                                                                     \
        {                                                                                        \
            AddBytesToCounter(::mi_malloc_size(pRet));                                           \
        }                                                                                        \
        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)
    {
        SysAssertDbg(length);
        void *pRet;

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

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

        return pRet;
    }

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

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

        if (!pRet)
        {
            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)
    {
        SysAssertDbg(length);
        void *pRet;

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

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

        return pRet;
    }

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

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

        if (!pRet)
        {
            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
        {
            RemoveBytesFromCounter(::mi_malloc_size(pHead));
            ::mi_free(pHead);
        }
    }

    AUKN_SYM AuInt64 GetLastOutOfMemoryTimeNS()
    {
        return tlsLastOutOfMemory;
    }

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