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

    File: AuHeapInternal.cpp
    File: AuHeap.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuHeap.hpp"
#include "AuHeapInternal.hpp"
#include "AuHeapDeletable.hpp"

typedef struct FragmentHeader
{
    void * next;
    void * prev;
    size_t size;
    bool   used;
} FragmentHeader;

namespace Aurora::Memory
{
    InternalHeap::InternalHeap() :
        pHeap_(nullptr),
        uAllocCount_(0)
    { }

    InternalHeap::InternalHeap(const MemoryViewWrite &memory)
    {
        SysAssert(this->Init(memory));
    }

    InternalHeap::InternalHeap(AuUInt uLength)
    {
        SysAssert(this->Init(uLength), "Couldn't initialize inline AuHeap! [OOM]");
    }

    InternalHeap::~InternalHeap()
    {
        SysAssertDbgExp(this->uAllocCount_ == 0);

        if (this->pBase_)
        {
            if (this->pHeap_)
            {
                this->pHeap_ = nullptr;
            }

            if (this->bOwnsMemory_)
            {
                SysAllocateFree(this->pBase_, this->uLength_);
                this->pBase_ = nullptr;
            }
        }
    }

    AuUInt InternalHeap::GetHeapSize(const void *ptr)
    {
        return reinterpret_cast<const FragmentHeader *>(ptr)[-1].size;
    }

    bool InternalHeap::Init(const MemoryViewWrite &memory)
    {
        if (memory.HasControlBlock())
        {
            this->pBloat_ = AuMakeShared<MemoryViewWrite>(memory);
        }

        return this->Init(memory.length, memory.ptr);
    }

    bool InternalHeap::Init(AuUInt uLength, void *ptr)
    {
        SysAssert(!this->pBase_, "heap already initialized");
        SysAssert(uLength, "invalid heap allocation");

        if (ptr)
        {
            this->pBase_       = ptr;
            this->uLength_     = uLength;
            this->bOwnsMemory_ = false;
        }
        else
        {
            if (uLength <= 4096)
            {
                uLength = 4096;
            }

            if (!(this->pBase_ = SysAllocateLarge(uLength)))
            {
                return false;
            }

            this->bOwnsMemory_ = true;
            this->uLength_     = uLength;
        }

        if (!(this->pHeap_ = o1heapInit(this->pBase_, uLength)))
        {
            return false;
        }

        return true;
    }

    Types::size_t InternalHeap::GetChunkSize(const void *head)
    {
        return InternalHeap::GetHeapSize(head);
    }

    AuSPtr<Heap> InternalHeap::AllocateDivision(AuUInt32 heap, AuUInt32 alignment)
    {
        return AllocateDivisionGlobal(this, heap, alignment);
    }

    AuSPtr<Heap> AllocateDivisionGlobal(Heap *heap, AuUInt32 uLength, AuUInt32 alignment)
    {
        auto ptr = heap->ZAlloc<void *>(uLength, alignment);
        if (!ptr)
        {
            return {};
        }

        auto ret = AuMakeShared<DeletableHeap>(heap, ptr);
        if (!ret)
        {
            heap->Free(ptr);
            return {};
        }

        if (!ret->Init(uLength, ptr))
        {
            return {};
        }

        return ret;
    }

    void *InternalHeap::_FAlloc(Types::size_t uLength)
    {
        if (!this->pHeap_)
        {
            return nullptr;
        }

        auto ret = o1heapAllocate(this->pHeap_, uLength);
        if (ret)
        {
            AuAtomicAdd(&this->uAllocCount_, 1u);
        }
        return ret;
    }

    void *InternalHeap::_FAlloc(Types::size_t uLength, Types::size_t uAlign)
    {
        SysAssert(uAlign <= O1HEAP_ALIGNMENT, "heap wrapping is unsupported, alignment past the supported 2^x alignment is not possible");
        return this->_FAlloc(uLength);
    }

    void *InternalHeap::_ZAlloc(Types::size_t uLength)
    {
        if (!this->pHeap_)
        {
            return nullptr;
        }

        auto ptr = this->_FAlloc(uLength);
        if (!ptr)
        {
            return nullptr;
        }

        AuMemset(ptr, 0, uLength);
        return ptr;
    }

    void *InternalHeap::_ZAlloc(Types::size_t uLength, Types::size_t uAlign)
    {
        SysAssert(uAlign <= O1HEAP_ALIGNMENT, "heap wrapping is unsupported, alignment past the supported 2^x alignment is not possible");
        return _ZAlloc(uLength);
    }

    void *InternalHeap::_ZRealloc(void *pBuffer, Types::size_t uLength)
    {
        auto prevLength = GetHeapSize(pBuffer);
        auto alloc = this->_ZAlloc(uLength);
        if (!alloc)
        {
            return nullptr;
        }

        AuMemcpy(alloc, pBuffer, AuMin(prevLength, uLength));
        this->_Free(pBuffer);
        return alloc;
    }

    void *InternalHeap::_ZRealloc(void *pBuffer, Types::size_t uLength, Types::size_t uAlign)
    {
        SysAssert(uAlign <= O1HEAP_ALIGNMENT, "heap wrapping is unsupported, alignment past the supported 2^x alignment is not possible");
        return this->_ZRealloc(pBuffer, uLength);
    }

    void *InternalHeap::_FRealloc(void *pBuffer, Types::size_t uLength)
    {
        auto prevLength = GetHeapSize(pBuffer);
        auto alloc = this->_FAlloc(uLength);
        if (!alloc)
        {
            return nullptr;
        }

        AuMemcpy(alloc, pBuffer, AuMin(prevLength, uLength));
        this->_Free(pBuffer);
        return alloc;
    }

    void *InternalHeap::_FRealloc(void *pBuffer, Types::size_t uLength, Types::size_t uAlign)
    {
        SysAssert(uAlign <= O1HEAP_ALIGNMENT, "heap wrapping is unsupported, alignment past the supported 2^x alignment is not possible");
        return this->_FRealloc(pBuffer, uLength);
    }

    void   InternalHeap::_Free(void *pBuffer)
    {
        if (pBuffer == nullptr)
        {
            return;
        }

        o1heapFree(this->pHeap_, pBuffer);
        DecrementUsers();
    }

    void InternalHeap::DecrementUsers()
    {
        if (AuAtomicSub(&this->uAllocCount_, 1u) == 0)
        {
            TryRelease();
        }
    }

    void InternalHeap::TryRelease()
    {
        if (!this->bIsDangling_)
        {
            return;
        }

        if (AuAtomicLoad(&uAllocCount_) == 0)
        {
            delete this;
        }
    }

    void InternalHeap::RequestTermination()
    {
        if (AuAtomicLoad(&this->uAllocCount_))
        {
            SysPushErrorMemory("Heap life was less than its allocations, waiting for final free");
            SysPushErrorMemory("Reporting using mayday!");

            // Write a crash dump for later review, and do not panic.
            // We just have a leak with no sign of corruption 
            Telemetry::Mayday();

            this->bIsDangling_ = true;
        }
        else
        {
            delete this;
        }
    }

    void InternalHeap::UpdateStats()
    {
        auto pDiag = o1heapGetDiagnostics(this->pHeap_);

        this->stats.uBytesLiveCounter = pDiag.allocated;
        this->stats.uBytesCapacity    = pDiag.capacity;
        this->stats.uBytesPeakCounter = pDiag.peak_allocated;
    }

    void InternalHeap::WalkHeap(bool(*fCallback)(void *, void *), void *pSecondArg)
    {
        o1heapTraverseHeap(this->pHeap_, fCallback, pSecondArg);
    }

    AuSPtr<Heap> InternalHeap::GetSelfReference()
    {
        try
        {
            return AuSharedFromThis();
        }
        catch (...)
        {
            return {};
        }
    }

    Heap *InternalHeap::GetSelfReferenceRaw()
    {
        return this;
    }

    AUKN_SYM Heap *AllocHeapNew(AuUInt uSize)
    {
        auto pHeap = _new InternalHeap();
        if (!pHeap)
        {
            return nullptr;
        }

        if (!pHeap->Init(uSize, nullptr))
        {
            delete pHeap;
            return nullptr;
        }

        return pHeap;
    }

    AUKN_SYM void AllocHeapRelease(Heap *pHeap)
    {
        static_cast<InternalHeap *>(pHeap)->RequestTermination();
    }

    AUKN_SYM Heap *RequestHeapOfRegionNew(const MemoryViewWrite &memory)
    {
        if (!memory)
        {
            SysPushErrorArg();
            return nullptr;
        }

        auto pHeap = _new InternalHeap();
        if (!pHeap)
        {
            return nullptr;
        }

        if (!pHeap->Init(memory))
        {
            delete pHeap;
            return nullptr;
        }

        return pHeap;
    }

    AUKN_SYM void RequestHeapOfRegionRelease(Heap *pHeap)
    {
        static_cast<InternalHeap *>(pHeap)->RequestTermination();
    }

    AUROXTL_INTERFACE_SOO_SRC_EX(AURORA_SYMBOL_EXPORT, RequestHeapOfRegion, InternalHeap, (const MemoryViewWrite &, memory))
    AUROXTL_INTERFACE_SOO_SRC_EX(AURORA_SYMBOL_EXPORT, AllocHeap, InternalHeap, (AuUInt, uLength))
}