#include <Base/Memory/MemoryBase.h>
#include <Base/Memory/ScratchAllocator.h>

class SetupMemoryHeapsUtil
{
    DefaultHeap         defaultHeap;
    MemTraceProxyHeap   globalHeap;
    MemTraceProxyHeap   processHeap;
    MemTraceProxyHeap   clumpsHeap;
    MemTraceProxyHeap   stringsHeap;
public:
    SetupMemoryHeapsUtil();
    ~SetupMemoryHeapsUtil();
};
class SetupScratchHeap : SetupMemoryHeapsUtil
{
    ScratchAllocator  scratchHeap;
    ThreadSafeProxyHeap threadSafeWrapper;
    AHeap * oldScratchHeap;
public:
    SetupScratchHeap( AHeap & _allocator, U32 sizeMiB );
    ~SetupScratchHeap();
};

/// A proxy allocator that simply draws upon another allocator.
/// http://bitsquid.blogspot.ru/2010/09/custom-memory-allocation-in-c.html
struct ProxyAllocator : AHeap, TLinkedList< ProxyAllocator >
{
    AHeap & m_allocator;
    const char * m_name;

    LONG    m_currently_allocated;  // doesn't take alignment into account
    LONG    m_peak_memory_usage;    // ditto
    LONG    m_total_allocations;
    LONG    m_total_frees;

    static ProxyAllocator::Head s_all;

public:
    ProxyAllocator( const char* _name, AHeap & _allocator )
        : m_allocator( _allocator ), m_name( _name )
    {
        _next = s_all;
        s_all = this;
    }
    ~ProxyAllocator()
    {
        mxASSERT(m_currently_allocated == 0);
    }
    virtual void* Allocate( U32 _bytes, U32 _alignment ) override;
    virtual void Deallocate( void* _memory ) override;
    virtual U32 allocated_size( const void* _memory ) const override;
};








SetupMemoryHeapsUtil::SetupMemoryHeapsUtil()
    : globalHeap( "global", defaultHeap )
    , processHeap( "process", defaultHeap )
    , clumpsHeap( "clumps", defaultHeap )
    , stringsHeap( "strings", defaultHeap )
{
    Memory_Initialize();

    for( int iHeap = 0; iHeap < g_heaps.num(); iHeap++ ) {
        g_heaps[ iHeap ] = &defaultHeap;
    }
    g_heaps[ HeapClumps ] = &clumpsHeap;
    g_heaps[ HeapString ] = &stringsHeap;
    //g_heaps[ HeapTemporary ] = &scratchHeap;
    g_heaps[ HeapProcess ] = &processHeap;
}

SetupMemoryHeapsUtil::~SetupMemoryHeapsUtil()
{
    Memory_Shutdown();
}

SetupScratchHeap::SetupScratchHeap( AHeap & _allocator, U32 sizeMiB )
    : scratchHeap( _allocator, mxMiB(sizeMiB), "scratch" )
    , threadSafeWrapper( scratchHeap )
{
    threadSafeWrapper.Initialize();
    oldScratchHeap = g_heaps[ HeapTemporary ];
    g_heaps[ HeapTemporary ] = &threadSafeWrapper;
}
SetupScratchHeap::~SetupScratchHeap()
{
    g_heaps[ HeapTemporary ] = oldScratchHeap;
    threadSafeWrapper.Shutdown();
}

//=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=

/*static*/ ProxyAllocator::Head ProxyAllocator::s_all = nil;

void* ProxyAllocator::Allocate( U32 _bytes, U32 _alignment ) {
    void * ptr = m_allocator.Allocate( _bytes, _alignment );
    if( ptr ) {
        U32 real_size = m_allocator.allocated_size( ptr );
        mxASSERT( real_size != SIZE_NOT_TRACKED );
        InterlockedIncrement( &m_total_allocations );
        LONG allocated_now = _InterlockedExchangeAdd( &m_currently_allocated, real_size );
        // interlocked max(a,b); NOTE: creates deadlocks on platforms without priority inversion
        LONG old_value, new_value;
        do {
            old_value = m_peak_memory_usage;
            new_value = largest( old_value, allocated_now );
        } while(
            _InterlockedCompareExchange( &m_peak_memory_usage, new_value, old_value )
            != old_value
        );
    }
    return ptr;
}
void ProxyAllocator::Deallocate( void* _memory ) {
    if( !_memory ) {
        return;
    }
    U32 size = m_allocator.allocated_size( _memory );
    mxASSERT( size != SIZE_NOT_TRACKED );

    _InterlockedExchangeAdd( &m_currently_allocated, -LONG(size) );
    InterlockedIncrement( &m_total_frees );

    m_allocator.Deallocate( _memory );
}
U32 ProxyAllocator::allocated_size( const void* _memory ) const {
    return m_allocator.allocated_size( _memory );
}