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

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

namespace Aurora::Memory
{
    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 is_always_equal = AuTrueType;

    #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)
        { }

        inline CppHeapWrapper(const CppHeapWrapper &fuckCpp)
        {
            this->pFuckCppRetardsFixYourWorthlessSpec = (CppHeapWrapper *)&fuckCpp;
            fuckCpp.pFuckCppRetardsFixYourWorthlessSpec2 = 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)

        constexpr void deallocate(const T *pType,
                                  const size_t count)
        {
            this->GetHeapRaw()->Free((T *)pType);
        }

        constexpr AU_ALLOC T *allocate(const size_t count)
        {
            if (!count)
            {
                return nullptr;
            }
            
            auto pData = this->GetHeapRaw()->FAlloc(count * sizeof(T), alignof(T));
            if (!pData)
            {
                throw std::bad_alloc();
            }
            return (T *)pData;
        }

        template <class U, class... Args >
        void construct(U *p, Args&&... args)
        {
            if constexpr (AuIsClass_v<T>)
            {
                new ((void *)p) T(AuForward<Args &&>(args)...);
            }
        }

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

    #if defined(AU_LANG_CPP_23)
        constexpr 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

        std::shared_ptr<Heap> GetHeap()
        {
            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()
        {
            if (this->pFuckCppRetardsFixYourWorthlessSpec)
            {
                return this->pFuckCppRetardsFixYourWorthlessSpec->GetHeapRaw();
            }
            else
            {
            #if defined(AU_NO_COMPRESS_CPPHEAP_WRAPPER)
                return this->pHeap;
            #else
                return this->spHeap.get();
            #endif
            }
        }

        void SetHeap(std::shared_ptr<Heap> pHeap)
        {
            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->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
            }
        }

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