#include "traster.h"
#include "tbigmemorymanager.h"
#include "timagecache.h"
#include "tsystem.h"
#include "tconvert.h"
#include <set>
#include "tfilepath_io.h"

#ifdef _DEBUG
std::set<TRaster *> Rasters;
#endif

#ifdef TNZCORE_LIGHT
#ifdef MessageBox
#undef MessageBox
#endif

#define MessageBox MessageBoxA
#endif

class Chunkinfo
{
public:
	TUINT32 m_size;
	//int m_locks;
	std::vector<TRaster *> m_rasters;
	//bool m_putInNormalMemory;
	Chunkinfo(TUINT32 size, TRaster *ras) //, bool putInNormalMemory=false)
		: m_size(size)
		  //, m_locks(0)
		  ,
		  m_rasters()
	//, m_putInNormalMemory(putInNormalMemory)
	{
		if (ras)
			m_rasters.push_back(ras);
	}

	Chunkinfo()
		: m_size(0)
		  //, m_locks(0)
		  ,
		  m_rasters()
	/*, m_putInNormalMemory(false)*/ {}
};

//------------------------------------------------------------

//! Sets the global callback handler for the 'Run out of contiguous memory' event
//! The callback receives the size (in bytes) of the raster which caused the problem.
void TBigMemoryManager::setRunOutOfContiguousMemoryHandler(void (*callback)(unsigned long))
{
	m_runOutCallback = callback;
}

//------------------------------------------------------------------------------

namespace
{
int allocationPeakKB = 0;
unsigned long long allocationSumKB = 0;
unsigned long allocationCount = 0;
}

//------------------------------------------------------------------------------

//! Returns the \b peak size, in KB, of the allocated rasters in current Toonz session.
int TBigMemoryManager::getAllocationPeak()
{
	return allocationPeakKB;
}

//------------------------------------------------------------------------------

//! Returns the \b mean size, in KB, of the allocated rasters in current Toonz session.
int TBigMemoryManager::getAllocationMean()
{
	return allocationSumKB / allocationCount;
}

//------------------------------------------------------------------------------

TBigMemoryManager *TBigMemoryManager::instance()
{
	static TBigMemoryManager *theManager = 0;

	if (theManager)
		return theManager;

	return theManager = new TBigMemoryManager();
}

//------------------------------------------------------------------------------
/*
TBigMemoryManager::~TBigMemoryManager()
{

if (m_theMemory==0) return;
QMutexLocker sl(m_mutex);
assert(m_chunks.empty());

free(m_theMemory);
theManager = 0;
}
*/

//------------------------------------------------------------------------------

UCHAR *TBigMemoryManager::allocate(UINT &size)
{
	TThread::MutexLocker sl(&m_mutex);
	UCHAR *chunk = (UCHAR *)calloc(size, 1);
	while (chunk == 0 && size > 128 * 1024 * 1024) {
		size -= 128 * 1024 * 1024;
		chunk = (UCHAR *)calloc(size, 1);
	}
	return chunk;
}

//------------------------------------------------------------------------------

TBigMemoryManager::TBigMemoryManager()
	: m_chunks(), m_theMemory(0), m_availableMemory(0), m_allocatedMemory(0)
#ifdef _DEBUG
	  ,
	  m_totRasterMemInKb(0)
#endif
{
}

//------------------------------------------------------------------------------

bool TBigMemoryManager::init(TUINT32 sizeinKb)
{
	TThread::MutexLocker sl(&m_mutex);

	if (sizeinKb == 0)
		return true;

	if (sizeinKb >= 2 * 1024 * 1024) {
		//  MessageBox( NULL, "TRONCO!!!", "Warning", MB_OK);
		sizeinKb = (TUINT32)(1.8 * 1024 * 1024);
	}

	m_availableMemory = sizeinKb * 1024;

	m_theMemory = allocate(m_availableMemory);
	m_allocatedMemory = m_availableMemory;

	if (!m_theMemory) {
		//  MessageBox( NULL, "Ouch!can't allocate Big Chunk!", "Warning", MB_OK);
		m_theMemory = 0;
		m_availableMemory = 0;
		return false;
	}
	//char str[1024];
	//sprintf_s(str, "chiesto %d MB, allocato un big chunk di %d MB", sizeinKb/1024, m_availableMemory/(1024*1024));
	//MessageBox( NULL, str, "Warning", MB_OK);

	m_chunks[m_theMemory + m_availableMemory] = Chunkinfo(0, 0);
	return true;
}

//------------------------------------------------------------------------------
/*
void TBigMemoryManager::lock(UCHAR *buffer)
{
QMutexLocker sl(m_mutex);
if (m_theMemory==0) return;
assert(buffer);
map<UCHAR*, Chunkinfo>::iterator it = m_chunks.find(buffer);

if (it==m_chunks.end()) return;

++it->second.m_locks;
}
*/
//------------------------------------------------------------------------------

/*
void TBigMemoryManager::unlock(UCHAR *buffer)
{
QMutexLocker sl(m_mutex);
if (m_theMemory==0) return;
assert(buffer);
map<UCHAR*, Chunkinfo>::iterator it = m_chunks.find(buffer);

if (it==m_chunks.end()) return;
int locks = it->second.m_locks;

assert(locks>0);

--it->second.m_locks;
}
*/

//------------------------------------------------------------------------------

UCHAR *TBigMemoryManager::getBuffer(UINT size)
{
	if (m_theMemory == 0)
		return (UCHAR *)calloc(size, 1);

	std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.begin();
	UCHAR *buffer = m_theMemory;
	TUINT32 chunkSize = 0;
	UCHAR *address = 0;
	while (it != m_chunks.end()) {
		/*if (it->second.m_putInNormalMemory)
    {it++; continue;}*/

		if ((TUINT32)((it->first) - (buffer + chunkSize)) >= size) {
			address = buffer + chunkSize;
			break;
		}

		buffer = it->first;
		chunkSize = it->second.m_size;
		it++;
	}
	if (address)
		memset(address, 0x00, size);
	return address;
}

//---------------------------------------------------------------------------------

#ifdef _DEBUG

void TBigMemoryManager::getRasterInfo(int &rasterCount, TUINT32 &totRasterMemInKb, int &notCachedRasterCount, TUINT32 &notCachedRasterMemInKb)
{
	totRasterMemInKb = 0;
	notCachedRasterMemInKb = 0;
	notCachedRasterCount = 0;
	rasterCount = Rasters.size();

	std::set<TRaster *>::iterator it = Rasters.begin();
	while (it != Rasters.end()) {
		TRaster *r = *it;
		assert(r->m_parent == 0);
		totRasterMemInKb += r->getLx() * r->getLy() * r->getPixelSize() >> 10;
		if (!(*it)->m_cashed && r->getLy() > 1) {
			notCachedRasterCount++;
			notCachedRasterMemInKb += r->getLx() * r->getLy() * r->getPixelSize() >> 10;
		}
		++it;
	}
}

#endif

bool TBigMemoryManager::putRaster(TRaster *ras, bool canPutOnDisk)
{

	if (!ras->m_parent && ras->m_buffer) {
#ifdef _DEBUG
		if (ras->m_bufferOwner)
			Rasters.insert(ras);
#endif
		return true;
	}

#ifdef _DEBUG
	checkConsistency();
#endif

	TUINT32 size = ras->getLx() * ras->getLy() * ras->getPixelSize();

	if (size == 0) {
		ras->m_buffer = 0;
		return true;
	}

	if (m_theMemory == 0) //il bigmemorymanager e' inattivo
	{
		if (!ras->m_parent) {
			int sizeKB = size >> 10;
			allocationPeakKB = std::max(allocationPeakKB, sizeKB);
			allocationSumKB += sizeKB;
			allocationCount++;
		}

		if (!ras->m_parent && !(ras->m_buffer = (UCHAR *)calloc(size, 1))) {
			//MessageBox( NULL, "Ouch!can't allocate!", "Warning", MB_OK);
			//non c'e' memoria; provo a comprimere
			/*TImageCache::instance()->doCompress(); 
    if (!(ras->m_buffer = (UCHAR *)malloc(size)))*/ //andata male pure cosi'; metto tutto su disco
			//TImageCache::instance()->outputMap(size, "C:\\logCacheFailure");
			TINT64 availMemInKb = TSystem::getFreeMemorySize(true);
			if (availMemInKb > (size >> 10)) {
				//char str[1024];
				//sprintf_s(str, "Non alloco, ma : memoria (KB) richiesta : %d - memoria disponibile : %d", size>>10, availMemInKb);
				//MessageBox( NULL, (LPCSTR)str, (LPCSTR)"Segmentation!", MB_OK);
			}
			ras->m_buffer = TImageCache::instance()->compressAndMalloc(size);

			if (!ras->m_buffer) {
				//char str[1024];
				//sprintf_s(str, "E' andata male: faccio il log della cache.");
				//MessageBox( NULL, (LPCSTR)str, (LPCSTR)"Segmentation!", MB_OK);
				TImageCache::instance()->outputMap(size, "C:\\logCacheTotalFailure");
			} else {
#ifdef _DEBUG
				m_totRasterMemInKb += size >> 10;
				Rasters.insert(ras);
#endif
			}
			return ras->m_buffer != 0;
		} else {
			if (!ras->m_parent) {
#ifdef _DEBUG
				m_totRasterMemInKb += size >> 10;
				Rasters.insert(ras);
#endif
			}
			return true;
		}
	}

	//il bigmemorymanager e' attivo
	TThread::MutexLocker sl(&m_mutex);

	/*
if (m_availableMemory<size && !ras->m_parent)
  {
  TImageCache::instance()->compressAndMalloc(size);
  if (m_availableMemory>=size)
    return TBigMemoryManager::instance()->putRaster(ras);
  else
   return (ras->m_buffer = (UCHAR *)malloc(size))!=0;
  
    
  }*/

	if (ras->m_parent) {
		std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.find(ras->m_parent->m_buffer);

		if (it != m_chunks.end()) {

#ifdef _DEBUG
//assert(!it->second.m_rasters.empty());
//for (UINT i=0; i<it->second.m_rasters.size(); i++)
//  assert (it->second.m_rasters[i]!=ras);
#endif

			it->second.m_rasters.push_back(ras);
		}
#ifdef _DEBUG
		checkConsistency();
#endif
		return true;
	}

	UCHAR *address = 0;
	bool remapped = false;

	assert(m_chunks.size() > 0);

	/*if (m_chunks.size()==1) //c'e' solo l'elemento che marca la fine del bufferone
  {
  if ((TUINT32)(m_chunks.begin()->first-m_theMemory)>=size)
    address = m_theMemory;
  }
else
  {*/

	address = getBuffer(size);

#ifdef _DEBUG
	checkConsistency();
#endif

	if (address == 0 && m_availableMemory >= size) {
		address = remap(size);
		remapped = true;
#ifdef _DEBUG
		checkConsistency();
#endif
	} else if (address == 0) {
		printLog(size);
		//assert(!"la bigmemory e' piena...scritto log");
	}
	//}
	if (address == 0) {
		if (canPutOnDisk)
			address = TImageCache::instance()->compressAndMalloc(size);
		if (address == 0)
			return (ras->m_buffer = (UCHAR *)calloc(size, 1)) != 0;
	}

	//assert(address);

	ras->m_buffer = address;
	m_chunks[address] = Chunkinfo(size, ras);
	assert(m_availableMemory >= size);
	m_availableMemory -= size;
#ifdef _DEBUG
	checkConsistency();
#endif
	return true;
}

//------------------------------------------------------------------------------

TRaster *TBigMemoryManager::findRaster(TRaster *ras)
{
	//return 0;

	std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.begin();
	while (it != m_chunks.end()) {
		for (UINT i = 0; i < it->second.m_rasters.size(); i++)
			if (it->second.m_rasters[i] == ras)
				return ras;
		it++;
	}
	return 0;
}

//------------------------------------------------------------------------------
#ifdef _DEBUG

void TBigMemoryManager::printMap()
{
	std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.begin();
	TSystem::outputDebug("BIGMEMORY chunks totali: " + std::to_string((int)m_chunks.size()) + "\n");

	int count = 0;

	while (it != m_chunks.end()) {
		TSystem::outputDebug("chunk #" +
							 std::to_string((int)count++) +
							 "dimensione(kb):" +
							 std::to_string((int)(it->second.m_size >> 10)) +
							 "num raster:" +
							 std::to_string((int)(it->second.m_rasters.size())) +
							 "\n");

		it++;
	}
}

#endif

bool TBigMemoryManager::releaseRaster(TRaster *ras)
{
	TThread::MutexLocker sl(&m_mutex);
	UCHAR *buffer = (ras->m_parent) ? (ras->m_parent->m_buffer) : (ras->m_buffer);
	std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.find(buffer);

	if (m_theMemory == 0 || it == m_chunks.end()) {
		assert(buffer);
		if (!ras->m_parent && ras->m_bufferOwner) {
			free(buffer);
#ifdef _DEBUG
			m_totRasterMemInKb -= (ras->getPixelSize() * ras->getLx() * ras->getLy()) >> 10;
			Rasters.erase(ras);
#endif
		}

		//assert(findRaster(ras)==0);

		return false;
	}

	assert(ras->m_lockCount == 0);

	if (it->second.m_rasters.size() > 1) //non e' il solo raster ad usare il buffer; non libero
	{
		std::vector<TRaster *>::iterator it2 = it->second.m_rasters.begin();
		for (; it2 != it->second.m_rasters.end(); ++it2) {
			if (ras == *it2) {
				it->second.m_rasters.erase(it2);

#ifdef _DEBUG
//assert(!it->second.m_rasters.empty());
//assert(findRaster(ras)==0);
#endif

				return true;
			}
		}
		assert(false);
		return false;
	} else if (ras->m_bufferOwner) //libero!
	{
#ifdef _DEBUG
		checkConsistency();
#endif
		/*if (it->second.m_putInNormalMemory && ras->m_bufferOwner)
		free(it->first);
	else */
		m_availableMemory += it->second.m_size;
		m_chunks.erase(it);
	}

#ifdef _DEBUG
	//assert(findRaster(ras)==0);
	checkConsistency();
#endif

	return true;
}
//------------------------------------------------------------------------------

void TBigMemoryManager::checkConsistency()
{
	return;

	//QMutexLocker sl(m_mutex);

	int count = 0;
	//int size = m_chunks.size();
	std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.begin();
	UCHAR *endAddress = m_theMemory;
	TUINT32 freeMem = 0, allocMem = 0;

	while (it != m_chunks.end()) {
		count++;
		//assert(it->second.m_rasters.size()==0 || it->second.m_rasters.size()>0);

		if (endAddress != 0 /*&& !it->second.m_putInNormalMemory*/) {
			freeMem += (TUINT32)(it->first - endAddress);
			allocMem += it->second.m_size;
		}
		assert(endAddress <= it->first);
		endAddress = it->first + it->second.m_size;
		for (UINT i = 0; i < it->second.m_rasters.size(); i++) {

			TRaster *ras = it->second.m_rasters[i];
			it->second.m_rasters[i] = 0; //ogni raster deve apparire una sola volta
			assert(findRaster(ras) == 0);
			it->second.m_rasters[i] = ras;

			UCHAR *buf1 = (ras->m_parent) ? ras->m_parent->m_buffer : ras->m_buffer;
			UCHAR *buf2 = it->first;
			assert(buf1 == buf2);
			UINT size;
			if (ras->m_parent)
				size = ras->m_parent->getLx() * ras->m_parent->getLy() * ras->m_parent->getPixelSize();
			else
				size = ras->getLx() * ras->getLy() * ras->getPixelSize();
			assert(size == it->second.m_size);
		}
		it++;
	}

	if (m_theMemory) {
		assert(allocMem + m_availableMemory == m_allocatedMemory);

		assert(freeMem == m_availableMemory);
	}
}

//------------------------------------------------------------------------------

std::map<UCHAR *, Chunkinfo>::iterator TBigMemoryManager::shiftBlock(const std::map<UCHAR *, Chunkinfo>::iterator &it, TUINT32 offset)
{
	UCHAR *newAddress = it->first - offset;

	if (offset > it->second.m_size)
		memcpy(newAddress, it->first, it->second.m_size); //se NON overlappano.
	else
		memmove(newAddress, it->first, it->second.m_size); //se overlappano.

	m_chunks[newAddress] = Chunkinfo(it->second.m_size, it->second.m_rasters[0]);
	std::map<UCHAR *, Chunkinfo>::iterator it1 = m_chunks.find(newAddress);

	assert(it1->first < it1->second.m_rasters[0]->m_buffer);
	UINT i = 0;
	for (i = 0; i < it->second.m_rasters.size(); i++) //prima rimappo i subraster, senza toccare il buffer del parent...
	{
		TRaster *ras = it->second.m_rasters[i];
		assert(i > 0 || !ras->m_parent);
		if (!ras->m_parent)
			continue;
		assert(ras->m_parent->m_buffer == it->first);
		ras->remap(newAddress);
		if (i > 0)
			it1->second.m_rasters.push_back(ras);
	}

	it->second.m_rasters[0]->remap(newAddress); //ORA rimappo il parent

#ifdef _DEBUG
	for (i = 1; i < it->second.m_rasters.size(); i++) //..poi i raster padri
	{
		//TRaster*ras = it->second.m_rasters[i];

		assert(it->second.m_rasters[i]->m_parent);
		//ras->remap(newAddress);
		//if (i>0)
		//  it1->second.m_rasters.push_back(ras);
	}
	assert(it1->second.m_rasters.size() == it->second.m_rasters.size());
#endif

	m_chunks.erase(it);
	it1 = m_chunks.find(newAddress); //non dovrebbe servire, ma per prudenza...
	assert(it1->first == it1->second.m_rasters[0]->m_buffer);
	return it1;
}

//------------------------------------------------------------------------------

UCHAR *TBigMemoryManager::remap(TUINT32 size) //size==0 -> remappo tutto
{
	bool locked = false;
//QMutexLocker sl(m_mutex); //gia' scopata

#ifdef _DEBUG
	checkConsistency();
#endif
	UINT i;
	std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.begin();

	try {
		UCHAR *buffer = m_theMemory;
		TUINT32 chunkSize = 0;
		while (it != m_chunks.end()) {
			/*if (it->second.m_putInNormalMemory)
      {it++; continue;}*/

			TUINT32 gap = (TUINT32)((it->first) - (buffer + chunkSize));
			if (size > 0 && gap >= size) //trovato chunk sufficiente
				return buffer + chunkSize;
			else if (gap > 0 && it->second.m_size > 0) //c'e' un frammento di memoria, accorpo; ma solo se non sto in fondo
			{
				std::vector<TRaster *> &rasters = it->second.m_rasters;
				assert(rasters[0]->m_parent == 0);

				//devo controllare il lockCount solo sul parent, la funzione lock() locka solo il parent;
				for (i = 0; i < rasters.size(); i++)
					rasters[i]->beginRemapping();

				if (rasters[0]->m_lockCount == 0)
					it = shiftBlock(it, gap);
				else
					locked = true;

				for (i = 0; i < rasters.size(); i++)
					rasters[i]->endRemapping();

				//rasters.clear();
			}

			buffer = it->first;
			chunkSize = it->second.m_size;
			it++;
		}

	} catch (...) {
		for (i = 0; i < it->second.m_rasters.size(); i++)
			it->second.m_rasters[i]->endRemapping();
	}

	if (size > 0) //e' andata male...non liberato un blocco di grandezza sufficiente
	{
		printLog(size);
		assert(!"Niente memoria! scritto log");
		if (!locked)
			assert(false); //se entro nella remap, di sicuro c'e' 'size'  memoria disponibile; basta deframmentarla
	}
	return 0;
}

//------------------------------------------------------------------------------

void TBigMemoryManager::printLog(TUINT32 size)
{
	TFilePath fp("C:\\memorymaplog.txt");
	Tofstream os(fp);

	os << "memoria totale: " << m_allocatedMemory / 1024 << " KB\n";
	os << "memoria richiesta: " << size / 1024 << " KB\n";
	os << "memoria libera: " << m_availableMemory / 1024 << " KB\n\n\n";

	std::map<UCHAR *, Chunkinfo>::iterator it = m_chunks.begin();
	UCHAR *buffer = m_theMemory;
	UINT chunkSize = 0;
	for (; it != m_chunks.end(); it++) {
		TUINT32 gap = (TUINT32)((it->first) - (buffer + chunkSize));
		if (gap > 0)
			os << "- gap di " << gap / 1024 << " KB\n";
		if (it->second.m_size > 0)
			os << "- raster di " << it->second.m_size / 1024 << " KB" << ((it->second.m_rasters[0]->m_lockCount > 0) ? " LOCCATO!\n" : "\n");
		buffer = it->first;
		chunkSize = it->second.m_size;
	}
}