#ifndef JKRHEAP_H
#define JKRHEAP_H

#include "dolphin/os/OSMutex.h"
#include "JSystem/JKernel/JKRDisposer.h"

#include "types.h"

#ifdef __cplusplus
typedef void JKRHeapErrorHandler(void*, u32, int);

class JKRHeap : public JKRDisposer
{
public:
    enum EAllocMode
    {
        HEAPALLOC_Unk1 = 1,
    };

    struct TState
    { // NB: this struct doesn't agree with TP's struct
        struct TLocation
        {
            TLocation() : _00(nullptr), _04(-1)
            {
            }

            void* _00; // _00
            int _04;   // _04
        };

        struct TArgument
        {
            TArgument(const JKRHeap* heap, u32 p2, bool p3)
                : mHeap((heap) ? heap : JKRHeap::sCurrentHeap), mId(p2), mIsCompareOnDestructed(p3)
            {
            }

            const JKRHeap* mHeap;        // _00
            u32 mId;                     // _04
            bool mIsCompareOnDestructed; // _08
        };

        TState(const JKRHeap* heap, u32 id, bool isCompareOnDestructed)
            : mUsedSize(0), mCheckCode(0), mArgument(heap, id, isCompareOnDestructed)
        {
            mArgument.mHeap->state_register(this, mArgument.mId);
        }

        TState(JKRHeap* heap)
            : mUsedSize(0), mCheckCode(0), mArgument(heap, 0xFFFFFFFF, true)
        {
        }

        ~TState();
        void dump() const { mArgument.mHeap->state_dump(*this); }
        bool isVerbose() { return bVerbose_; };
        bool isCompareOnDestructed() const { return mArgument.mIsCompareOnDestructed; };
        u32 getUsedSize() const { return mUsedSize; }
        u32 getCheckCode() const { return mCheckCode; }
        const JKRHeap* getHeap() const { return mArgument.mHeap; }
        u32 getId() const { return mArgument.mId; }

        // unused/inlined:
        TState(const JKRHeap::TState::TArgument& arg, const JKRHeap::TState::TLocation& location);
        TState(const JKRHeap::TState& other, bool p2);
        TState(const JKRHeap::TState& other, const JKRHeap::TState::TLocation& location, bool p3);

        static bool bVerbose_;

        u32 mUsedSize;       // _00
        u32 mCheckCode;      // _04, plausibly TLocation when combined with _00
        u32 mBuf;            // _08
        u8 _0C[0x4];         // _0C
        TArgument mArgument; // _10
        TLocation mLocation; // _1C
    };

public:
    JKRHeap(void*, u32, JKRHeap*, bool);

    bool setErrorFlag(bool errorFlag);
    bool isSubHeap(JKRHeap* heap) const;

    virtual ~JKRHeap();
    virtual void callAllDisposer();
    virtual void* do_alloc(u32, int) = 0;
    virtual void do_free(void*) = 0;
    virtual void do_freeAll() = 0;
    virtual void do_freeTail() = 0;
    virtual s32 do_resize(void*, u32) = 0;
    virtual s32 do_getSize(void*) = 0;
    virtual s32 do_getFreeSize() = 0;
    virtual s32 do_getTotalFreeSize() = 0;
    virtual u32 getHeapType() = 0;
    virtual bool check() = 0;
    virtual bool dump_sort() { return true; }
    virtual bool dump() = 0;
    virtual s32 do_changeGroupID(u8 newGroupID) { return 0; }
    virtual u8 do_getCurrentGroupId() { return 0; }
    virtual void state_register(JKRHeap::TState*, u32) const;
    virtual bool state_compare(JKRHeap::TState const&, JKRHeap::TState const&) const;
    virtual void state_dump(JKRHeap::TState const&) const;

    JKRHeap* becomeSystemHeap();
    JKRHeap* becomeCurrentHeap();
    void destroy();
    void* alloc(u32, int);
    void free(void*);
    void freeAll();
    void freeTail();
    void fillFreeArea();
    void resize(void*, u32);
    static s32 getSize(void*, JKRHeap*);

    // ... more functions

    s32 getSize(void* ptr);
    s32 getFreeSize();
    void* getMaxFreeBlock();
    s32 getTotalFreeSize();
    u8 getCurrentGroupId();
    s32 changeGroupID(u8 newGroupId);
    u32 getMaxAllocatableSize(int alignment);
    JKRHeap* find(void*) const;        // 0x80084640
    JKRHeap* findAllHeap(void*) const; // 0x8008492c
    void dispose_subroutine(u32 begin, u32 end);
    bool dispose(void*, u32);    // 0x80084b9c
    void dispose(void*, void*); // 0x80084c2c
    void dispose();               // 0x80084cb8

    void appendDisposer(JKRDisposer* disposer)
    {
        mDisposerList.append(&disposer->mPointerLinks);
    }

    void removeDisposer(JKRDisposer* disposer)
    {
        mDisposerList.remove(&disposer->mPointerLinks);
    }

    void setDebugFill(bool debugFill) { mDebugFill = debugFill; }
    bool getDebugFill() const { return mDebugFill; }
    void* getStartAddr() const { return (void*)mStart; }
    void* getEndAddr() const { return (void*)mEnd; }
    u32 getHeapSize() const { return mSize; }
    bool getErrorFlag() const { return mErrorFlag; }
    void callErrorHandler(JKRHeap* heap, u32 size, int alignment)
    {
        if (mErrorHandler)
        {
            (*mErrorHandler)(heap, size, alignment);
        }
    }

    // TState related
    static u32 getState_buf_(TState* state) { return state->mBuf; } // might instead be a pointer to a next state?
    static void setState_u32ID_(TState* state, u32 id)
    {
        state->mArgument.mId = id;
    }
    static void setState_uUsedSize_(TState* state, u32 usedSize)
    {
        state->mUsedSize = usedSize;
    }
    static void setState_u32CheckCode_(TState* state, u32 checkCode) { state->mCheckCode = checkCode; }

    void lock() const { OSLockMutex(const_cast<OSMutex*>(&mMutex)); }
    void unlock() const { OSUnlockMutex(const_cast<OSMutex*>(&mMutex)); }

    JKRHeap* getParent()
    {
        return mChildTree.getParent()->getObject();
    }

    const JSUTree<JKRHeap>& getHeapTree()
    {
        return mChildTree;
    }

    // Unused
    void checkMemoryFilled(u8*, u32 size, u8);

    static void destroy(JKRHeap* heap); // fabricated
    static bool initArena(char**, u32*, int);
    static void* alloc(u32, int, JKRHeap*);
    static void copyMemory(void*, void*, u32);
    static void free(void*, JKRHeap*);
    static void state_dumpDifference(const TState&, const TState&);
    static JKRHeap* findFromRoot(void*);
    static JKRHeapErrorHandler* setErrorHandler(JKRHeapErrorHandler*);

    static void* getCodeStart()
    {
        return mCodeStart;
    }

    static void* getCodeEnd()
    {
        return mCodeEnd;
    }

    static void* getUserRamStart()
    {
        return mUserRamStart;
    }

    static void* getUserRamEnd()
    {
        return mUserRamEnd;
    }

    static u32 getMemorySize()
    {
        return mMemorySize;
    }

    static JKRHeap* getCurrentHeap()
    {
        return sCurrentHeap;
    }

    static JKRHeap* getRootHeap()
    {
        return sRootHeap;
    }

    static JKRHeap* getSystemHeap()
    {
        return sSystemHeap;
    }

    static void* mCodeStart;
    static void* mCodeEnd;
    static void* mUserRamStart;
    static void* mUserRamEnd;
    static u32 mMemorySize;

    static JKRHeap* sSystemHeap;
    static JKRHeap* sCurrentHeap;
    static JKRHeap* sRootHeap;

    static bool sDefaultFillFlag;
    static bool sDefaultFillCheckFlag;

    static JKRHeapErrorHandler* mErrorHandler;

protected:
    /* 0x00 */ // vtable
    /* 0x04 */ // JKRDisposer
    /* 0x18 */ OSMutex mMutex;
    /* 0x30 */ void* mStart;
    /* 0x34 */ void* mEnd;
    /* 0x38 */ u32 mSize;
    /* 0x3C */ bool mDebugFill;
    /* 0x3D */ bool mCheckMemoryFilled;
    /* 0x3E */ u8 mAllocationMode; // EAllocMode?
    /* 0x3F */ u8 mGroupId;
    /* 0x40 */ JSUTree<JKRHeap> mChildTree;
    /* 0x5C */ JSUList<JKRDisposer> mDisposerList;
    /* 0x68 */ bool mErrorFlag;
    /* 0x69 */ bool mInitFlag;
    /* 0x6A */ u8 padding_0x6a[2];
};

inline JKRHeap* JKRGetCurrentHeap()
{
    return JKRHeap::getCurrentHeap();
}

inline JKRHeap* JKRGetSystemHeap()
{
    return JKRHeap::getSystemHeap();
}

inline JKRHeap* JKRGetRootHeap()
{
    return JKRHeap::getRootHeap();
}

inline void* JKRAllocFromSysHeap(u32 size, int alignment)
{
    return JKRHeap::getSystemHeap()->alloc(size, alignment);
}

inline void* JKRAllocFromHeap(JKRHeap* heap, u32 size, int alignment)
{
    return JKRHeap::alloc(size, alignment, heap);
}

inline void JKRFree(void* pBuf)
{
    JKRHeap::free(pBuf, nullptr);
}

inline void JKRFreeToHeap(JKRHeap* heap, void* ptr)
{
    JKRHeap::free(ptr, heap);
}

inline void JKRFreeToSysHeap(void* buf)
{
    JKRHeap::getSystemHeap()->free(buf);
}

void JKRDefaultMemoryErrorRoutine(void*, u32, int);

void* operator new(size_t);
void* operator new(size_t, s32);
void* operator new(size_t, JKRHeap*, int);
inline void* operator new(size_t, void* buf) { return buf; } // i believe this is actually part of MSL_C?

void* operator new[](size_t);
void* operator new[](size_t, s32);
void* operator new[](size_t, JKRHeap*, int);

void operator delete(void*);
void operator delete[](void*);

#endif
#endif // !JKRHEAP_H