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

    File: Heap.hpp
    Date: 2021-6-9
    Author: Reece
***/
#pragma once

namespace Aurora::Memory
{
    struct Heap
    {
        virtual AuSPtr<Heap> AllocateDivision(AuUInt32 heap, AuUInt32 alignment = 32) = 0;
        virtual Types::size_t GetChunkSize(const void *head) = 0;
        virtual HeapStats &GetStats() = 0;

        template<typename T = void *>
        T ZAlloc(Types::size_t length)
        {
            return reinterpret_cast<T>(_ZAlloc(length));
        }

        template<typename T = void *>
        T ZAlloc(Types::size_t length, Types::size_t align)
        {
            return reinterpret_cast<T>(_ZAlloc(length, align));
        }

        template<typename T>
        T *ZAlloc()
        {
            return reinterpret_cast<T *>(_ZAlloc(sizeof(T)));
        }

        template<typename T>
        T *NewArray(Types::size_t count)
        {
            return ZAlloc<T *>(count * sizeof(T));
        }

        template<typename T>
        T *NewArray(Types::size_t count, Types::size_t align)
        {
            return ZAlloc<T *>(count * sizeof(T), align);
        }

        /// Fast, unsafe alloc
        template<typename T = void *>
        T FAlloc(Types::size_t length)
        {
            return reinterpret_cast<T>(_FAlloc(length));
        }

        template<typename T = void *>
        T FAlloc(Types::size_t length, Types::size_t align)
        {
            return reinterpret_cast<T>(_FAlloc(length, align));
        }

        template<typename T>
        T ZRealloc(T in, Types::size_t length)
        {
            return reinterpret_cast<T>(_ZRealloc(reinterpret_cast<void *>(in), length));
        }

        template<typename T>
        T ZRealloc(T in, Types::size_t length, Types::size_t alloc)
        {
            return reinterpret_cast<T>(_ZRealloc(reinterpret_cast<void *>(in), length), alloc);
        }

        template<typename T>
        T FRealloc(T in, Types::size_t length)
        {
            return reinterpret_cast<T>(_FRealloc(reinterpret_cast<void *>(in), length));
        }

        template<typename T>
        T FRealloc(T in, Types::size_t length, Types::size_t alloc)
        {
            return reinterpret_cast<T>(_FRealloc(reinterpret_cast<void *>(in), length), alloc);
        }

        template<typename T>
        void Free(T in)
        {
            _Free(reinterpret_cast<void *>(in));
        }

        template <class T, class ...Args>
        AuSPtr<T> NewClass(Args &&...args)
        {
            void *pPtr;

            if constexpr (AuIsClass_v<T>)
            {
                pPtr = this->FAlloc<void *>(sizeof(T));
                if (pPtr)
                {
                    new (pPtr) T(AuForward<Args &&>(args)...);
                }
            }
            else
            {
                pPtr = this->ZAlloc<void *>(sizeof(T));
            }

            if (!pPtr)
            {
                return {};
            }

            return ToSmartPointer((T *)pPtr, true);
        }

        template<typename T>
        AuSPtr<T> ToSmartPointer(T *in, bool bPinThis)
        {
            if (in == nullptr) return {};
            auto heapHandle = bPinThis ? GetSelfReference() : AuSPtr<Heap> {};
            return AuSPtr<T>(in,
                   [heapHandle, in, this](T *delt)
                   {
                       if constexpr (AuIsClass_v<T>)
                       {
                           delt->~T();
                       }
                       this->Free(delt);
                   });
        }

        template<typename T>
        static AuSPtr<T> ToSmartPointer(AuSPtr<Heap> heap, T *in, bool pinHeap = true)
        {
            auto handle = pinHeap ? heap : AuSPtr<Heap> {};
            auto ptr = heap.get();
            return AuSPtr<T>(in,
                   [handle,  ptr](T *delt)
                   {
                        if constexpr (AuIsClass_v<T>)
                        {
                            delt->~T();
                        }

                        ptr->Free(delt);
                   });
        }

    private:

        virtual AuSPtr<Heap> GetSelfReference() = 0; // may return empty/default. not all heaps are sharable.

        virtual AU_ALLOC void *_ZAlloc(Types::size_t uLength) = 0;
        virtual AU_ALLOC void *_ZAlloc(Types::size_t uLength, Types::size_t align) = 0;
        virtual AU_ALLOC void *_FAlloc(Types::size_t uLength) = 0;
        virtual AU_ALLOC void *_FAlloc(Types::size_t uLength, Types::size_t align) = 0;
        virtual AU_ALLOC void *_ZRealloc(void *pBase, Types::size_t uLength, Types::size_t uAlign) = 0;
        virtual AU_ALLOC void *_ZRealloc(void *pBase, Types::size_t uLength) = 0;
        virtual AU_ALLOC void *_FRealloc(void *pBase, Types::size_t uLength, Types::size_t uAlign) = 0;
        virtual AU_ALLOC void *_FRealloc(void *pBase, Types::size_t uLength) = 0;
        virtual void           _Free(void* pBase) = 0;
    };

    /**
        Returns a heap interface backed by the default allocator
    */
    AUKN_SHARED_API(GetDefaultDiscontiguousHeap, Heap);

    /**
        Allocates uLength amount of contiguous virtual memory

        @warning Heaps are guaranteed to outlive their allocations; heaps are the one object that effectively own a single reference count on themselves.
                 Requesting termination before all of its' memory has been free will result, in at worst, a warning.
                 Expect to leak unless all allocs have been paired by a free.

                 I do not expect to implement force frees simply because all our primary use cases keep track of dtors to forcefully release leaked objects.
                 Use RequestHeapOfRegion to be backed by caller owned memory.

        @return a heap backed by uLength bytes of virtual memory
    */
    AUKN_SHARED_API(AllocHeap, Heap, AuUInt uLength);

    AUKN_SHARED_API(RequestHeapOfRegion, Heap, void *pPtr, AuUInt uLength);

    // AllocHeap but use mimalloc (or the default allocator) instead
    AUKN_SHARED_API(AllocHeapMimalloc, Heap, AuUInt uLength);
}