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

    File: HeapProxy.cpp
    Date: 2024-1-16
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "Memory.hpp"
#include "Heap.hpp"
#include "HeapProxy.hpp"

namespace Aurora::Memory
{
    ProxyHeap::ProxyHeap(std::shared_ptr<Heap> pHeap,
                  LeakFinderAlloc_f pAlloc,
                  LeakFinderFree_f pFree) :
        pHeap(pHeap),
        pAlloc(pAlloc),
        pFree(pFree)
    {

    }

    void ProxyHeap::UpdateStats()
    {
        auto &stats = this->pHeap->GetStats();

        this->stats.uBytesLiveCounter        = this->uBytesAllocated;
        this->stats.qwBytesAllocatedLifetime = this->uBytesLifetime;
        this->stats.uBytesPeakCounter        = this->uBytesPeak;
        this->stats.qwBytesFreeLifetime      = this->uBytesFree;
        this->stats.uBytesCapacity           = stats.uBytesCapacity;
    }

    AuSPtr<Heap> ProxyHeap::AllocateDivision(AuUInt32 heap, AuUInt32 alignment)
    { 
        return this->pHeap->AllocateDivision(heap, alignment);
    }

    void *ProxyHeap::_ZAlloc(Types::size_t uLength)
    { 
        if (auto pThat = this->pHeap->ZAlloc(uLength))
        {
            auto uLengthCurrent = this->GetChunkSize(pThat);
            AuAtomicAdd(&this->uBytesLifetime, uLengthCurrent);
            AuAtomicAdd(&this->uBytesAllocated, uLengthCurrent);
            this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
            if (this->pAlloc)
            {
                this->pAlloc(pThat, uLength);
            }
            return pThat;
        }
        else
        {
            return {};
        }
    }

    void *ProxyHeap::_ZAlloc(Types::size_t uLength, Types::size_t align)
    {
        if (auto pThat = this->pHeap->ZAlloc(uLength, align))
        {
            auto uLengthCurrent = this->GetChunkSize(pThat);
            AuAtomicAdd(&this->uBytesLifetime, uLengthCurrent);
            AuAtomicAdd(&this->uBytesAllocated, uLengthCurrent);
            this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
            if (this->pAlloc)
            {
                this->pAlloc(pThat, uLength);
            }
            return pThat;
        }
        else
        {
            return {};
        }
    }

    Types::size_t ProxyHeap::GetChunkSize(const void *head)
    {
        return this->pHeap->GetChunkSize(head);
    }

    void *ProxyHeap::_FAlloc(Types::size_t uLength)
    {
        if (auto pThat = this->pHeap->_FAlloc(uLength))
        {
            auto uLengthCurrent = this->GetChunkSize(pThat);
            AuAtomicAdd(&this->uBytesLifetime, uLengthCurrent);
            AuAtomicAdd(&this->uBytesAllocated, uLengthCurrent);
            this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
            if (this->pAlloc)
            {
                this->pAlloc(pThat, uLength);
            }
            return pThat;
        }
        else
        {
            return {};
        }
    }

    void *ProxyHeap::_FAlloc(Types::size_t uLength, Types::size_t align)
    {
        if (auto pThat = this->pHeap->_FAlloc(uLength, align))
        {
            auto uLengthCurrent = this->GetChunkSize(pThat);
            AuAtomicAdd(&this->uBytesLifetime, uLengthCurrent);
            AuAtomicAdd(&this->uBytesAllocated, uLengthCurrent);
            this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
            if (this->pAlloc)
            {
                this->pAlloc(pThat, uLength);
            }
            return pThat;
        }
        else
        {
            return {};
        }
    }

    void *ProxyHeap::_ZRealloc(void *pHead, Types::size_t uLength, Types::size_t align)
    {
        if (pHead)
        {
            auto uLengthCurrent = this->GetChunkSize(pHead);

            if (auto pThat = this->pHeap->_ZRealloc(pHead, uLength, align))
            {
                auto uLengthNext = this->GetChunkSize(pThat);
                AuAtomicAdd(&this->uBytesFree, uLengthCurrent);
                AuAtomicAdd(&this->uBytesLifetime, uLengthNext);
                AuAtomicAdd(&this->uBytesAllocated, decltype(this->uBytesAllocated)(uLengthNext) - decltype(this->uBytesAllocated)(uLengthCurrent));
                this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
                if (this->pAlloc)
                {
                    this->pFree(pHead, uLengthCurrent);
                    this->pAlloc(pThat, uLength);
                }
                return pThat;
            }
            else
            {
                return {};
            }
        }
        else
        {
            return nullptr;
        }
    }

    void *ProxyHeap::_ZRealloc(void *pHead, Types::size_t uLength)
    {
        if (pHead)
        {
            auto uLengthCurrent = this->GetChunkSize(pHead);

            if (auto pThat = this->pHeap->_ZRealloc(pHead, uLength))
            {
                auto uLengthNext = this->GetChunkSize(pThat);
                AuAtomicAdd(&this->uBytesFree, uLengthCurrent);
                AuAtomicAdd(&this->uBytesLifetime, uLengthNext);
                AuAtomicAdd(&this->uBytesAllocated, decltype(this->uBytesAllocated)(uLengthNext) - decltype(this->uBytesAllocated)(uLengthCurrent));
                this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
                if (this->pAlloc)
                {
                    this->pFree(pHead, uLengthCurrent);
                    this->pAlloc(pThat, uLength);
                }
                return pThat;
            }
            else
            {
                return {};
            }
        }
        else
        {
            return nullptr;
        }
    }

    void *ProxyHeap::_FRealloc(void *pHead, Types::size_t uLength, Types::size_t align)
    {
        if (pHead)
        {
            auto uLengthCurrent = this->GetChunkSize(pHead);

            if (auto pThat = this->pHeap->_FRealloc(pHead, uLength, align))
            {
                auto uLengthNext = this->GetChunkSize(pThat);
                AuAtomicAdd(&this->uBytesFree, uLengthCurrent);
                AuAtomicAdd(&this->uBytesLifetime, uLengthNext);
                AuAtomicAdd(&this->uBytesAllocated, decltype(this->uBytesAllocated)(uLengthNext) - decltype(this->uBytesAllocated)(uLengthCurrent));
                this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
                if (this->pAlloc)
                {
                    this->pFree(pHead, uLengthCurrent);
                    this->pAlloc(pThat, uLength);
                }
                return pThat;
            }
            else
            {
                return {};
            }
        }
        else
        {
            return nullptr;
        }
    }

    void *ProxyHeap::_FRealloc(void *pHead, Types::size_t uLength)
    {
        if (pHead)
        {
            auto uLengthCurrent = this->GetChunkSize(pHead);

            if (auto pThat = this->pHeap->_FRealloc(pHead, uLength))
            {
                auto uLengthNext = this->GetChunkSize(pThat);
                AuAtomicAdd(&this->uBytesFree, uLengthCurrent);
                AuAtomicAdd(&this->uBytesLifetime, uLengthNext);
                AuAtomicAdd(&this->uBytesAllocated, decltype(this->uBytesAllocated)(uLengthNext) - decltype(this->uBytesAllocated)(uLengthCurrent));
                this->uBytesPeak = AuMax(this->uBytesPeak, this->uBytesAllocated);
                if (this->pAlloc)
                {
                    this->pFree(pHead, uLengthCurrent);
                    this->pAlloc(pThat, uLength);
                }
                return pThat;
            }
            else
            {
                return {};
            }
        }
        else
        {
            return nullptr;
        }
    }

    void  ProxyHeap::_Free(void *pHead)
    {
        if (pHead)
        {
            auto uLengthCurrent = this->GetChunkSize(pHead);
            AuAtomicAdd(&this->uBytesFree, uLengthCurrent);
            AuAtomicSub(&this->uBytesAllocated, uLengthCurrent);
            if (this->pAlloc)
            {
                this->pFree(pHead, uLengthCurrent);
            }
        }
    }

    AuSPtr<Heap> ProxyHeap::GetSelfReference()
    {
        return this->pHeap;
    }

    Heap *ProxyHeap::GetSelfReferenceRaw()
    {
    #if defined(AU_CFG_ID_SHIP) || defined(AU_FORCE_SAFE_PROXY_SHARED_HEAP_FREE)
        return this->pHeap.get();
    #else
        return this;
    #endif
    }

    AUKN_SYM Heap *HeapProxyNew(const AuSPtr<Heap> &pHeap)
    {
        if (!pHeap)
        {
            SysPushErrorArg();
            return {};
        }

        auto pRetHeap = _new ProxyHeap(pHeap);
        if (!pRetHeap)
        {
            SysPushErrorMemory();
            return nullptr;
        }

        return pRetHeap;
    }

    AUKN_SYM void HeapProxyRelease(Heap *heap)
    {
        AuSafeDelete<ProxyHeap *>(heap);
    }

    AUKN_SYM Heap *HeapProxyExNew(const AuSPtr<Heap> &pHeap,
                                  LeakFinderAlloc_f pAlloc,
                                  LeakFinderFree_f pFree)
    {
        if (!pHeap ||
            !pAlloc ||
            !pFree)
        {
            SysPushErrorArg();
            return {};
        }

        auto pRetHeap = _new ProxyHeap(pHeap, pAlloc, pFree);
        if (!pRetHeap)
        {
            SysPushErrorMemory();
            return nullptr;
        }

        return pRetHeap;
    }

    AUKN_SYM void HeapProxyExRelease(Heap *heap)
    {
        AuSafeDelete<ProxyHeap *>(heap);
    }
}