#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <windows.h>
#include <psapi.h>
#pragma comment(linker, "/DEFAULTLIB:psapi.lib")
#include <time.h>
#include "jni.h"
#include "jvmti.h"

#include "agent_util.h"

#define MAX(X,Y) ((X) > (Y) ? (X) : (Y))
#define DELAY_BETWEEN_CHECKS 1000

static FILE *logFile;
static jrawMonitorID lock;

static int 
getNativeMemoryUsage()
{
    PROCESS_MEMORY_COUNTERS pmc;

    if ( GetProcessMemoryInfo( GetCurrentProcess(), &pmc, sizeof(pmc)) )
    {
//        stdout_message( , pmc.PageFaultCount );
//        stdout_message( "\tPeakWorkingSetSize: 0x%08X\n", pmc.PeakWorkingSetSize );
//        stdout_message( "\tWorkingSetSize: 0x%08X\n", pmc.WorkingSetSize );
//        stdout_message( "\tQuotaPeakPagedPoolUsage: 0x%08X\n", pmc.QuotaPeakPagedPoolUsage );
//        stdout_message( "\tQuotaPagedPoolUsage: 0x%08X\n", pmc.QuotaPagedPoolUsage );
//        stdout_message( "\tQuotaPeakNonPagedPoolUsage: 0x%08X\n", pmc.QuotaPeakNonPagedPoolUsage );
//        stdout_message( "\tQuotaNonPagedPoolUsage: 0x%08X\n", pmc.QuotaNonPagedPoolUsage );
//        stdout_message( "\tPagefileUsage: %8d\n", pmc.PagefileUsage ); 
//        stdout_message( "\tPeakPagefileUsage: %8d\n", pmc.PeakPagefileUsage );"\tPageFaultCount: 0x%08X\n"
		return (int)pmc.WorkingSetSize;
    } else {
		stdout_message("Can't get process memory usage\n");
		return -1;
	}
}

static int
getHeapMemoryUsage(JNIEnv *env) 
{
    jclass runtimeClass;
    jmethodID getRuntimeMethodId;
	jmethodID totalMemoryMethodId;
	jmethodID freeMemoryMethodId;
    jobject runtimeObject;
	jlong total, free;

    runtimeClass = (*env)->FindClass(env, "java/lang/Runtime");
    if ( runtimeClass == NULL ) {
		fatal_error("Cannot find Runtime class\n");
    }
	
    getRuntimeMethodId = (*env)->GetStaticMethodID(env, runtimeClass, "getRuntime", "()Ljava/lang/Runtime;");
    if ( getRuntimeMethodId == NULL ) {
		fatal_error("Cannot find getRuntime() method\n");
    }	
	
	totalMemoryMethodId = (*env)->GetMethodID(env, runtimeClass, "totalMemory", "()J");
    if ( totalMemoryMethodId == NULL ) {
		fatal_error("Cannot find totalMemory() method\n");
    }
	
	freeMemoryMethodId = (*env)->GetMethodID(env, runtimeClass, "freeMemory", "()J");
    if ( freeMemoryMethodId == NULL ) {
		fatal_error("Cannot find freeMemory() method\n");
    }
	
    runtimeObject = (*env)->CallStaticObjectMethod(env, runtimeClass, getRuntimeMethodId);
    if ( runtimeObject == NULL ) {
		fatal_error("Cannot get instance of java.lang.Runtime\n");
    }
	
	total = (*env)->CallLongMethod(env, runtimeObject, totalMemoryMethodId);
	free = (*env)->CallLongMethod(env, runtimeObject, freeMemoryMethodId);
	
    return (int)(total - free);
}

/* Worker thread that waits for garbage collections */
static void JNICALL
worker(jvmtiEnv* jvmti, JNIEnv* jni, void *p) {	
	clock_t start, checkStart;
	int nativeUsage, heapUsage, i;
	
	stdout_message("Worker thread started.\n");
	start = clock();
	
	if (logFile != NULL) {
		fprintf(logFile, "iteration; clocks_since_start; page_file_usage; heap_usage\n");
	}
	
	for(i = 1; ;i++) {		
		if (logFile != NULL) {
			checkStart = clock();
			
			nativeUsage = getNativeMemoryUsage();
			heapUsage = getHeapMemoryUsage(jni);
			
	
			fprintf(logFile, "%d; %d; %d; %d \n", 
					(i-1), 
					(checkStart - start),
					nativeUsage,
					heapUsage);
		} else {
			stdout_message("Can't write to file.\n");
			return;
		}
		
		Sleep(MAX((i*DELAY_BETWEEN_CHECKS) - (clock() - start), 1)); // FIXME: assume, that CLOCKS_PER_SEC == 1000
	}
}

/* Get heap memory usage */
static jthread 
alloc_thread(JNIEnv *env) 
{
    jclass    thrClass;
    jmethodID cid;
    jthread   res;

    thrClass = (*env)->FindClass(env, "java/lang/Thread");
    if ( thrClass == NULL ) {
	fatal_error("Cannot find Thread class\n");
    }
    cid      = (*env)->GetMethodID(env, thrClass, "<init>", "()V");
    if ( cid == NULL ) {
	fatal_error("Cannot find Thread constructor method\n");
    }
    res      = (*env)->NewObject(env, thrClass, cid);
    if ( res == NULL ) {
	fatal_error("Cannot create new Thread object\n");
    }
    return res;
}

/* Callback for JVMTI_EVENT_VM_INIT */
static void JNICALL 
vm_init(jvmtiEnv *jvmti, JNIEnv *env, jthread thread)
{
    jvmtiError err;
    jint       runtime_version;
	
    err = (*jvmti)->GetVersionNumber(jvmti, &runtime_version);
    check_jvmti_error(jvmti, err, "get version number");
	
	err = (*jvmti)->RunAgentThread(jvmti, alloc_thread(env), &worker, NULL,	JVMTI_THREAD_MAX_PRIORITY);
    check_jvmti_error(jvmti, err, "running agent thread");
}

/* Agent_OnLoad() is called first, we prepare for a VM_INIT event here. */
JNIEXPORT jint JNICALL
Agent_OnLoad(JavaVM *vm, char *options, void *reserved)
{
    jint                rc;
    jvmtiError          err;
    jvmtiEventCallbacks callbacks;
    jvmtiEnv           *jvmti;
    
    /* Get JVMTI environment */
    rc = (*vm)->GetEnv(vm, (void **)&jvmti, JVMTI_VERSION);
    if (rc != JNI_OK) {
	fatal_error("ERROR: Unable to create jvmtiEnv, GetEnv failed, error=%d\n", rc);
	return -1;
    }

    /* Set callbacks and enable event notifications */
    memset(&callbacks, 0, sizeof(callbacks));
    callbacks.VMInit                  = &vm_init;
    err = (*jvmti)->SetEventCallbacks(jvmti, &callbacks, sizeof(callbacks));
    check_jvmti_error(jvmti, err, "set event callbacks");
    err = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE, 
			JVMTI_EVENT_VM_INIT, NULL);
    check_jvmti_error(jvmti, err, "set event notify");
	
	logFile = fopen ("memoryusage.csv","w");
	
    return 0;
}

/* Agent_OnUnload() is called last */
JNIEXPORT void JNICALL
Agent_OnUnload(JavaVM *vm)
{
	if (logFile != NULL) {
		fclose(logFile);
		logFile = NULL;
	}
}

