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

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

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

  }

  if (address) memset(address, 0x00, size);

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

    std::map<uchar *,="" chunkinfo="">::iterator it =</uchar>

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

  std::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() {

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

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

  {

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

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

  }

}

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

std::map<uchar *,="" chunkinfo="">::iterator TBigMemoryManager::shiftBlock(</uchar>

    const std::map<uchar *,="" chunkinfo="">::iterator &it, TUINT32 offset) {</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]);

  std::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;

  std::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

      {

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

  std::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;

  }

}