#include "traster.h"

#include "tbigmemorymanager.h"

#include "timagecache.h"

#include "tsystem.h"

#include "tconvert.h"

#include <set></set>

#include "tfilepath_io.h"

#ifdef _DEBUG

std::set<traster *=""> Rasters;</traster>

#endif

#ifdef TNZCORE_LIGHT

#ifdef MessageBox

#undef MessageBox

#endif

#define MessageBox MessageBoxA

#endif

class Chunkinfo

{

public:

	TUINT32 m_size;

	//int m_locks;

	vector<traster *=""> m_rasters;</traster>

	//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);</uchar*,>

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);</uchar*,>

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);

	map<uchar *,="" chunkinfo="">::iterator it = m_chunks.begin();</uchar>

	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++;

	}

	return address;

}

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

#ifdef _DEBUG

void TBigMemoryManager::getRasterInfo(int &rasterCount, TUINT32 &totRasterMemInKb, int ¬CachedRasterCount, TUINT32 ¬CachedRasterMemInKb)

{

	totRasterMemInKb = 0;

	notCachedRasterMemInKb = 0;

	notCachedRasterCount = 0;

	rasterCount = Rasters.size();

	std::set<traster *="">::iterator it = Rasters.begin();</traster>

	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 = tmax(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)</size>

  {

  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) {

		map<uchar *,="" chunkinfo="">::iterator it = m_chunks.find(ras->m_parent->m_buffer);</uchar>

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

#ifdef _DEBUG

//assert(!it->second.m_rasters.empty());

//for (UINT i=0; i<it->second.m_rasters.size(); i++)</it->

//  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;

	map<uchar *,="" chunkinfo="">::iterator it = m_chunks.begin();</uchar>

	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()

{

	map<uchar *,="" chunkinfo="">::iterator it = m_chunks.begin();</uchar>

	TSystem::outputDebug("BIGMEMORY chunks totali: " + toString((int)m_chunks.size()) + "\n");

	int count = 0;

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

		TSystem::outputDebug("chunk #" +

							 toString((int)count++) +

							 "dimensione(kb):" +

							 toString((int)(it->second.m_size >> 10)) +

							 "num raster:" +

							 toString((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);

	map<uchar *,="" chunkinfo="">::iterator it = m_chunks.find(buffer);</uchar>

	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

	{

		vector<traster *="">::iterator it2 = it->second.m_rasters.begin();</traster>

		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();

	map<uchar *,="" chunkinfo="">::iterator it = m_chunks.begin();</uchar>

	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);

	}

}

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

map<uchar *,="" chunkinfo="">::iterator TBigMemoryManager::shiftBlock(const map<uchar *,="" chunkinfo="">::iterator &it, TUINT32 offset)</uchar></uchar>

{

	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]);

	map<uchar *,="" chunkinfo="">::iterator it1 = m_chunks.find(newAddress);</uchar>

	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;

	map<uchar *,="" chunkinfo="">::iterator it = m_chunks.begin();</uchar>

	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

			{

				vector<traster *=""> &rasters = it->second.m_rasters;</traster>

				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";

	map<uchar *,="" chunkinfo="">::iterator it = m_chunks.begin();</uchar>

	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;

	}

}