#ifndef EGG_HEAP_H
#define EGG_HEAP_H

// This file is ported from https://github.com/riidefi/mkw/blob/master/source/egg/core/eggHeap.hpp

#include "egg/core/eggDisposer.h"
#include "egg/core/eggThread.h"
#include "egg/prim/eggBitFlag.h"
#include "nw4r/ut.h"
#include "rvl/MEM.h"
#include "rvl/OS.h"
#include <common.h>

namespace EGG {

// class ExpHeap;
class Allocator;

struct HeapAllocArg {
    void *userArg; // 00
    u32 size;      // 04
    int align;     // 08
    void *ptr;     // 0C the allocated ptr
    Heap *heap;    // 10 heap to allocate in

    inline HeapAllocArg() : userArg(0), size(0), align(0), heap(nullptr), ptr(nullptr) {}
};
typedef void (*HeapAllocCallback)(HeapAllocArg *arg);

struct HeapErrorArg {
    const char *msg;
    void *userdata;

    inline HeapErrorArg() {}
};
typedef void (*ErrorCallback)(HeapErrorArg *);

struct HeapFreeArg {
    void *userArg;
    void *ptr;
    Heap *heap;
};
typedef void (*HeapFreeCallback)(HeapFreeArg *);

typedef void (*HeapCreateCallback)(void *);
typedef void (*HeapDestroyCallback)(void *);

class Heap : public Disposer {
public:
    enum eHeapKind {
        HEAP_KIND_NONE = 0,
        HEAP_KIND_EXPANDED,
        HEAP_KIND_FRAME,
        HEAP_KIND_UNIT,
        HEAP_KIND_ASSERT,
    };
    // vtable at 0x0 | 8056e950
    /* vt 0x08 | 804954c0 */ virtual ~Heap();
    /* vt 0x0C | 00000000 */ virtual eHeapKind getHeapKind() const = 0;
    /* vt 0x10 | 80495a40 */ virtual void initAllocator(Allocator *allocator, s32 align);
    /* vt 0x14 | 00000000 */ virtual void *alloc(u32 size, s32 align) = 0;
    /* vt 0x18 | 00000000 */ virtual void free(void *block) = 0;
    /* vt 0x1C | 00000000 */ virtual void destroy() = 0;
    /* vt 0x20 | 00000000 */ virtual u32 resizeForMBlock(void *block, u32 size) = 0;
    /* vt 0x24 | 00000000 */ virtual u32 getTotalFreeSize() = 0;
    /* vt 0x28 | 00000000 */ virtual u32 getAllocatableSize(s32 align) = 0;
    /* vt 0x2C | 00000000 */ virtual u32 adjust() = 0;

public:
    void setName(const char *name) {
        mName = name;
    }
    inline bool isExpHeap() {
        return getHeapKind() == HEAP_KIND_EXPANDED;
    }
    // inline Heap* getParentHeap() { return mParentHeap; } // not part of ss
    inline void *getStartAddress() {
        return this;
    }

public:
    // members
    /* 0x10 */ MEMiHeapHead *mHeapHandle;
    /* 0x14 */ void *mParentBlock;
    /* 0x18 */ TBitFlag<u16> mFlag;
    /* 0x1C */ nw4r::ut::Node mNode;
    /* 0x24 */ nw4r::ut::List mChildren;
    /* 0x30 */ const char *mName; // set to "NoName" in ctor

    enum HeapFlag {
        // tstDisableAllocation, enableAllocation, disableAllocation
        // setBit__Q23EGG12TBitFlag<Us>FUc
        HEAP_FLAG_LOCKED = (1 << 0)
    };

public:
    /* 804953f0 */ static void initialize();
    /* 80495430 */ Heap(MEMiHeapHead *heapHandle);
    /* 80495690 */ static Heap *findHeap(MEMiHeapHead *heapHandle);
    /* 80495730 */ Heap *findParentHeap();
    /* 80495780 */ static Heap *findContainHeap(const void *memBlock);
    /* 80495560 */ static void *alloc(size_t size, int align, Heap *heap);
    /* 804957c0 */ static void free(void *memBlock, Heap *heap);
    /* 80495830 */ void dispose();
    /* 804958a0 */ void dump();
    /* 804958b0 */ static void dumpAll();
    /* 804959a0 */ Heap *becomeCurrentHeap();
    /* 80495a00 */ Heap *_becomeCurrentHeapWithoutLock();

public:
    template <typename T>
    static T *alloc(u32 count, Heap *heap, int align = 4) {
        return reinterpret_cast<T *>(alloc(count * sizeof(T), align, heap));
    }
    template <typename T>
    static T *alloc(Heap *heap, int align = 4) {
        return reinterpret_cast<T *>(alloc(sizeof(T), align, heap));
    }

public:
    static void *addOffset(void *begin, u32 size) {
        return reinterpret_cast<char *>(begin) + size;
    }

    inline void appendDisposer(Disposer *disposer) {
        nw4r::ut::List_Append(&mChildren, disposer);
    }
    inline void removeDisposer(Disposer *disposer) {
        nw4r::ut::List_Remove(&mChildren, disposer);
    }

    inline MEMiHeapHead *getHeapHandle() {
        return mHeapHandle;
    }

    static inline Heap *getCurrentHeap() {
        return sCurrentHeap;
    }

    inline int getArenaEnd() {
        return (int)mHeapHandle->end;
    }

    inline const char *getName() {
        return mName;
    }

    /* 80673ae8 */ static nw4r::ut::List sHeapList;
    /* 80673af8 */ static OSMutex sRootMutex;
    /* 80576740 */ static Heap *sCurrentHeap;
    /* 80576744 */ static int sIsHeapListInitialized;
    /* 80576748 */ static Heap *sAllocatableHeap;
    /* 8057674c */ static ErrorCallback sErrorCallback;
    /* 80576750 */ static HeapAllocCallback sAllocCallback;
    /* 80576754 */ static HeapFreeCallback sFreeCallback;
    /* 80576758 */ static void *sErrorCallbackArg;
    /* 8057675c */ static void *sAllocCallbackArg;
    /* 80576760 */ static void *sFreeCallbackArg;
    /* 80576764 */ static HeapCreateCallback sCreateCallback;
    /* 80576764 */ static HeapDestroyCallback sDestroyCallback;
};
} // namespace EGG

/* 80495a60 */ void *operator new(size_t, void *p);
/* 80495a70 */ void *operator new(size_t size, EGG::Heap *heap, int align);
/* 80495a80 */ void *operator new(size_t size, EGG::Allocator *alloc);
/* 80495a90 */ void *operator new[](size_t size, int align);
/* 80495aa0 */ void *operator new[](size_t size, EGG::Heap *heap, int align);

#endif