#include <windows.h>

#include <assert.h>

#include <psapi.h>

#include <stdio.h>

#include <tchar.h>

#include <time.h>

#include "Shlwapi.h"

#pragma comment(lib, "psapi.lib")

#pragma comment(lib, "shlwapi.lib")

bool gSingleProcess = true;

int gQueryInterval = 5; // seconds

time_t gDuration = 0;   // seconds

LPTSTR gCommandLine;

HRESULT ProcessArgs(int argc, TCHAR *argv[]);

HRESULT PrintUsage();

void UseImage(void (functionForQueryType(HANDLE)));

void QueryContinuously(HANDLE hProcess);

time_t ElapsedTime(time_t startTime);

unsigned int OneQuery(HANDLE hProcess);

unsigned int OneQueryMP(HANDLE hProcess);

int __cdecl _tmain (int argc, TCHAR *argv[])

{

    HRESULT result = ProcessArgs(argc, argv);

    if (FAILED(result))

        return result;

    UseImage(QueryContinuously);

    return S_OK;

}

HRESULT ProcessArgs(int argc, TCHAR *argv[])

{

    LPTSTR argument;

    for( int count = 1; count < argc; count++ ) {

        argument = argv[count] ;

        if (wcsstr(argument, _T("-h")) ||

            wcsstr(argument, _T("--help")))

            return PrintUsage();

        else if (wcsstr(argument, _T("--exe"))) {

            gCommandLine = argv[++count];

            if (wcsstr(gCommandLine, _T("chrome.exe")))

                gSingleProcess = false;

        } else if (wcsstr(argument, _T("-i")) ||

            wcsstr(argument, _T("--interval"))) {

            gQueryInterval = _wtoi(argv[++count]);

            if (gQueryInterval < 1) {

                printf("ERROR: invalid interval\n");

                return E_INVALIDARG;

            }

        } else if (wcsstr(argument, _T("-d")) ||

            wcsstr(argument, _T("--duration"))) {

            gDuration = _wtoi(argv[++count]);

            if (gDuration < 1) {

                printf("ERROR: invalid duration\n");

                return E_INVALIDARG;

            }

        } else {

            _tprintf(_T("ERROR: unrecognized argument \"%s\"\n"), (LPCTSTR)argument);

            return PrintUsage();

        }

    }

    if (argc < 2 || !wcslen(gCommandLine) ) {

        printf("ERROR: executable path is required\n");

        return PrintUsage();

    }

    return S_OK;

}

HRESULT PrintUsage()

{

    printf("record-memory-win --exe EXE_PATH\n");

    printf("    Launch an executable and print the memory usage (in Private Bytes)\n");

    printf("    of the process.\n\n");

    printf("Usage:\n");

    printf("-h [--help]         : Print usage\n");

    printf("--exe arg           : Launch specified image.  Required\n");

    printf("-i [--interval] arg : Print memory usage every arg seconds.  Default: 5 seconds\n");

    printf("-d [--duration] arg : Run for up to arg seconds.  Default: no limit\n\n");

    printf("Examples:\n");

    printf("    record-memory-win --exe \"C:\\Program Files\\Safari\\Safari.exe\"\n");

    printf("    record-memory-win --exe Safari.exe -i 10 -d 7200\n");

    return E_FAIL;

}

void UseImage(void (functionForQueryType(HANDLE)))

{

    STARTUPINFO si = {0};

    si.cb = sizeof(STARTUPINFO);

    PROCESS_INFORMATION pi = {0};

    // Start the child process. 

    if(!CreateProcess( NULL,   // No module name (use command line)

        gCommandLine,        // Command line

        NULL,           // Process handle not inheritable

        NULL,           // Thread handle not inheritable

        FALSE,          // Set handle inheritance to FALSE

        0,              // No creation flags

        NULL,           // Use parent's environment block

        NULL,           // Use parent's starting directory 

        &si,            // Pointer to STARTUPINFO structure

        &pi ))          // Pointer to PROCESS_INFORMATION structure

        printf("CreateProcess failed (%d)\n", GetLastError());

    else {

        printf("Created process\n");

        functionForQueryType(pi.hProcess);

        // Close process and thread handles. 

        CloseHandle( pi.hProcess );

        CloseHandle( pi.hThread );

    }

}

void QueryContinuously(HANDLE hProcess)

{

    Sleep(2000); // give the process some time to launch

    bool pastDuration = false;

    time_t startTime = time(NULL);

    unsigned int memUsage = gSingleProcess ? OneQuery(hProcess) : OneQueryMP(hProcess);

    while(memUsage && !pastDuration) {

        printf( "%u\n", memUsage );

        Sleep(gQueryInterval*1000);

        memUsage = gSingleProcess ? OneQuery(hProcess) : OneQueryMP(hProcess);

        pastDuration = gDuration > 0 ? ElapsedTime(startTime) > gDuration : false;

    } 

}

// returns elapsed time in seconds

time_t ElapsedTime(time_t startTime)

{

    time_t currentTime = time(NULL);

    return currentTime - startTime;

}

// returns Commit Size (Private Bytes) in bytes

unsigned int OneQuery(HANDLE hProcess)

{

    PROCESS_MEMORY_COUNTERS_EX pmc;

    if (NULL == hProcess)

        return 0;

    if (GetProcessMemoryInfo(hProcess, (PPROCESS_MEMORY_COUNTERS)&pmc, sizeof(pmc)))

        return (unsigned)pmc.PrivateUsage;

    return 0;

}

// returns Commit Size (Private Bytes) in bytes for multi-process executables

unsigned int OneQueryMP(HANDLE hProcess)

{

    unsigned int memUsage = 0;

    TCHAR monitoredProcessName[MAX_PATH];

    GetProcessImageFileName(hProcess, monitoredProcessName, sizeof(monitoredProcessName)/sizeof(TCHAR));

    LPTSTR shortProcessName = PathFindFileName(monitoredProcessName);

    DWORD aProcesses[1024], cbNeeded, cProcesses;

    HANDLE hFoundProcess;

    if (!EnumProcesses(aProcesses, sizeof(aProcesses), &cbNeeded))

        return 0;

    // Calculate how many process identifiers were returned.

    cProcesses = cbNeeded / sizeof(DWORD);

    // find existing process

    for (unsigned int i = 0; i < cProcesses; i++)

        if (aProcesses[i] != 0) {

            DWORD retVal = 0;

            TCHAR foundProcessName[MAX_PATH];

            // Get a handle to the process.

            hFoundProcess = OpenProcess(PROCESS_QUERY_INFORMATION |

                                   PROCESS_VM_READ,

                                   FALSE, aProcesses[i]);

            // Get the process name.

            if (NULL != hFoundProcess) {

                HMODULE hMod;

                DWORD cbNeeded;

                if (EnumProcessModules(hFoundProcess, &hMod, sizeof(hMod), &cbNeeded)) {

                    GetModuleBaseName(hFoundProcess, hMod, foundProcessName, sizeof(foundProcessName)/sizeof(TCHAR));

                    if (wcsstr(foundProcessName, shortProcessName))

                        memUsage += OneQuery(hFoundProcess);

                }

            }

            CloseHandle(hFoundProcess);

        }

    return memUsage;

}