#include "diagnostics_win.h"
#include "procinfo.h"

using std::vector;

// NtQuerySystemInformation
typedef NTSTATUS (WINAPI *tNTQSI)(
    ULONG SystemInformationClass,
    PVOID SystemInformation,
    ULONG SystemInformationLength,
    PULONG ReturnLength
);

static int get_process_information(PVOID* ppBuffer, PULONG pcbBuffer) {
    int      retval = 0;
    NTSTATUS Status = STATUS_INFO_LENGTH_MISMATCH;
    HANDLE   hHeap  = GetProcessHeap();
    HMODULE  hNTDllLib = NULL;
    tNTQSI   pNTQSI = NULL;

    hNTDllLib = GetModuleHandle("ntdll.dll");
    pNTQSI = (tNTQSI)GetProcAddress(hNTDllLib, "NtQuerySystemInformation");

    do {
        *ppBuffer = HeapAlloc(hHeap, HEAP_ZERO_MEMORY, *pcbBuffer);
        if (ppBuffer == NULL) {
            retval = ERROR_NOT_ENOUGH_MEMORY;
        }

        Status = pNTQSI(
            SystemProcessAndThreadInformation,
            *ppBuffer,
            *pcbBuffer,
            pcbBuffer
        );

        if (Status == STATUS_INFO_LENGTH_MISMATCH) {
            HeapFree(hHeap, NULL, *ppBuffer);
            *pcbBuffer *= 2;
        } else if (!NT_SUCCESS(Status)) {
            HeapFree(hHeap, NULL, *ppBuffer);
            retval = Status;
        }
    } while (Status == STATUS_INFO_LENGTH_MISMATCH);

    return retval;
}

// Note: the following will work on both NT and XP,
// because the NT process structure differs only at the end
//
int get_procinfo_XP(vector<PROCINFO>& pi) {
    ULONG                   cbBuffer = 32*1024;    // 32k initial buffer
    PVOID                   pBuffer = NULL;
    PSYSTEM_PROCESSES       pProcesses = NULL;

    get_process_information(&pBuffer, &cbBuffer);
    pProcesses = (PSYSTEM_PROCESSES)pBuffer;
    while (pProcesses) {
        PROCINFO p;
		p.id = pProcesses->ProcessId;
		p.parentid = pProcesses->InheritedFromProcessId;
        p.swap_size = pProcesses->VmCounters.PagefileUsage;
        p.working_set_size = pProcesses->VmCounters.WorkingSetSize;
		p.page_fault_count = pProcesses->VmCounters.PageFaultCount;
        p.user_time = ((double) pProcesses->UserTime.QuadPart)/1e7;
        p.kernel_time = ((double) pProcesses->KernelTime.QuadPart)/1e7;
		p.id = pProcesses->ProcessId;
		p.parentid = pProcesses->InheritedFromProcessId;
		p.is_boinc_app = false;
        pi.push_back(p);

        if (!pProcesses->NextEntryDelta) {
            break;
        }
        pProcesses = (PSYSTEM_PROCESSES)(((LPBYTE)pProcesses) + pProcesses->NextEntryDelta);
    }

    if (pBuffer) HeapFree(GetProcessHeap(), NULL, pBuffer);
    return 0;
}

// get a list of all running processes.
//
int procinfo_setup(vector<PROCINFO>& pi) {
    OSVERSIONINFO osvi; 
    osvi.dwOSVersionInfoSize = sizeof(osvi);
    GetVersionEx(&osvi);

    pi.clear();

    switch(osvi.dwPlatformId) {
    case VER_PLATFORM_WIN32_WINDOWS:
        // Win95, Win98, WinME
        return 0;   // not supported
    case VER_PLATFORM_WIN32_NT:
        return get_procinfo_XP(pi);
    }
    return 0;
}

// scan the process table from the given point,
// adding in CPU time and mem usage
// 
void add_proc_totals(PROCINFO& pi, vector<PROCINFO>& piv, int pid, int start) {
	unsigned int i;
	for (i=start; i<piv.size(); i++) {
		PROCINFO& p = piv[i];
		if (p.id == pid || p.parentid == pid) {
			pi.kernel_time += p.kernel_time;
			pi.user_time += p.user_time;
			pi.swap_size += p.swap_size;
			pi.working_set_size += p.working_set_size;
			pi.page_fault_count += p.page_fault_count;
			p.is_boinc_app = true;
		} 
		if (p.parentid == pid) {
			add_proc_totals(pi, piv, p.id, i+1);	// recursion - woo hoo!
		}
	}
}

// fill in the given PROCINFO (which initially is zero except for id)
// with totals from that process and all its descendants
//
void procinfo_app(PROCINFO& pi, vector<PROCINFO>& piv) {
	add_proc_totals(pi, piv, pi.id, 0);
}

void procinfo_other(PROCINFO& pi, vector<PROCINFO>& piv) {
	unsigned int i;
	memset(&pi, 0, sizeof(pi));
	for (i=0; i<piv.size(); i++) {
		PROCINFO& p = piv[i];
		if (!p.is_boinc_app) {
			pi.kernel_time += p.kernel_time;
			pi.user_time += p.user_time;
			pi.swap_size += p.swap_size;
			pi.working_set_size += p.working_set_size;
			p.is_boinc_app = true;
		}
	}
}