#include <list>
#include <map>
#include <stdarg.h>
#include <string>
#include <time.h>

#include "KString.h"

#include "KMemory.h" // last include

#pragma warning(disable:4996)

///////////////////////////////////////////////////////////////////////////////
//
//	KANJI_DUMP_LEAKED_MEM will print out the memory block that was leaked.
//	This is helpful when leaked objects have string identifiers.
//
///////////////////////////////////////////////////////////////////////////////
//#define KANJI_DUMP_LEAKED_MEM


///////////////////////////////////////////////////////////////////////////////
//
//	KANJI_TRACK_MEM_USAGE will print out all memory allocations.
//
///////////////////////////////////////////////////////////////////////////////
//#define KANJI_TRACK_MEM_USAGE



///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
// Our memory system is thread-safe, but instead of linking massive libraries,
// we attempt to use C++11 std::mutex.
#ifdef USE_CPP11_MUTEX_DISABLED
#include <mutex>
typedef std::recursive_mutex KSysLock; // rentrant
struct KAutoLock {
	KAutoLock(KSysLock& lock) :mLock(lock) { mLock.lock(); }	// acquire 
	~KAutoLock() { mLock.unlock(); }							// release

	KSysLock& mLock;
};
#else // Fallback to unsafe.  don't spawn threads
typedef int KSysLock;
struct KAutoLock {
	KAutoLock(KSysLock) {}	// acquire 
	~KAutoLock() {}			// release
};
#endif



///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
struct KANJI_ALLOC_INFO
{
    size_t  size;
    std::string file;
    int     line;
};
static bool showLeaks = false;
class KAllocMap : public std::map<void*, KANJI_ALLOC_INFO>
{
public:
    KSysLock crit;
    static bool allocMapValid;

public:
    KAllocMap() { allocMapValid = true; }
    ~KAllocMap()
    {
        if (showLeaks)
            KMemoryDumpUnfreed();

        allocMapValid = false;
    }
};
bool KAllocMap::allocMapValid = false;
static KAllocMap allocMap; // once this static object destructs, it dumps unfreed memory


///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
#ifdef KANJI_TRACK_MEM_USAGE
void KMemoryDumpUsage(); // forward declaration
class KAllocStat
{
public:
    typedef std::map<int, int> allocCount; // [size] = count
    typedef std::map<std::pair<std::string, int>, allocCount> allocInfo; // [file, line] = allocCount

    allocInfo memInfo;
    static bool allocMapValid;

public:

    KAllocStat()
    {
        allocMapValid = true;
    }
    ~KAllocStat()
    {
        if (showLeaks)
            KMemoryDumpUsage();

        allocMapValid = false;
    }
    void addTrack(const char* fname, int lnum, int asize)
    {
        allocCount& info = memInfo[std::pair<std::string, int>(fname, lnum)];
        info[asize]++;
    }
};
bool KAllocStat::allocMapValid = false;
static KAllocStat allocStat;
#endif // KANJI_TRACK_MEM_USAGE



///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
extern void KMemoryAddTrack( void* addr, size_t asize, const char* fname, int lnum )
{
    if (!KAllocMap::allocMapValid || asize == 0)
        return;

    KAutoLock aCrit(allocMap.crit);
    showLeaks = true;

    KANJI_ALLOC_INFO& info = allocMap[addr];
    info.file = fname;
    info.line = lnum;
    info.size = asize;

#ifdef KANJI_TRACK_MEM_USAGE
    if (KAllocStat::allocMapValid)
        allocStat.addTrack(fname, lnum, asize);
#endif
};

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
void KMemoryRemoveTrack(void* addr)
{
    if (!KAllocMap::allocMapValid)
        return;

    KAutoLock aCrit(allocMap.crit);
    KAllocMap::iterator anItr = allocMap.find(addr);
    if (anItr != allocMap.end())
        allocMap.erase(anItr);
};

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
void KMemoryDumpUnfreed()
{
    if (!KAllocMap::allocMapValid)
        return;

    KAutoLock aCrit(allocMap.crit); // prevent modification of the map while iterating

    size_t totalSize = 0;
    char buf[8192];

    FILE* f = fopen("mem_leaks.txt", "wt");
    if (!f)
        return;

    time_t aTime = time(NULL);
    sprintf(buf, "Memory Leak Report for %s\n", asctime(localtime(&aTime)));
    fprintf(f, "%s", buf);
    KOutputDebug(DEBUGLVL, "\n");
    KOutputDebug(INFOLVL, buf);
    for(KAllocMap::iterator i = allocMap.begin(); i != allocMap.end(); ++i)
    {
        sprintf(buf, "%s(%d) : Leak %u byte%s @0x%08X\n", i->second.file.c_str(), i->second.line, i->second.size, i->second.size > 1 ? "s" : "", (size_t)i->first);
        KOutputDebug(ERRORLVL, buf);
        fprintf(f, "%s", buf);

#ifdef KANJI_DUMP_LEAKED_MEM
        unsigned char* data = (unsigned char*)i->first;
        int count = 0;
        char hex_dump[1024];
        char ascii_dump[1024];

        for (int index = 0; index < i->second.size; index++)
        {
            unsigned char _c = *data;

            if (count == 0)
                sprintf(hex_dump, "\t%02X ", _c);
            else
                sprintf(hex_dump, "%s%02X ", hex_dump, _c); // technically, this is undefined behavior

            if ((_c < 32) || (_c > 126))
                _c = '.';

            if (count == 7)
                sprintf(ascii_dump, "%s%c ", ascii_dump, _c);
            else
                sprintf(ascii_dump, "%s%c", count == 0 ? "\t" : ascii_dump, _c); // technically, this is undefined behavior


            if (++count == 16)
            {
                count = 0;
                sprintf(buf, "%s\t%s\n", hex_dump, ascii_dump);
                fprintf(f, buf);

                memset((void*)hex_dump, 0, 1024);
                memset((void*)ascii_dump, 0, 1024);
            }

            data++;
        }

        if (count != 0)
        {
            fprintf(f, hex_dump);
            for (int index = 0; index < 16 - count; index++)
                fprintf(f, "\t");

            fprintf(f, ascii_dump);

            for (int index = 0; index < 16 - count; index++)
                fprintf(f, ".");
        }

        count = 0;
        fprintf(f, "\n\n");
        memset((void*)hex_dump, 0, 1024);
        memset((void*)ascii_dump, 0, 1024);

#endif // KANJI_DUMP_LEAKED_MEM

        totalSize += i->second.size;
    }

	ErrorLevel lvl = (totalSize > 0) ? ERRORLVL : INFOLVL;

    sprintf(buf, "-----------------------------------------------------------\n");
    fprintf(f, "%s", buf);
    KOutputDebug(lvl, buf);
    sprintf(buf, "Total Unfreed: %u bytes (%luKB)\n\n", totalSize, totalSize / 1024);
    KOutputDebug(lvl, buf);
    fprintf(f, "%s", buf);
    fclose(f);
}

#ifdef KANJI_TRACK_MEM_USAGE
void KMemoryDumpUsage()
{
    if (!KAllocStat::allocMapValid)
        return;

    char buf[8192];
    FILE* f = fopen("mem_usage.txt", "wt");

    time_t aTime = time(NULL);
    sprintf(buf, "Memory Usage Report for %s\n", asctime(localtime(&aTime)));
    if (f) fprintf(f, "%s", buf);
    KOutputDebug("\n");
    KOutputDebug(buf);

    for(KAllocStat::allocInfo::iterator i = allocStat.memInfo.begin(); i != allocStat.memInfo.end(); ++i)
    {
        int aBytesTotal = 0;
        int aCallsTotal = 0;
        for (KAllocStat::allocCount::iterator index = i->second.begin(); index != i->second.end(); ++index)
        {
            aBytesTotal += index->first;
            aCallsTotal += index->second;
            sprintf(buf, "%s(%d) : %d bytes (%d %s)\n", i->first.first.c_str(), i->first.second, index->first, index->second, index->second == 1 ? "call" : "calls");
            if (f) fprintf(f, "%s", buf);
            KOutputDebug(buf);
        }

        if (i->second.size() > 1)
        {
            sprintf(buf, "    %s(%d) : %d KB total (%d calls)\n", i->first.first.c_str(), i->first.second, aBytesTotal / 1024, aCallsTotal);
            if (f) fprintf(f, "%s", buf);
            KOutputDebug(buf);
        }
    }
    if (f) fclose(f);
}
#endif // KANJI_TRACK_MEM_USAGE

size_t KMemoryAllocated()
{
    if (!KAllocMap::allocMapValid)
        return 0;

    KAutoLock aCrit(allocMap.crit);

    size_t size = 0;
    for(auto i = allocMap.begin(); i != allocMap.end(); ++i)
    {
        KANJI_ALLOC_INFO& info = i->second;
        size += info.size;
    }
    return size;
}