/* Reid Hochstedler
Pros: Keeps track of all alocations. Know's the difference between new and new[].
Cons: Can not keep track of what type is used, because new only knows the size of the object it is allocating.
*/

#ifndef _MEMCHECK_H_
#define _MEMCHECK_H_

#include <iostream>
#include <fstream>
#include <map>
#include <new>

#ifdef _DEBUG

struct memory_info
{
        void*         address;
        unsigned int  size;
        int           line;
  char          file[128];
  char          reason[48];
 
    inline memory_info::memory_info() :
    address(NULL), size(0), line(0)
    {
      memset( file,   0,  128 );
      memset( reason, 0,  48  );
    }
};

class Memory
{
public:
  static Memory* Instance();
  void Add(unsigned int size, const char* file, int line, void* address);
  void AddArray(unsigned int size, const char* file, int line, void* address);
  void Remove(void* address);
  void RemoveArray(void* address);
 
private:
  std::ofstream out;
  std::map<void *, memory_info> memory_map;
  std::map<void *, memory_info> array_map;

  void WriteFile();
  Memory();
  Memory(const Memory&);
  virtual ~Memory();
  Memory& operator= ( const Memory& );
};


inline Memory::Memory()
{
}

inline Memory::Memory( const Memory&)
{
}

inline Memory::~Memory()
{
  WriteFile();
}

inline Memory& Memory::operator = ( const Memory& )
{
  return *Instance();
}

inline Memory* Memory::Instance()
{
  static Memory instance;
  return &instance;
}

inline void Memory::Add(unsigned int size, const char* file, int line, void* address)
{
        memory_info info;

  info.address  = address;
  info.size     = size;
  info.line     = line;
  strncpy( info.file, file, 127     );
  strncpy( info.reason, "Pointer not deleted.", 48  );

  memory_map.insert( std::make_pair( address, info ) );
}

inline void Memory::AddArray(unsigned int size, const char* file, int line, void* address)
{
        memory_info info;

  info.address  = address;
  info.size     = size;
  info.line     = line;
  strncpy( info.file, file, 127     );
  strncpy( info.reason, "Pointer to array not deleted", 48  );

  array_map.insert( std::make_pair( address, info ) );
}

inline void Memory::Remove(void* address)
{
  std::map<void *, memory_info>::iterator i = memory_map.find( address );

  if( i != memory_map.end() ) {
    memory_map.erase( i );
  }
}

inline void Memory::RemoveArray(void* address)
{
  std::map<void *, memory_info>::iterator i = array_map.find( address );

  if( i != array_map.end() ) {
    array_map.erase( i );
  }
}

inline void Memory::WriteFile()
{
  long totalSize  = 0;
  char file_name[]="memory_leaks.csv";

  out.open( file_name, std::ios::out );

  if(!out.is_open()) {
    std::cerr <<"ERROR\ncould not open file " <<file_name << std::endl;
  } else {
    out << "Address, Size, File, Line, Reason" << std::endl;
    std::map<void *, memory_info>::iterator i;
   
    for( i = memory_map.begin(); i != memory_map.end(); ++i) {
      out << i->first       << ","
          << i->second.size << ","
          << i->second.file << ","
          << i->second.line << ","
          << i->second.reason
          << std::endl;

      totalSize += i->second.size;
    }
   
    for( i = array_map.begin(); i != array_map.end(); ++i) {
      out << i->first       << ","
          << i->second.size << ","
          << i->second.file << ","
          << i->second.line << ","
          << i->second.reason
          << std::endl;

      totalSize += i->second.size;
    }
    out << std::endl << "Total Size: " << totalSize << " bytes leaked." << std::endl;
    out.close();
  }
}

void* operator new(unsigned int size, const char* file, int line)
{
  //  C++ standard states that allocation requests of size = 0, will still return a valid value
  if( size == 0 )
    size  = 1;

  for(;;) {
    //  C++ standard says that we loop continously because the new_handler may free up memory
    void *ptr = (void *)malloc(size);
   
    if(ptr) {
      Memory::Instance()->Add(size, file, line, ptr);
      return  ptr;
    }
   
    //  No way to know new_handler so set it to NULL and then set it back.
    new_handler nh = std::set_new_handler( NULL );
    std::set_new_handler(nh);
   
    if( nh ) {
      (*nh)();
    } else {
      throw std::bad_alloc();
    }
  }
}

void* operator new[](unsigned int size, const char* file = "Unknown ", int line = 0)
{
  //  C++ standard states that allocation requests of size = 0, will still return a valid value
  if( size == 0 )
    size  = 1;

  for(;;) {
    //  C++ standard says that we loop continously because the new_handler may free up memory
    void *ptr = (void *)malloc(size);
   
    if(ptr) {
      Memory::Instance()->AddArray(size, file, line, ptr);
      return  ptr;
    }
   
    //  No way to know new_handler so set it to NULL and then set it back.
    new_handler nh = std::set_new_handler( NULL );
    std::set_new_handler(nh);
   
    if( nh ) {
      (*nh)();
    } else {
      throw std::bad_alloc();
    }
  }
}

inline void operator delete(void* ptr)
{
  if( !ptr )
    return;

  Memory::Instance()->Remove(ptr);
  free(ptr); 
}

inline void operator delete[](void* ptr)
{
  if( !ptr )
    return;

  Memory::Instance()->RemoveArray(ptr);
  free(ptr); 
}

#ifdef new
  #undef new
#endif

#define DEBUG_NEW new(__FILE__, __LINE__)
#define new DEBUG_NEW

#endif  //  _DEBUG
#endif  //  _MEMCHECK_H_