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

    File: Heap.cpp
    Date: 2021-6-12
    Author: Reece
***/
#include <RuntimeInternal.hpp>
#include "Memory.hpp"
#include "Heap.hpp"

#include "mimalloc.h"
#include "o1heap.hpp"

namespace Aurora::Memory
{
    class InternalHeap : public Heap
    {
    public:
        InternalHeap() : base_(nullptr), mutex_(nullptr), heap_(nullptr), _count(0)
        {
        }

        ~InternalHeap();

        bool Init(AuUInt length);

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

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

        Types::size_t GetChunkSize(const void *head) override;

        void *_FAlloc(Types::size_t length) override;
        void *_FAlloc(Types::size_t length, Types::size_t align) override;
        void *_ZAlloc(Types::size_t length) override;
        void *_ZAlloc(Types::size_t length, Types::size_t align) override;
        void *_ZRealloc(void *buffer, Types::size_t length) override;
        void *_ZRealloc(void *buffer, Types::size_t length, Types::size_t align) override;
        void *_FRealloc(void *buffer, Types::size_t length) override;
        void *_FRealloc(void *buffer, Types::size_t length, Types::size_t align) override;
        void  _Free(void *buffer) override;
        
        void TryRelease();
        void RequestTermination();

    private:
        Threading::Primitives::MutexUnique_t mutex_;
        void *base_{};
        O1HeapInstance *heap_{};
        int _count{};
        bool _isDangling{};
    };
    
    InternalHeap::~InternalHeap()
    {
        SysAssertDbgExp(_count == 0);

        mutex_.reset();

        if (base_)
        {
            o1HeapReleaseCpp(heap_);// ->~O1HeapInstance(); // TODO: make free func

            Memory::Free(base_);
            base_ = nullptr;
            mi_collect(false);
        }
    }

    bool InternalHeap::Init(AuUInt length)
    {
        SysAssert(!base_, "heap already initialized");
        SysAssert(!mutex_, "heap already initialized");

        SysAssert(length, "invalid heap allocation");

        mutex_ = Threading::Primitives::MutexUnique();
        if (!mutex_) return false;

        base_ = Memory::FAlloc<void *>(length);
        if (!base_) return false;

        heap_ = o1heapInit(base_, length,
                           [this](const O1HeapInstance *const handle) -> void
        {
            SysAssertDbg(this->mutex_ ? true : false, "missing mutex");
            this->mutex_->Lock();
        },
                           [this](const O1HeapInstance *const handle) -> void
        {
            SysAssertDbg(this->mutex_ ? true : false, "missing mutex");
            this->mutex_->Unlock();
        });
        return true;
    }

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

    void *InternalHeap::_FAlloc(Types::size_t length)
    {
        if (!heap_) return nullptr;
        auto ret = o1heapAllocate(heap_, length);
        if (ret) _count++;
        return ret;
    }

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

    void *InternalHeap::_ZAlloc(Types::size_t length)
    {
        if (!heap_) return nullptr;
        auto ptr = _FAlloc(length);
        if (!ptr) return nullptr;
        std::memset(ptr, 0, length);
        return ptr;
    }

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

    void *InternalHeap::_ZRealloc(void *buffer, Types::size_t length)
    {
        auto prevLength = GetHeapSize(buffer);
        auto alloc = _ZAlloc(length);
        if (!alloc) return nullptr;
        std::memcpy(alloc, buffer, std::min(prevLength, length));
        _Free(buffer);
        return alloc;
    }

    void *InternalHeap::_ZRealloc(void *buffer, Types::size_t length, Types::size_t align)
    {
        SysAssert(align < O1HEAP_ALIGNMENT, "heap wrapping is unsupported, alignment past the supported 2^x alignment is not possible");
        return _ZRealloc(buffer, length);
    }

    void *InternalHeap::_FRealloc(void *buffer, Types::size_t length)
    {
        auto prevLength = GetHeapSize(buffer);
        auto alloc = _FAlloc(length);
        if (!alloc) return nullptr;
        std::memcpy(alloc, buffer, std::min(prevLength, length));
        _Free(buffer);
        return alloc;
    }

    void *InternalHeap::_FRealloc(void *buffer, Types::size_t length, Types::size_t align)
    {
        SysAssert(align < O1HEAP_ALIGNMENT, "heap wrapping is unsupported, alignment past the supported 2^x alignment is not possible");
        return _FRealloc(buffer, length);
    }

    void   InternalHeap::_Free(void *buffer)
    {
        if (buffer == nullptr) return;
        o1heapFree(heap_, buffer);
        _count--;
    }

    void InternalHeap::TryRelease()
    {
        if (!_isDangling) return;

        if (_count == 0)
        {
            delete this;
            return;
        }
    }

    void InternalHeap::RequestTermination()
    {
        if (_count)
        {
            LogWarn("Heap life was less than its allocations, waiting for final free");
            LogWarn("Reporting using mayday!");
            Telemetry::Mayday();

            _isDangling = true;
            TryRelease();
        }
        else
        {
            delete this;
        }
    }

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

        if (!heap->Init(size))
        {
            delete heap;
            return nullptr;
        }

        return heap;
    }

    AUKN_SYM void AllocHeapRelease(Heap *heap) 
    {
        dynamic_cast<InternalHeap *>(heap)->RequestTermination();
    }
}