/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkMemoryUsageObserver.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even 
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR 
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

#include "itkConfigure.h"
#include "itkMemoryUsageObserver.h"

#if defined(WIN32) || defined(_WIN32)
  #include <windows.h>
  #if defined(SUPPORT_PSAPI)
    #include <psapi.h>
  #endif
#endif // defined(WIN32) || defined(_WIN32)

#ifdef linux
  #include "itkSmapsFileParser.h"
#endif // linux

#if defined(__APPLE__) && MAC_OS_X_VERSION >= MAC_OS_X_VERSION_10_2
  #include "itkSmapsFileParser.h"
#endif // Mac OS X

#if defined(__SUNPRO_CC) || defined (__sun__)
  #include <unistd.h>
  #include <stdio.h>
  #include <string>
  #include <sstream>
#endif // !defined(__SUNPRO_CC) && !defined (__sun__)

#if !defined(WIN32) && !defined(_WIN32)
  #include <sys/resource.h>     // getrusage()
  #if !defined( __APPLE__ ) && !defined( __SUNPRO_CC ) && !defined ( __sun__ ) && !defined( __FreeBSD__ ) \
  && !defined( __OpenBSD__ )
    #include <malloc.h>           // mallinfo()
  #endif // !defined(__APPLE__) && !defined(__SUNPRO_CC) && !defined (__sun__)
#endif // !defined(WIN32) && !defined(_WIN32)

#if defined( __OpenBSD__ )
#include <stdlib.h>
#endif

namespace itk
{

MemoryUsageObserverBase::~MemoryUsageObserverBase()
{
}


#if defined(WIN32) || defined(_WIN32)

/**         ----         Windows Memory Usage Observer       ----       */ 

WindowsMemoryUsageObserver::WindowsMemoryUsageObserver()
{
#if defined(SUPPORT_TOOLHELP32)
  m_hNTLib = ::LoadLibraryA("ntdll.dll");
  if(m_hNTLib)
    {
    // load the support function from the kernel
    ZwQuerySystemInformation = (PZwQuerySystemInformation)::GetProcAddress(m_hNTLib, 
                                                                           "ZwQuerySystemInformation");
    }
#endif
}

WindowsMemoryUsageObserver::~WindowsMemoryUsageObserver()
{
#if defined (SUPPORT_TOOLHELP32)
  if(m_hNTLib)
    {
    FreeLibrary(m_hNTLib);
    }
#endif
}

#if defined(SUPPORT_TOOLHELP32)

#define STATUS_INFO_LENGTH_MISMATCH ((NTSTATUS)0xC0000004L)

typedef LONG    KPRIORITY;
#define SystemProcessesAndThreadsInformation    5

typedef struct _CLIENT_ID {
  DWORD        UniqueProcess;
  DWORD        UniqueThread;
} CLIENT_ID;

typedef struct _UNICODE_STRING {
  USHORT        Length;
  USHORT        MaximumLength;
  PWSTR        Buffer;
} UNICODE_STRING;

typedef struct _VM_COUNTERS {
#ifdef _WIN64
  // the following was inferred by painful reverse engineering
  SIZE_T         PeakVirtualSize;     // not actually
  SIZE_T         PageFaultCount;
  SIZE_T         PeakWorkingSetSize;
  SIZE_T         WorkingSetSize;
  SIZE_T         QuotaPeakPagedPoolUsage;
  SIZE_T         QuotaPagedPoolUsage;
  SIZE_T         QuotaPeakNonPagedPoolUsage;
  SIZE_T         QuotaNonPagedPoolUsage;
  SIZE_T         PagefileUsage;
  SIZE_T         PeakPagefileUsage;
  SIZE_T         VirtualSize;         // not actually
#else
  SIZE_T         PeakVirtualSize;
  SIZE_T         VirtualSize;
  ULONG          PageFaultCount;
  SIZE_T         PeakWorkingSetSize;
  SIZE_T         WorkingSetSize;
  SIZE_T         QuotaPeakPagedPoolUsage;
  SIZE_T         QuotaPagedPoolUsage;
  SIZE_T         QuotaPeakNonPagedPoolUsage;
  SIZE_T         QuotaNonPagedPoolUsage;
  SIZE_T         PagefileUsage;
  SIZE_T         PeakPagefileUsage;
#endif
} VM_COUNTERS;

typedef struct _SYSTEM_THREADS {
  LARGE_INTEGER   KernelTime;
  LARGE_INTEGER   UserTime;
  LARGE_INTEGER   CreateTime;
  ULONG           WaitTime;
  PVOID           StartAddress;
  CLIENT_ID       ClientId;
  KPRIORITY       Priority;
  KPRIORITY       BasePriority;
  ULONG           ContextSwitchCount;
  LONG            State;
  LONG            WaitReason;
} SYSTEM_THREADS, * PSYSTEM_THREADS;

typedef struct _SYSTEM_PROCESSES { // Information Class 5
  ULONG NextEntryDelta;
  ULONG ThreadCount;
  ULONG Reserved1[6];
  LARGE_INTEGER CreateTime;
  LARGE_INTEGER UserTime;
  LARGE_INTEGER KernelTime;
  UNICODE_STRING ProcessName;
  KPRIORITY BasePriority;
#ifdef _WIN64
  ULONG  pad1;
  ULONG  ProcessId;
  ULONG  pad2;
  ULONG  InheritedFromProcessId;
  ULONG  pad3;
  ULONG  pad4;
  ULONG  pad5;
#else
  ULONG  ProcessId;
  ULONG  InheritedFromProcessId;
#endif
  ULONG HandleCount;
  ULONG Reserved2[2];
  VM_COUNTERS VmCounters;
#if defined(_WIN64) || _WIN32_WINNT >= 0x500
  IO_COUNTERS        IoCounters;
#endif
  SYSTEM_THREADS Threads[1];
} SYSTEM_PROCESSES, *PSYSTEM_PROCESSES;
#endif
  
MemoryUsageObserverBase::MemoryLoadType 
WindowsMemoryUsageObserver::GetMemoryUsage()
{

  MemoryLoadType mem = 0;

#if defined(SUPPORT_PSAPI)
  DWORD pid = GetCurrentProcessId();
  PROCESS_MEMORY_COUNTERS memoryCounters;

  HANDLE  hProcess = OpenProcess(  PROCESS_QUERY_INFORMATION |
                                   PROCESS_VM_READ,
                                   FALSE, pid );

  if (NULL == hProcess)
    {
    // Can't determine memory usage.
    return 0;
    }

  GetProcessMemoryInfo( hProcess, &memoryCounters, sizeof(memoryCounters));

  mem = static_cast<MemoryLoadType>( 
    static_cast<double>( memoryCounters.PagefileUsage )
    / 1024.0 ); 
#elif defined(SUPPORT_TOOLHELP32)

  /* Retrieve memory usage using Windows Native API. For more information, 
   * read the book "Windows NT 2000 Native API Reference" 
  */

  if(!m_hNTLib)
    {
    itkGenericExceptionMacro( << "Can't find ntdll.dll. "
                              << "You should probably disable SUPPORT_TOOLHELP32" );
    }
  // the ntdll.dll library could not have been opened (file not found?)
  if ( !ZwQuerySystemInformation )
    {
    itkGenericExceptionMacro( << "The file ntdll.dll is not supported. "
                              << "You should probably disable SUPPORT_TOOLHELP32" );
    return mem;
    }

  DWORD pid = GetCurrentProcessId();
  ULONG n = 50;
  PSYSTEM_PROCESSES sp = new SYSTEM_PROCESSES[n];
  // as we can't know how many processes running, we loop and test a new size everytime.
  while (ZwQuerySystemInformation(SystemProcessesAndThreadsInformation,
                                  sp, n * sizeof *sp, 0)
         == STATUS_INFO_LENGTH_MISMATCH)
    {
    delete [] sp;
    n = n * 2;
    sp = new SYSTEM_PROCESSES[n];
    }
  bool done = false;
  for ( PSYSTEM_PROCESSES spp = sp; 
        !done; 
        spp = PSYSTEM_PROCESSES(PCHAR(spp) + spp->NextEntryDelta))
    {
    // only the current process is interesting here
    if (spp->ProcessId == pid)
      {
      mem = static_cast<MemoryLoadType>( 
        static_cast<double>( spp->VmCounters.PagefileUsage - sizeof(*sp)) / 1024);
      break;
      }
    done = (spp->NextEntryDelta == 0);
    }
  delete [] sp;
  
#else

  /* This solution is not optimal as it returns the system memory usage 
   * instead of the process memory usage. 
  */

  MEMORYSTATUSEX statex;

  statex.dwLength = sizeof (statex);

  GlobalMemoryStatusEx (&statex);

  mem   = static_cast<MemoryLoadType>( 
    static_cast<double>( statex.ullTotalPhys - statex.ullAvailPhys) / 1024);
#endif
  return mem;
}

#endif // WIN32

#if linux

/**         ----         Linux Memory Usage Observer       ----       */ 

LinuxMemoryUsageObserver::~LinuxMemoryUsageObserver()
{
}

MemoryUsageObserverBase::MemoryLoadType 
LinuxMemoryUsageObserver::GetMemoryUsage()
{
  SmapsFileParser<SmapsData_2_6> m_ParseSmaps;
  m_ParseSmaps.ReadFile();
  return m_ParseSmaps.GetHeapUsage() + m_ParseSmaps.GetStackUsage();
}

#endif // linux

#if defined(__APPLE__) && MAC_OS_X_VERSION >= MAC_OS_X_VERSION_10_2

/**         ----         Mac OS X Memory Usage Observer       ----       */ 

MacOSXMemoryUsageObserver::~MacOSXMemoryUsageObserver()
{
}

MemoryUsageObserverBase::MemoryLoadType 
MacOSXMemoryUsageObserver::GetMemoryUsage()
{
  VMMapFileParser<VMMapData_10_2> m_ParseVMMmap;
  m_ParseVMMmap.ReadFile();
  return m_ParseVMMmap.GetHeapUsage() + m_ParseVMMmap.GetStackUsage();
}

#endif // Mac OS X

#if defined(__SUNPRO_CC) || defined (__sun__)

/**         ----         Sun Solaris Memory Usage Observer       ----       */ 

SunSolarisMemoryUsageObserver::~SunSolarisMemoryUsageObserver()
{
}
/** On Sun Solaris machines, the system call pmap returns information on process.
 *  calling "pmap PID", the output shall be like the following:
 *  102905:    *my_app*
 *  00010000    192K r-x--  /usr/bin/my_app
 *  00042000     40K rwx--    [ heap ]
 *  FF180000    664K r-x--  /usr/lib/libc.so.1
 *  FF236000     24K rwx--  /usr/lib/libc.so.1
 *  FF23C000      8K rwx--  /usr/lib/libc.so.1
 *  FF250000      8K rwx--    [ anon ]
 *  ...       ...    ...    ...
 *  FF3F6000      8K rwx--  /usr/lib/ld.so.1
 *  FFBFC000     16K rw---    [ stack ]
 *   total     1880K
 */
MemoryUsageObserverBase::MemoryLoadType 
SunSolarisMemoryUsageObserver::GetMemoryUsage()
{
  MemoryLoadType mem = 0;
  int pid = getpid();

  FILE * fp = NULL;
  std::stringstream command;
  command << "pmap " << pid << std::endl;
  
  if ((fp = popen(command.str().c_str(), "r")) == NULL)
    {
    itkGenericExceptionMacro( << "Error using pmap. Can execute pmap command" );
    }
  char remaining[256];
  int pmappid = -1;
  fscanf(fp,"%d:%s",&pmappid,remaining);
  //the first word shall be the process ID
  if ( pmappid != pid)
    {
    itkGenericExceptionMacro( << "Error using pmap. 1st line output shall be PID: name" );
    }
  bool heapNotFound = true;
  char address[64],perms[32];
  int memUsage = 0;
  std::string mapping;
  while (heapNotFound)
    {
    if ( fscanf(fp,"%s %dK %s",address, &memUsage, perms) != 3 )
      break;
    if ( fgets(remaining,256,fp) != NULL )
      {
      mapping = remaining;
      if ( mapping.find("[ heap ]",0) != std::string::npos)
        {
        mem = memUsage;
        heapNotFound = false;
        break;
        }
      // if no [ heap ] token is defined, accumulate all the [ xxx ] tokens
      else if ( mapping.find("[ ",0) != std::string::npos && 
                mapping.find(" ]",0) != std::string::npos )
        {  
        mem += memUsage;
        }
      }
    else
      {
      if (ferror (fp))
        {
        itkGenericExceptionMacro( << "Error using pmap. Corrupted pmap output" );
        }
      }
    }
  if (pclose(fp) == -1)
    {
    itkGenericExceptionMacro( << "Error using pmap. Can't close pmap output file." );
    }
  return mem;
}
#endif //defined(__SUNPRO_CC) || defined (__sun__)

#if !defined( WIN32 ) && !defined( _WIN32 ) || defined( __OpenBSD__ )

/**         ----         SysResource Memory Usage Observer       ----       */ 

SysResourceMemoryUsageObserver::~SysResourceMemoryUsageObserver()
{
}

MemoryUsageObserverBase::MemoryLoadType 
SysResourceMemoryUsageObserver::GetMemoryUsage()
{
  // Maybe use getrusage() ??
  rusage resourceInfo;

  int who = RUSAGE_SELF;
  if (getrusage(who, &resourceInfo) == 0)
    {
    return resourceInfo.ru_ixrss;
    }

  return 0;
}

#if !defined( __APPLE__ ) && !defined( __SUNPRO_CC ) && !defined ( __sun__ ) && !defined( __FreeBSD__ ) \
  && !defined( __OpenBSD__ )

/**         ----         Mallinfo Memory Usage Observer       ----       */ 

MallinfoMemoryUsageObserver::~MallinfoMemoryUsageObserver()
{
}

MemoryUsageObserverBase::MemoryLoadType 
MallinfoMemoryUsageObserver::GetMemoryUsage()
{
  struct mallinfo minfo = mallinfo(); 

  MemoryLoadType mem = static_cast<MemoryLoadType>( 
    static_cast<double>( minfo.uordblks ) / 1024.0 );
  return mem;
}

#endif //  !defined(__APPLE__) && !defined(__SUNPRO_CC) && !defined (__sun__) &&
       // !defined(__FreeBSD__) && !defined(__OpenBSD__)

#endif // Unix and Mac Platforms !defined(WIN32) && !defined(_WIN32) 

}//end namespace itk