// __CR__
// Copyright (c) 2021 LongdaFeng All Rights Reserved
//
// This software contains the intellectual property of LongdaFeng
// or is licensed to LongdaFeng from third parties.  Use of this
// software and the intellectual property contained therein is
// expressly limited to the terms and conditions of the License Agreement
// under which it is provided by or on behalf of LongdaFeng.
// __CR__
//
// Created by Longda on 2010
//

#ifdef MEM_DEBUG

#include <cstddef>
#include <stdlib.h>
#include <string.h>

#include "lang/mutex.h"
#include "mm/lmem.h"

#define MEM_ID_HASH(p) (((unsigned long)(p) >> 8) % MEM_HASHTABLE_SIZE)

pthread_mutex_t CLMemTrace::mMutex = PTHREAD_MUTEX_INITIALIZER;
u64_t CLMemTrace::mUsedSize = 0;
MemID *CLMemTrace::mMemIDs[MEM_HASHTABLE_SIZE] = {0};
bool CLMemTrace::mVerbose = false;
;

void *CLMemTrace::malloc(size_t size, const char *file, const int line,
                         bool retry) throw(std::bad_alloc) {

  size_t allocSize = size + sizeof(MemID);
  void *usedPointer = NULL;

  do {
    MemID *ptr = (MemID *)::malloc(allocSize);
    if (ptr) {
      // successfully alloc memory

      // set the MemID
      strncpy(ptr->mFile, file, MemID::MEM_FILENAME_LEN - 1);
      ptr->mFile[MemID::MEM_FILENAME_LEN - 1] = '\0';
      ptr->mLine = line;
      ptr->mSize = size;

      usedPointer = (char *)ptr + sizeof(MemID);
      u64_t hashIndex = (u64_t)MEM_ID_HASH(usedPointer);

      MUTEX_LOCK(&mMutex);

      ptr->mNext = mMemIDs[hashIndex];
      mMemIDs[hashIndex] = ptr;

      mUsedSize += size;

      MUTEX_UNLOCK(&mMutex);

      if (mVerbose) {
        LOG_INFO("%s:%d alloc %llu memory %p", file, line, size, usedPointer);
      }
      return usedPointer;
    }

    if (retry == false) {
      throw std::bad_alloc();
      break;
    }

    std::new_handler allocHandler = getNewHandler();

    if (allocHandler) {
      (*allocHandler)();
    } else {
      throw std::bad_alloc();
      break;
    }
  } while (retry);

  return NULL;
}

void *CLMemTrace::realloc(void *pointer, size_t size, const char *file,
                          const int line) {
  if (pointer == NULL && size == 0) {
    return NULL;
  } else if (pointer == NULL && size != 0) {
    try {
      return malloc(size, file, line);
    } catch (std::bad_alloc &e) {
      LOG_WARN("NO memory to alloc for %llu", size);
      return NULL;
    }
  } else if (pointer && size == 0) {
    free(pointer);
    return NULL;
  }

  // the left case, ptr && size
  MemID *pMemID = NULL;
  MemID *pLast = NULL;

  MemID *pFreeMemID = NULL;
  MemID *pNewMemID = NULL;

  MemID oldMemID;

  bool foundOld = false;

  // use big lock
  MUTEX_LOCK(&mMutex);
  u64_t hashIndex = MEM_ID_HASH(pointer);
  pMemID = mMemIDs[hashIndex];
  while (pMemID) {
    if ((char *)pMemID + sizeof(MemID) != pointer) {
      // not found
      pLast = pMemID;
      pMemID = pMemID->mNext;
      continue;
    }

    // find
    foundOld = true;

    // backup old MemID firstly
    memcpy(&oldMemID, pMemID, sizeof(MemID));

    u64_t allocSize = size + sizeof(MemID);

    pNewMemID = (MemID *)::realloc(pMemID, allocSize);
    if (pNewMemID == NULL) {
      // case 1:no memory to alloc, free the old one
      if (pLast == NULL) {
        mMemIDs[hashIndex] = pMemID->mNext;
      } else {
        pLast->mNext = pMemID->mNext;
      }

      pFreeMemID = pMemID;
      mUsedSize -= oldMemID.mSize;
      break;
    }

    // set the new pNewMemID
    strncpy(pNewMemID->mFile, file, MemID::MEM_FILENAME_LEN - 1);
    pNewMemID->mFile[MemID::MEM_FILENAME_LEN - 1] = '\0';
    pNewMemID->mLine = line;
    pNewMemID->mSize = size;

    mUsedSize -= oldMemID.mSize;
    mUsedSize += size;

    if (pNewMemID == pMemID) {
      // case 2: just extension the old memory
      pFreeMemID = NULL;
      break;

    } else {
      // case 3: the old memory can't meet the requirement, alloc new
      /**
       * Firstly, remove the old one from table
       * don't add new before remove the old one
       */
      if (pLast == NULL) {
        mMemIDs[hashIndex] = pMemID->mNext;
      } else {
        pLast->mNext = pMemID->mNext;
      }
      pFreeMemID = pMemID;

      /**
       * Secondly, add the new one to table
       */
      u64_t newHashIndex =
          (u64_t)MEM_ID_HASH((char *)pNewMemID + sizeof(MemID));

      pNewMemID->mNext = mMemIDs[newHashIndex];
      mMemIDs[newHashIndex] = pNewMemID;

      /**
       * Third, do memory copy
       * to simplify the old logic, copy memory here
       */
      memcpy((char *)pNewMemID + sizeof(MemID),
             (char *)pFreeMemID + sizeof(MemID), pFreeMemID->mSize);
      break;
    }
  }

  MUTEX_UNLOCK(&mMutex);

  if (foundOld == false) {
    LOG_WARN(
        "Something is wrong, the old pointer %p isn't found, so alloc new one",
        pointer);
    try {
      return malloc(size, file, line, false);
    } catch (std::bad_alloc &e) {
      LOG_WARN("NO memory to alloc for %llu", size);
      return NULL;
    };
  }

  if (mVerbose) {
    LOG_INFO("Delete %p, file:%s, line:%u, size:%llu", pointer, oldMemID.mFile,
             oldMemID.mLine, oldMemID.mSize);
  }

  if (pFreeMemID) {
    ::free(pFreeMemID);
  }

  if (pNewMemID) {

    if (mVerbose) {
      LOG_INFO("Alloc %p, file:%s, line:%u, size:%llu",
               (char *)pNewMemID + sizeof(MemID), pNewMemID->mFile,
               pNewMemID->mLine, pNewMemID->mSize);
    }
    return pNewMemID;
  }

  return NULL;
}

void CLMemTrace::free(void *pointer) {
  if (pointer == NULL) {
    LOG_WARN("Free one empty pointer");
    return;
  }

  u64_t hashIndex = MEM_ID_HASH(pointer);

  MemID *pMemID = NULL;
  MemID *pLast = NULL;

  // use big lock
  MUTEX_LOCK(&mMutex);

  pMemID = mMemIDs[hashIndex];
  while (pMemID) {
    if ((char *)pMemID + sizeof(MemID) == pointer) {
      // find
      if (pLast == NULL) {
        mMemIDs[hashIndex] = pMemID->mNext;
      } else {
        pLast->mNext = pMemID->mNext;
      }

      mUsedSize -= pMemID->mSize;
      break;
    } else {
      pLast = pMemID;
      pMemID = pMemID->mNext;
    }
  }
  MUTEX_UNLOCK(&mMutex);

  if (pMemID) {

    if (mVerbose) {
      LOG_INFO("Delete %p, file:%s, line:%u, size:%llu", pointer, pMemID->mFile,
               pMemID->mLine, pMemID->mSize);
    }
    ::free(pMemID);
    return;
  } else {
    // not found
    LOG_ERROR("Double free for pointer :%p", pointer);
  }

  return;
}

std::new_handler CLMemTrace::getNewHandler() {
  std::new_handler newHandler = NULL;
  MUTEX_LOCK(&mMutex);

  newHandler = std::set_new_handler(0);
  std::set_new_handler(newHandler);

  MUTEX_UNLOCK(&mMutex);

  return newHandler;
}

void CLMemTrace::output() {
  for (int i = 0; i < MEM_HASHTABLE_SIZE; ++i) {
    // Don't lock outside of the loop
    // 1. avoid output too long to alloc/free memory
    // 2. if LOG_INFO alloc memory, it will leading to dead loop
    MUTEX_LOCK(&mMutex);
    MemID *ptr = mMemIDs[i];
    if (ptr == NULL) {
      MUTEX_UNLOCK(&mMutex);
      continue;
    }
    while (ptr) {
      // if LOG_INFO alloc memory, it will easy leading to dead lock
      LOG_INFO("Exist %p, file:%s, line:%u, size:%llu",
               (char *)ptr + sizeof(MemID), ptr->mFile, ptr->mLine, ptr->mSize);
      ptr = ptr->mNext;
    }

    MUTEX_UNLOCK(&mMutex);
  }
}

void *operator new(std::size_t size, const char *file, int line) {
  return CLMemTrace::malloc(size, file, line, true);
}

void *operator new[](std::size_t size, const char *file, int line) {
  return operator new(size, file, line);
}

void *operator new(std::size_t size) throw(std::bad_alloc) {
  return operator new(size, "<Unknown>", 0);
}

void *operator new[](std::size_t size) throw(std::bad_alloc) {
  return operator new(size);
}

void *operator new(std::size_t size, const std::nothrow_t &) throw() {
  void *pointer = NULL;
  try {
    pointer = operator new(size);
  } catch (std::bad_alloc &e) {
    LOG_WARN("Failed to alloc memory");
    return NULL;
  }
  return pointer;
}

void *operator new[](std::size_t size, const std::nothrow_t &) throw() {
  void *pointer = NULL;
  try {
    pointer = operator[] new(size);
  } catch (std::bad_alloc &e) {
    LOG_WARN("Failed to alloc memory");
    return NULL;
  }
  return pointer;
}

void operator delete(void *pointer) { CLMemTrace::free(pointer); }

void operator delete[](void *pointer) { operator delete(pointer); }

// Some older compilers like Borland C++ Compiler 5.5.1 and Digital Mars
// Compiler 8.29 do not support placement delete operators.
// NO_PLACEMENT_DELETE needs to be defined when using such compilers.
// Also note that in that case memory leakage will occur if an exception
// is thrown in the initialization (constructor) of a dynamically
// created object.
void operator delete(void *pointer, const char *file, int line) {
  operator delete(pointer);
}

void operator delete[](void *pointer, const char *file, int line) {
  operator delete(pointer, file, line);
}

void operator delete(void *pointer, const std::nothrow_t &) {
  operator delete(pointer, "<Unknown>", 0);
}

void operator delete[](void *pointer, const std::nothrow_t &) {
  operator delete(pointer, std::nothrow);
}

void *Lcalloc(size_t nmemb, size_t size, const char *file, const int line) {
  try {
    void *point = CLMemTrace::malloc(size * nmemb, file, line, false);
    if (point) {
      memset(point, 0, size * nmemb);
    }
  } catch (std::bad_alloc &e) {
    LOG_WARN("Failed to alloc memory");
    return NULL;
  }
  return pointer;
}

void *Lmalloc(size_t size, const char *file, const int line) {
  try {
    void *point = CLMemTrace::malloc(size, file, line, false);
  } catch (std::bad_alloc &e) {
    LOG_WARN("Failed to alloc memory");
    return NULL;
  }
  return pointer;
}

void Lfree(void *ptr) { CLMemTrace::free(pointer); }
void *Lrealloc(void *ptr, size_t size, const char *file, const int line) {
  // simplify the logic
  return CLMemTrace::realloc(ptr, size, file, line);
}

#endif /* MEM_DEBUG */