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

    File: Heap.inl
    Date: 2024-7-14
    Date: 2021-6-9
    Author: Reece
***/
#pragma once

namespace Aurora::Memory
{
    template<typename T>
    T Heap::ZAlloc(Types::size_t uLength)
    {
        if constexpr (AuIsVoid_v<AuRemovePointer_t<T>>)
        {
            return reinterpret_cast<T>(_ZAlloc(uLength));
        }
        else
        {
            return reinterpret_cast<T>(_ZAlloc(uLength, alignof(AuRemovePointer_t<T>)));
        }
    }

    template<typename T>
    T Heap::ZAlloc(Types::size_t uLength, Types::size_t uAlignment)
    {
        return reinterpret_cast<T>(_ZAlloc(uLength, uAlignment));
    }

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

    template<typename T>
    T *Heap::NewArray(Types::size_t uLength)
    {
        return ZAlloc<T *>(uLength * sizeof(T), alignof(T));
    }

    template<typename T>
    T *Heap::NewArray(Types::size_t uLength, Types::size_t uAlignment)
    {
        return ZAlloc<T *>(uLength * sizeof(T), uAlignment);
    }

    template<typename T>
    T Heap::FAlloc(Types::size_t uLength)
    {
        if constexpr (AuIsVoid_v<AuRemovePointer_t<T>>)
        {
            return reinterpret_cast<T>(_FAlloc(uLength));
        }
        else
        {
            return reinterpret_cast<T>(_FAlloc(uLength, alignof(AuRemovePointer_t<T>)));
        }
    }

    template<typename T>
    T Heap::FAlloc(Types::size_t uLength, Types::size_t uAlignment)
    {
        return reinterpret_cast<T>(_FAlloc(uLength, uAlignment));
    }

    template<typename T>
    T *Heap::FAlloc()
    {
        return reinterpret_cast<T *>(_FAlloc(sizeof(T), alignof(T)));
    }

    // Reallocs
    template<typename T>
    T Heap::ZRealloc(T pHead, Types::size_t uLength)
    {
        if constexpr (AuIsVoid_v<AuRemovePointer_t<T>>)
        {
            return reinterpret_cast<T>(_ZRealloc(reinterpret_cast<void *>(pHead), uLength));
        }
        else
        {
            return reinterpret_cast<T>(_ZRealloc(reinterpret_cast<void *>(pHead), uLength, alignof(AuRemovePointer_t<T>)));
        }
    }

    template<typename T>
    T Heap::ZRealloc(T pHead, Types::size_t uLength, Types::size_t uAlignment)
    {
        return reinterpret_cast<T>(_ZRealloc(reinterpret_cast<void *>(pHead), uLength, uAlignment));
    }

    template<typename T>
    T Heap::FRealloc(T pHead, Types::size_t uLength)
    {
        if constexpr (AuIsVoid_v<AuRemovePointer_t<T>>)
        {
            return reinterpret_cast<T>(_FRealloc(reinterpret_cast<void *>(pHead), uLength));
        }
        else
        {
            return reinterpret_cast<T>(_FRealloc(reinterpret_cast<void *>(pHead), uLength, alignof(AuRemovePointer_t<T>)));
        }
    }

    template<typename T>
    T Heap::FRealloc(T pHead, Types::size_t uLength, Types::size_t uAlignment)
    {
        return reinterpret_cast<T>(_FRealloc(reinterpret_cast<void *>(pHead), uLength, uAlignment));
    }

    // Free
    template<typename T>
    void Heap::Free(T pHead)
    {
        _Free(reinterpret_cast<void *>(pHead));
    }

    template <typename T>
    void Heap::DeleteThat(T *pThat)
    {
        static const auto kAlignment = AuMax(alignof(T), sizeof(void *));

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyDestructible_v<T>)
        {
            pThat->~T();
        }

        auto &pHeap = *(Heap **)(((char *)pThat) - kAlignment);
        pHeap->_Free(&pHeap);
    }

    template <typename T>
    void Heap::DeleteThatArray(T *pThat)
    {
        static const auto kAlignment = AuMax(alignof(T), sizeof(void *) * 2);

        auto pVoids = (void **)(((char *)pThat) - kAlignment);
        auto pHeap  = (Heap *)pVoids[0];
        auto uLength = (AuUInt)pVoids[1];

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyDestructible_v<T>)
        {
            for (AU_ITERATE_N(i, uLength))
            {
                auto &refElement = pThat[i];
                refElement.~T();
            }
        }

        pHeap->_Free(pVoids);
    }
        
    template <typename T>
    void Heap::DeleteThatArray2(T *pThat)
    {
        auto pBase = ((void **)pThat)[-1];

        auto pVoids = (void **)pBase;
        auto pHeap  = (Heap *)pVoids[0];
        auto uLength = (AuUInt)pVoids[1];

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyDestructible_v<T>)
        {
            for (AU_ITERATE_N(i, uLength))
            {
                auto &refElement = pThat[i];
                refElement.~T();
            }
        }

        pHeap->_Free(pVoids);
    }
    template <typename T, typename Z>
    void Heap::DeleteThatCastedOnce(T *pThat)
    {
        static const auto kAlignment = AuMax(alignof(Z), sizeof(void *));

        auto pBaseClass = AuStaticCast<Z>(pThat);

        if constexpr (AuIsClass_v<Z> &&
                      !AuIsTriviallyDestructible_v<Z>)
        {
            pBaseClass->~Z();
        }

        auto &pHeap = *(Heap **)(((char *)pBaseClass) - kAlignment);
        pHeap->_Free(&pHeap);
    }

    template <typename T>
    void Heap::RetardedSpecWrittenByRetards(T *pThat)
    {

    }

    template <class T, class ...Args>
    AuSPtr<T> Heap::NewClass(Args &&...args)
    {
        static const auto kAlignment = AuMax(alignof(T), sizeof(void *));
        AuUInt8 *pPtr;

        auto pThat = this->GetSelfReferenceRaw();
        if (!pThat)
        {
            pThat = this;
        }

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyConstructible_v<T, Args...>)
        {
            pPtr = pThat->FAlloc<AuUInt8 *>(sizeof(T) + kAlignment, kAlignment);
            if (pPtr)
            {
                new (pPtr + kAlignment) T(AuForward<Args>(args)...);
            }
        }
        else
        {
            pPtr = pThat->ZAlloc<AuUInt8 *>(sizeof(T) + kAlignment, kAlignment);
        }

        if (!pPtr)
        {
            return {};
        }
        
        *(void **)pPtr = pThat;

        auto pTThat = (T *)(pPtr + kAlignment);
        AUROXTL_COMMODITY_TRY
        {
            return AuSPtr<T>(pTThat, &Heap::DeleteThat<T>, CppHeapWrapper<T> { this });
        }
        AUROXTL_COMMODITY_CATCH
        {
            Heap::DeleteThat<T>(pTThat);
            return {};
        }
    }

    template <class T, class Z, class ...Args>
    AuUPtr<Z, decltype(&Heap::DeleteThat<Z>)> Heap::NewClassUnique(Args &&...args)
    {
        static const auto kAlignment = AuMax(alignof(T), sizeof(void *));
        AuUInt8 *pPtr;

        auto pThat = this->GetSelfReferenceRaw();
        if (!pThat)
        {
            pThat = this;
        }

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyConstructible_v<T, Args...>)
        {
            pPtr = pThat->FAlloc<AuUInt8 *>(sizeof(T) + kAlignment, kAlignment);
            if (pPtr)
            {
                new (pPtr + kAlignment) T(AuForward<Args>(args)...);
            }
        }
        else
        {
            pPtr = pThat->ZAlloc<AuUInt8 *>(sizeof(T) + kAlignment, kAlignment);
        }

        if (!pPtr)
        {
            return AuUPtr<Z, decltype(&Heap::DeleteThat<Z>)>(nullptr, &Heap::RetardedSpecWrittenByRetards<Z>);
        }
        
        *(void **)pPtr = pThat;
            
        if constexpr (AuIsSame_v<T, Z>)
        {
            return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>((T *)(pPtr + kAlignment), &Heap::DeleteThat<T>);
        }
        else
        {
            return Heap::CastPointer<Z>(AuMove(AuUPtr<T, decltype(&Heap::DeleteThat<T>)>((T *)(pPtr + kAlignment), &Heap::DeleteThat<T>)));
        }
    }

    template <class T, class ...Args>
    AuSPtr<T> Heap::NewClassArray(AuUInt uElements, Args &&... fillCtr)
    {
        const auto kAlignment = AuMax(alignof(T), sizeof(void *) * 2);
        AuUInt8 *pPtr;

        if (!uElements)
        {
            return {};
        }

        auto pThat = this->GetSelfReferenceRaw();
        if (!pThat)
        {
            pThat = this;
        }

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyConstructible_v<T, Args...>)
        {
            if (bool(pPtr = pThat->FAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                for (AU_ITERATE_N(i, uElements))
                {
                    new (pPtr + kAlignment + (sizeof(T) * i)) T(AuForward<Args>(fillCtr)...);
                }
            }
        }
        else
        {
            if (bool(pPtr = pThat->ZAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                if constexpr (sizeof...(Args) != 0)
                {
                #if defined(AURT_HEAP_NO_STL)
                    static_assert(false);
                #else
                    auto pElements = (T *)(pPtr + kAlignment);
                    std::fill(pElements, pElements + uElements, AuForward<Args>(fillCtr)...);
                #endif
                }
            }
        }

        if (!pPtr)
        {
            return {};
        }

        auto pVoids = (void **)pPtr;
        pVoids[0] = pThat;
        pVoids[1] = (void *)uElements;
            
        auto pTThat = (T *)(pPtr + kAlignment);
        AUROXTL_COMMODITY_TRY
        {
            return AuSPtr<T>(pTThat, &Heap::DeleteThatArray<T>, CppHeapWrapper<T> { this });
        }
        AUROXTL_COMMODITY_CATCH
        {
            Heap::DeleteThatArray<T>(pTThat);
            return {};
        }
    }
    
    template <class T, class ...Args>
    AuSPtr<T> Heap::NewClassArray2(AuUInt uElements, AuUInt uAlignment, Args &&... fillCtr)
    {
        const auto kAlignment = AuMax(uAlignment, sizeof(void *) * 4);
        AuUInt8 *pPtr;

        if (!uElements)
        {
            return {};
        }

        auto pThat = this->GetSelfReferenceRaw();
        if (!pThat)
        {
            pThat = this;
        }

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyConstructible_v<T, Args...>)
        {
            if (bool(pPtr = pThat->FAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                for (AU_ITERATE_N(i, uElements))
                {
                    new (pPtr + kAlignment + (sizeof(T) * i)) T(AuForward<Args>(fillCtr)...);
                }
            }
        }
        else
        {
            if (bool(pPtr = pThat->ZAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                if constexpr (sizeof...(Args) != 0)
                {
                #if defined(AURT_HEAP_NO_STL)
                    static_assert(false);
                #else
                    auto pElements = (T *)(pPtr + kAlignment);
                    std::fill(pElements, pElements + uElements, AuForward<Args>(fillCtr)...);
                #endif
                }
            }
        }

        if (!pPtr)
        {
            return {};
        }

        auto pVoids = (void **)pPtr;
        pVoids[0] = pThat;
        pVoids[1] = (void *)uElements;
        ((void **)(pPtr + kAlignment))[-1] = pPtr;
        auto pTThat = (T *)(pPtr + kAlignment);

        AUROXTL_COMMODITY_TRY
        {
            return AuSPtr<T>(pTThat, &Heap::DeleteThatArray2<T>, CppHeapWrapper<T> { this });
        }
        AUROXTL_COMMODITY_CATCH
        {
            Heap::DeleteThatArray2<T>(pTThat);
            return {};
        }
    }

    template <class T, class ...Args>
    AuUPtr<T, decltype(&Heap::DeleteThat<T>)> Heap::NewClassArrayUnique(AuUInt uElements, Args &&... fillCtr)
    {
        const auto kAlignment = AuMax(alignof(T), sizeof(void *) * 2);
        AuUInt8 *pPtr;

        if (!uElements)
        {
            return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>(nullptr, &Heap::RetardedSpecWrittenByRetards<T>);
        }

        auto pThat = this->GetSelfReferenceRaw();
        if (!pThat)
        {
            pThat = this;
        }

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyConstructible_v<T, Args...>)
        {
            if (bool(pPtr = pThat->FAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                for (AU_ITERATE_N(i, uElements))
                {
                    new (pPtr + kAlignment + (sizeof(T) * i)) T(AuForward<Args>(fillCtr)...);
                }
            }
        }
        else
        {
            if (bool(pPtr = pThat->ZAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                if constexpr (sizeof...(Args) != 0)
                {
                #if defined(AURT_HEAP_NO_STL)
                    static_assert(false);
                #else
                    auto pElements = (T *)(pPtr + kAlignment);
                    std::fill(pElements, pElements + uElements, AuForward<Args>(fillCtr)...);
                #endif
                }
            }
        }

        if (!pPtr)
        {
            return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>(nullptr, &Heap::RetardedSpecWrittenByRetards<T>);
        }

        auto pVoids = (void **)pPtr;
        pVoids[0] = pThat;
        pVoids[1] = (void *)uElements;

        return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>((T *)(pPtr + kAlignment), &Heap::DeleteThatArray<T>);
    }
    
    template <class T, class ...Args>
    AuUPtr<T, decltype(&Heap::DeleteThat<T>)> Heap::NewClassArray2Unique(AuUInt uElements, AuUInt uAlignment, Args &&... fillCtr)
    {
        const auto kAlignment = AuMax(uAlignment, sizeof(void *) * 4);
        AuUInt8 *pPtr;

        if (!uElements)
        {
            return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>(nullptr, &Heap::RetardedSpecWrittenByRetards<T>);
        }

        auto pThat = this->GetSelfReferenceRaw();
        if (!pThat)
        {
            pThat = this;
        }

        if constexpr (AuIsClass_v<T> &&
                      !AuIsTriviallyConstructible_v<T, Args...>)
        {
            if (bool(pPtr = pThat->FAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                for (AU_ITERATE_N(i, uElements))
                {
                    new (pPtr + kAlignment + (sizeof(T) * i)) T(AuForward<Args>(fillCtr)...);
                }
            }
        }
        else
        {
            if (bool(pPtr = pThat->ZAlloc<AuUInt8 *>((sizeof(T) * uElements) + kAlignment, kAlignment)))
            {
                if constexpr (sizeof...(Args) != 0)
                {
                #if defined(AURT_HEAP_NO_STL)
                    static_assert(false);
                #else
                    auto pElements = (T *)(pPtr + kAlignment);
                    std::fill(pElements, pElements + uElements, AuForward<Args>(fillCtr)...);
                #endif
                }
            }
        }

        if (!pPtr)
        {
            return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>(nullptr, &Heap::RetardedSpecWrittenByRetards<T>);
        }

        auto pVoids = (void **)pPtr;
        pVoids[0] = pThat;
        pVoids[1] = (void *)uElements;

        ((void **)(pPtr + kAlignment))[-1] = pPtr;
        auto pTThat = (T *)(pPtr + kAlignment);

        return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>(pTThat, &Heap::DeleteThatArray2<T>);
    }

    template <class T>
    AuUPtr<T, decltype(&Heap::DeleteThat<T>)> Heap::NullUniquePointer()
    {
        return AuUPtr<T, decltype(&Heap::DeleteThat<T>)>(nullptr, &Heap::RetardedSpecWrittenByRetards<T>);
    }

    template <class Z, class T>
    AuUPtr<Z, decltype(&Heap::DeleteThat<Z>)> Heap::CastPointer(AuUPtr<T, decltype(&Heap::DeleteThat<T>)> &&pInPointer)
    {
        if (!pInPointer)
        {
            return NullUniquePointer<Z>();
        }
        else if (pInPointer.get_deleter() == &Heap::DeleteThat<T>)
        {
            return AuUPtr<Z, decltype(&Heap::DeleteThat<Z>)>(AuStaticCast<Z>(pInPointer.release()), &Heap::DeleteThatCastedOnce<Z, T>);
        }
        else
        {
            return NullUniquePointer<Z>();
        }
    }

    namespace detail
    {
        inline AuSPtr<AuUInt8> AllocateArray(Heap *pHeap, AuUInt uLength, AuUInt32 uAlignment)
        {
            return pHeap->NewClassArray2<AuUInt8>(uLength, uAlignment);
        }
    }
}