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

    File: CppHeapWrapper.hpp
    Date: 2023-12-22
    Author: Reece
    Note: use AuPMRAllocator instead
***/
#pragma once

namespace Aurora::Memory
{
    namespace detail
    {
        struct AccessorICantEven
        {
            template <class T, class Z>
            cstatic void Set(const T &fuckCpp, Z that)
            {
                fuckCpp.pFuckCppRetardsFixYourWorthlessSpec2 = decltype(fuckCpp.pFuckCppRetardsFixYourWorthlessSpec2)(that);
            }
        };
    }

    /// @deprecated use ROXTL AuPMRAllocator instead
    /// Some legacy code treats the container as the handle, such that .get_allocator() still returns a reference.
    template <class T>
    struct CppHeapWrapper
    {
        using value_type = T;

        using size_type       = size_t;
        using difference_type = ptrdiff_t;

        using propagate_on_container_move_assignment = AuTrueType;
        using propagate_on_container_copy_assignment = AuTrueType;
        using is_always_equal = AuFalseType;

    #if defined(AU_LANG_CPP_17)
        using pointer         = T *;
        using const_pointer   = const T *;

        using reference       = T &;
        using const_reference = const T &;

        template <class Z>
        struct rebind
        {
            using other = CppHeapWrapper<Z>;
        };

        T *address(T &val) const noexcept
        {
            return std::addressof(val);
        }

        const T *address(const T &val) const noexcept
        {
            return std::addressof(val);
        }
    #endif

        inline CppHeapWrapper(std::shared_ptr<Heap> pHeap) :
        #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
            spHeap(pHeap),
            pHeap(pHeap.get())
        #else
            spHeap(pHeap)
        #endif
        { }

        inline CppHeapWrapper(Heap *pHeap, std::shared_ptr<void> pThat = {}) :
        #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
            spHeap(pThat),
            pHeap(pHeap)
        #else
            spHeap(AuSharedPointerFromShared(pHeap, pThat))
        #endif
        { }

        inline CppHeapWrapper(CppHeapWrapper *pFuckCppRetardsFixYourWorthlessSpec) :
            pFuckCppRetardsFixYourWorthlessSpec(pFuckCppRetardsFixYourWorthlessSpec)
        { }

        template <class B>
        inline CppHeapWrapper(const CppHeapWrapper<B> &fuckCpp)
        {
            this->pFuckCppRetardsFixYourWorthlessSpec = (CppHeapWrapper *)&fuckCpp;
            detail::AccessorICantEven::Set(fuckCpp, this);
        }

        inline CppHeapWrapper(const CppHeapWrapper &fuckCpp)
        {
            this->pFuckCppRetardsFixYourWorthlessSpec = (CppHeapWrapper *)&fuckCpp;
            detail::AccessorICantEven::Set(fuckCpp, this);
        }

        inline CppHeapWrapper(CppHeapWrapper &&fuckCpp)
        {
            if ((this->pFuckCppRetardsFixYourWorthlessSpec = fuckCpp.pFuckCppRetardsFixYourWorthlessSpec))
            {
                this->pFuckCppRetardsFixYourWorthlessSpec->pFuckCppRetardsFixYourWorthlessSpec2 = this;
            }
            fuckCpp.pFuckCppRetardsFixYourWorthlessSpec2 = fuckCpp.pFuckCppRetardsFixYourWorthlessSpec = nullptr;
        }

        inline ~CppHeapWrapper()
        {
            if (this->pFuckCppRetardsFixYourWorthlessSpec &&
                this->pFuckCppRetardsFixYourWorthlessSpec->pFuckCppRetardsFixYourWorthlessSpec2 == this)
            {
                this->pFuckCppRetardsFixYourWorthlessSpec->pFuckCppRetardsFixYourWorthlessSpec2 = this->pFuckCppRetardsFixYourWorthlessSpec2;
            }
            if (this->pFuckCppRetardsFixYourWorthlessSpec2)
            {
            #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
                this->pFuckCppRetardsFixYourWorthlessSpec2->pHeap = this->pHeap;
            #endif
                this->pFuckCppRetardsFixYourWorthlessSpec2->spHeap = this->spHeap;
                this->pFuckCppRetardsFixYourWorthlessSpec2->pFuckCppRetardsFixYourWorthlessSpec = this->pFuckCppRetardsFixYourWorthlessSpec;
                this->pFuckCppRetardsFixYourWorthlessSpec2 = nullptr;
            }
        }

        AU_DEF(CppHeapWrapper)
        AU_OPERATOR_COPY_MOVE(CppHeapWrapper)

        void *allocate_bytes(std::size_t nbytes,
                             std::size_t alignment = alignof(std::max_align_t))
        {
            if (auto pRet = this->GetHeapRaw()->FAlloc(nbytes, alignment))
            {
                return pRet;
            }
            else
            {
                throw std::bad_alloc();
            }
        }

        void deallocate_bytes(void *p,
                              std::size_t nbytes,
                              std::size_t alignment = alignof(std::max_align_t)) noexcept
        {
            return this->GetHeapRaw()->Free(p);
        }
        
    #if defined(AU_LANG_CPP_23_)
        std::allocation_result<T *> allocate_at_least(const size_t count)
        {
            auto pThat = this->allocate(count);
            return { (T *)pThat, this->GetHeapRaw()->GetChunkSize(pThat) / sizeof(T) };
        }
    #endif
        
        template <class U, class... Args >
        void construct(U *p, Args&&... args)
        {
            new ((void *)p) U(AuForward<Args>(args)...);
        }

        template <class U, class... Args >
        void construct_at(U *p, Args&&... args)
        {
            new ((void *)p) U(AuForward<Args>(args)...);
        }

        void deallocate(const T *pType,
                        const size_t count) noexcept
        {
            this->deallocate_bytes((void *)pType, 0, 0);
        }

        AU_ALLOC T *allocate(const size_t count)
        {
            if (!count)
            {
                return nullptr;
            }
            
            return (T *)this->allocate_bytes(count * sizeof(T), alignof(T));
        }

        template <class U>
        U *allocate_object(std::size_t n = 1)
        {
            return (U *)this->allocate_bytes(sizeof(U) * n, alignof(U));
        }

        template <class U>
        void deallocate_object(U *p, std::size_t n = 1)
        {
            this->deallocate_bytes(p, 0, 0);
        }

        template <class U, class... CtorArgs>
        U *new_object(CtorArgs&&... args)
        {
            U *p = this->allocate_object<U>();
            try
            {
                this->construct(p, AuForward<CtorArgs>(args)...);
            }
            catch (...)
            {
                this->deallocate_object(p);
                throw;
            }
            return p;
        }

        template <class U>
        void delete_object(U *p)
        {
            if constexpr (AuIsClass_v<U> &&
                          !AuIsTriviallyDestructible_v<U>)
            {
                p->~U();
            }

            this->deallocate_object(p);
        }

        bool IsInitialized() const
        {
        #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
            return bool(this->pHeap);
        #else
            return bool(this->spHeap);
        #endif
        }

        std::shared_ptr<Heap> GetHeap() const
        {
            if (this->pFuckCppRetardsFixYourWorthlessSpec)
            {
                return this->pFuckCppRetardsFixYourWorthlessSpec->GetHeap();
            }
            else
            {
            #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
                return AuSharedPointerFromShared(this->pHeap, this->spHeap);
            #else
                return this->spHeap;
            #endif
            }
        }

        Heap *GetHeapRaw() const
        {
            Heap *pRet {};
            
            if (this->pFuckCppRetardsFixYourWorthlessSpec)
            {
                pRet = this->pFuckCppRetardsFixYourWorthlessSpec->GetHeapRaw();
            }
            else
            {
            #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
                pRet = this->pHeap;
            #else
                pRet = this->spHeap.get();
            #endif
            }

            if (pRet)
            {
                return pRet;
            }

        #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
            return this->pHeap = __audetail::gDefaultDiscontiguousHeap;
        #else
            return (this->spHeap = AuUnsafeRaiiToShared(__audetail::gDefaultDiscontiguousHeap)).get();
        #endif
        }

        void SetHeap(std::shared_ptr<Heap> pHeap)
        {
            if (this->IsInitialized())
            {
                return;
            }
            if (this->pFuckCppRetardsFixYourWorthlessSpec)
            {
                this->pFuckCppRetardsFixYourWorthlessSpec->SetHeap(pHeap);
            }
            else
            {
                this->spHeap = pHeap;
            #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
                this->pHeap = pHeap.get();
            #endif
            }
        }

        void SetHeapRaw(Heap *pHeap)
        {
            if (this->IsInitialized())
            {
                return;
            }
            if (this->pFuckCppRetardsFixYourWorthlessSpec)
            {
                this->pFuckCppRetardsFixYourWorthlessSpec->SetHeapRaw(pHeap);
            }
            else
            {
            #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
                AuResetMember(this->spHeap);
                this->pHeap = pHeap;
            #else
                this->spHeap = AuUnsafeRaiiToShared(pHeap);
            #endif
            }
        }

        template <class Z>
        bool operator==(const Aurora::Memory::CppHeapWrapper<Z> &rhs) noexcept
        {
            return this->GetHeapRaw() == rhs.GetHeapRaw();
        }

    private:
        // should be sizeof(void *) * 4 = [pHeap, pControlBlock, pParent, pSingleThreadChild]
        // nor not. it doesnt matter.
    #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
        mutable std::shared_ptr<void> spHeap;
        mutable Heap *pHeap {};
    #else
        mutable std::shared_ptr<Heap> spHeap;
    #endif
        mutable CppHeapWrapper *pFuckCppRetardsFixYourWorthlessSpec {};
        mutable CppHeapWrapper *pFuckCppRetardsFixYourWorthlessSpec2 {};

        template <typename Z>
        friend struct CppHeapWrapper;
        friend struct detail::AccessorICantEven;
    };
}

template <class T, class Z>
inline bool operator==(const Aurora::Memory::CppHeapWrapper<T> &lhs,
                       const Aurora::Memory::CppHeapWrapper<Z> &rhs) noexcept
{
    return lhs.GetHeapRaw() == rhs.GetHeapRaw();
}

template <class T, class Z>
inline bool operator!=(const Aurora::Memory::CppHeapWrapper<T> &lhs,
                       const Aurora::Memory::CppHeapWrapper<Z> &rhs) noexcept
{
    return lhs.GetHeapRaw() != rhs.GetHeapRaw();
}