#pragma once

#include <cmath>
#include <carb/logging/Log.h>
#include <carb/Defines.h>
#include <carb/flatcache/Defines.h>
#include <carb/flatcache/Intrinsics.h>


#define ALLOCATOR_HEADER USE_IF( USING( DEVELOPMENT_BUILD ) )
#define ALLOCATOR_STATS USE_IF( USING( ALLOCATOR_HEADER ) ) // requires Header's byte tracking per-allocation
#define ALLOCATOR_LEAK_CHECK USE_IF( USING( ALLOCATOR_HEADER ) ) // requires Header's byte tracking per-allocation

namespace carb
{
namespace flatcache
{
inline const char* humanReadableSize( const uint64_t bytes ) noexcept
{
    auto va = [](auto ...params) -> const char* {
        static char tmp[1024];
#ifdef _WIN32
        _snprintf_s(tmp, sizeof(tmp), params...);
#else
        snprintf(tmp, sizeof(tmp), params...);
#endif
        return (const char*)&tmp;
    };

    constexpr const char SIZE_UNITS[64][3]{
        " B", " B", " B", " B", " B", " B", " B", " B", " B", " B",
        "KB", "KB", "KB", "KB", "KB", "KB", "KB", "KB", "KB", "KB",
        "MB", "MB", "MB", "MB", "MB", "MB", "MB", "MB", "MB", "MB",
        "GB", "GB", "GB", "GB", "GB", "GB", "GB", "GB", "GB", "GB",
        "TB", "TB", "TB", "TB", "TB", "TB", "TB", "TB", "TB", "TB",
        "PB", "PB", "PB", "PB", "PB", "PB", "PB", "PB", "PB", "PB",
        "EB", "EB", "EB", "EB"
    };
    static constexpr size_t B = 1ull;
    static constexpr size_t KB = 1024ull;
    static constexpr size_t MB = (1024ull*1024ull);
    static constexpr size_t GB = (1024ull*1024ull*1024ull);
    static constexpr size_t TB = (1024ull*1024ull*1024ull*1024ull);
    static constexpr size_t PB = (1024ull*1024ull*1024ull*1024ull*1024ull);
    static constexpr size_t EB = (1024ull*1024ull*1024ull*1024ull*1024ull*1024ull);
    constexpr const size_t SIZE_BASE[64]{
        B, B, B, B, B, B, B, B, B, B,
        KB, KB, KB, KB, KB, KB, KB, KB, KB, KB,
        MB, MB, MB, MB, MB, MB, MB, MB, MB, MB,
        GB, GB, GB, GB, GB, GB, GB, GB, GB, GB,
        TB, TB, TB, TB, TB, TB, TB, TB, TB, TB,
        PB, PB, PB, PB, PB, PB, PB, PB, PB, PB,
        EB, EB, EB, EB
    };
    const uint32_t power = bytes ? ( 64u - clz64( bytes ) ) - 1u : 0;
    const char *const units = SIZE_UNITS[power];
    const size_t base = SIZE_BASE[power];
    const size_t count = bytes / base;
    return va("%zu %s", count, units);
}


// A wrapper around malloc/free that aims to:
//
//   * Cheaply track allocation counts and bytes, and detect leaks automatically at ~Allocator()
//
//   * Cheaply track usage in terms of peak memory usage, and total lifetime usage broken down by size. Sample output:
//     dumped to console appears like so:
//
//     == Allocator 0x000000E67BEFCEA0 Stats ==
//     allocCount:                0
//     allocBytes:                0  B
//     peakAllocCount:         4002
//     peakAllocBytes:            4 GB
//     minAllocBytes:           312  B
//     maxAllocBytes:             6 MB
//     
//     Lifetime Allocation Histogram:
//       Normalized over TOTAL allocations: 13956
//       < 512  B|*****             29%         4002
//       <   1 KB|                   0%            0
//       <   2 KB|                   0%            0
//       <   4 KB|                   0%            0
//       <   8 KB|***               14%         2000
//       <  16 KB|                   0%            0
//       <  32 KB|***               14%         1994
//       <  64 KB|                   0%            0
//       < 128 KB|***               14%         1976
//       < 256 KB|                   0%            0
//       < 512 KB|***               14%         1904
//       <   1 MB|                   0%            0
//       <   2 MB|**                12%         1616
//       <   4 MB|                   0%            0
//       <   8 MB|*                  3%          464
//     ========================
struct Allocator
{
    Allocator();
    ~Allocator();

    void* alloc(const size_t bytes);
    void free(void *const ptr);

    template<typename T, typename ...Params>
    T* new_(Params&& ...params);

    template<typename T>
    void delete_(T*const t);

    void resetStats() noexcept;
    void reportUsage() noexcept;

    bool checkLeaks() noexcept;

private:
#if USING( ALLOCATOR_HEADER )
    struct BlockHeader
    {
        size_t bytes;
    };
#endif // #if USING( ALLOCATOR_HEADER )

#if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )
    size_t allocCount;
    size_t allocBytes;
#endif // #if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )

#if USING( ALLOCATOR_STATS )
    size_t peakAllocCount;
    size_t peakAllocBytes;
    size_t minAllocBytes;
    size_t maxAllocBytes;

    static constexpr size_t ALLOC_BUCKET_COUNT = 65;
    size_t lifetimeAllocCount;
    size_t lifetimeAllocBuckets[ALLOC_BUCKET_COUNT];
#endif // #if USING( ALLOCATOR_STATS )
};

struct AllocFunctor
{
    Allocator *allocator;

    void* operator()(const size_t bytes)
    {
        CARB_ASSERT(allocator);
        return allocator->alloc(bytes);
    }
};

struct FreeFunctor
{
    Allocator *allocator;

    void operator()(void *const ptr)
    {
        CARB_ASSERT(allocator);
        return allocator->free(ptr);
    }
};

inline Allocator::Allocator()
{
#if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )
    allocCount = 0;
    allocBytes = 0;
#endif // #if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )

    resetStats();
}


inline Allocator::~Allocator()
{
    checkLeaks();
    reportUsage();
}


inline void* Allocator::alloc(const size_t bytes)
{
#if USING( ALLOCATOR_HEADER )
    const size_t totalBytes = bytes + sizeof(BlockHeader);
#endif // #if USING( ALLOCATOR_HEADER )

#if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )
    CARB_ASSERT((allocCount + 1) > allocCount);
    CARB_ASSERT((allocBytes + totalBytes) > allocBytes);
    ++allocCount;
    allocBytes += totalBytes;
#endif // #if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )

#if USING( ALLOCATOR_STATS )
    if ( allocBytes > peakAllocBytes )
    {
        peakAllocBytes = allocBytes;
        peakAllocCount = allocCount;
    }
    if ( totalBytes < minAllocBytes )
    {
        minAllocBytes = totalBytes;
    }
    if ( totalBytes > maxAllocBytes )
    {
        maxAllocBytes = totalBytes;
    }
    const uint32_t bucket = ( 64u - clz64( totalBytes - 1ull ) );
    CARB_ASSERT(lifetimeAllocBuckets[bucket] + 1 > lifetimeAllocBuckets[bucket]);
    ++lifetimeAllocBuckets[bucket];
    ++lifetimeAllocCount;
#endif // #if USING( ALLOCATOR_STATS )

#if USING( ALLOCATOR_HEADER )
    BlockHeader *const header = (BlockHeader*)malloc(totalBytes);
    CARB_ASSERT(header);
    header->bytes = totalBytes;
    return header+1;
#else // #if USING( ALLOCATOR_HEADER )
    return malloc(bytes);
#endif // #if USING( ALLOCATOR_STATS )
}

inline void Allocator::free(void *const ptr)
{
#if USING( ALLOCATOR_HEADER )
    CARB_ASSERT(ptr);
    BlockHeader *header = (BlockHeader*)ptr;
    --header;

#if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )
    const size_t totalBytes = header->bytes;

    CARB_ASSERT((allocCount - 1) < allocCount);
    CARB_ASSERT((allocBytes - totalBytes) < allocBytes);
    --allocCount;
    allocBytes -= totalBytes;
#endif // #if USING( ALLOCATOR_STATS ) || USING( ALLOCATOR_LEAK_CHECK )

    ::free(header);
#else // #if USING( ALLOCATOR_STATS )
    ::free(ptr);
#endif // #if USING( ALLOCATOR_STATS )
}


template<typename T, typename ...Params>
inline T* Allocator::new_(Params&& ...params)
{
    T *const t = (T*)Allocator::alloc(sizeof(T));
    new (t) T(std::forward<Params>(params)...);
    return t;
}


template<typename T>
inline void Allocator::delete_(T*const t)
{
    CARB_ASSERT(t);
    t->~T();
#if USING( ALLOCATOR_HEADER )
    BlockHeader *header = (BlockHeader*)t;
    header--;
    CARB_ASSERT(header->bytes == (sizeof(BlockHeader) + sizeof(T)));
#endif // #if USING( ALLOCATOR_HEADER )
    Allocator::free(t);
}


inline void Allocator::resetStats() noexcept
{
#if USING( ALLOCATOR_STATS )
    peakAllocCount = 0;
    peakAllocBytes = 0;
    minAllocBytes = SIZE_MAX;
    maxAllocBytes = 0;
    lifetimeAllocCount = 0;
    for ( size_t i = 0; i < ALLOC_BUCKET_COUNT; ++i )
    {
        lifetimeAllocBuckets[i] = 0;
    }
#endif // #if USING( ALLOCATOR_STATS )
}


inline void Allocator::reportUsage() noexcept
{
#if USING( ALLOCATOR_STATS )
    CARB_LOG_INFO("== Allocator 0x%p Stats ==", this);

    if (!lifetimeAllocCount)
    {
        CARB_LOG_INFO("<no stats to report; unused allocator>");
        CARB_LOG_INFO("========================");
        return;
    }

    CARB_LOG_INFO("allocCount:     %12zu", allocCount);
    CARB_LOG_INFO("allocBytes:     %15s", humanReadableSize(allocBytes));
    CARB_LOG_INFO("peakAllocCount: %12zu", peakAllocCount);
    CARB_LOG_INFO("peakAllocBytes: %15s", humanReadableSize(peakAllocBytes));
    CARB_LOG_INFO("minAllocBytes:  %15s", humanReadableSize(minAllocBytes));
    CARB_LOG_INFO("maxAllocBytes:  %15s", humanReadableSize(maxAllocBytes));
    CARB_LOG_INFO("");
    CARB_LOG_INFO("Lifetime Allocation Histogram:");

    size_t begin = 0;
    for ( ; begin < ALLOC_BUCKET_COUNT; ++begin )
    {
        if ( lifetimeAllocBuckets[begin] )
        {
            break;
        }
    }
    size_t end = 0;
    for ( ; end < ALLOC_BUCKET_COUNT; ++end )
    {
        if ( lifetimeAllocBuckets[ALLOC_BUCKET_COUNT - end - 1] )
        {
            end = ALLOC_BUCKET_COUNT - end;
            break;
        }
    }

    CARB_LOG_INFO("  Normalized over TOTAL allocations: %zu", lifetimeAllocCount);
    size_t i;
    float normalized[ALLOC_BUCKET_COUNT];
    for ( i = begin; i < end; ++i )
    {
        normalized[i] = (float)lifetimeAllocBuckets[i] / (float)lifetimeAllocCount;
    }

    constexpr size_t WIDTH = 16;
    for ( i = begin; i < end; ++i )
    {
        char buf[WIDTH+1] = {};
        const size_t w = ( size_t )std::ceil(WIDTH * normalized[i]);
        for( size_t j = 0; j < w; ++j)
        {
            buf[j] = '*';
        }
        static_assert(WIDTH == 16, "Fix CARB_LOG_INFO below");
        CARB_LOG_INFO("  <%7s|%-16s %3.0f%% %12zu", humanReadableSize(1ull<<i), buf, (normalized[i] * 100.f), lifetimeAllocBuckets[i]);
    }

    CARB_LOG_INFO("========================");
#endif // #if USING( ALLOCATOR_STATS )
}


inline bool Allocator::checkLeaks() noexcept
{
#if USING( ALLOCATOR_LEAK_CHECK )
    if (allocCount || allocBytes)
    {
        CARB_LOG_ERROR("PathToAttributesMap detected a memory leak of %s!\n", humanReadableSize(allocBytes));
        CARB_ASSERT(false, "PathToAttributesMap detected a memory leak of %s!\n", humanReadableSize(allocBytes));
        return true;
    }
#endif // #if USING( ALLOCATOR_LEAK_CHECK )

    return false;
}

} // namespace flatcache

} // namespace carb
