
#include <string>
#include <sstream>
#include <iostream>
#include <dlfcn.h>
#include "agent.h"
#include "cpu_profiler.h"
#include "lock_profiler.h"
#include <chrono>
#include <thread>

using namespace std;

Agent Agent::instance;

void Agent::LogMessage(string level, string msg) {
	if (debug_mode) {
		cout << "StackImpact Agent: " << level << ": " << msg << endl;
	}
}


void Agent::LogInfo(string msg) {
	LogMessage("INFO", msg);
}


void Agent::LogError(string msg) {
	LogMessage("ERROR", msg);
}


bool Agent::CheckJVMTIError(jvmtiEnv *jvmti, jvmtiError error, string msg) {
	if (error != JVMTI_ERROR_NONE) {
		if (debug_mode) {
			char *error_str;
			error_str = NULL;
			jvmti->GetErrorName(error, &error_str);
			ostringstream error_message;
			error_message << "JVMTI: "
					<< (error_str == NULL ? "Unknown" : error_str)
					<< " (" << error << "): "
					<< msg;

			LogError(error_message.str());

			return false;
		}
	}

	return true;
}


void Agent::Init(JavaVM* vm) {
	Agent* agent = &Agent::instance;

	jvm = vm;

	int rc = jvm->GetEnv((void**)&agent->jni, JNI_VERSION_1_6);
	if (rc != JNI_OK && rc != JNI_EDETACHED) {
		agent->LogError("GetEnv, JNI_VERSION_1_6");
		return;
	}

	rc = jvm->GetEnv((void **)&agent->jvmti, JVMTI_VERSION_1_2);
	if (rc != JNI_OK) {
		agent->LogError("GetEnv, JVMTI_VERSION_1_2");
		return;
	}
}


void LoadMethodIDs(jvmtiEnv* jvmti, jclass klass) {
    jint num_methods;
    jmethodID* methods;
    if (jvmti->GetClassMethods(klass, &num_methods, &methods) == 0) {
        jvmti->Deallocate((unsigned char*)methods);
    }
}


void LoadAllMethodIDs(jvmtiEnv* jvmti) {
    jint num_classes;
    jclass* classes;
    if (jvmti->GetLoadedClasses(&num_classes, &classes) == 0) {
        for (int i = 0; i < num_classes; i++) {
            LoadMethodIDs(jvmti, classes[i]);
        }
        jvmti->Deallocate((unsigned char*)classes);
    }
}


void JNICALL ClassLoad(jvmtiEnv *jvmti_env, JNIEnv *jni_env, jthread thread,
		jclass klass) {
}


void JNICALL ClassPrepare(jvmtiEnv *jvmti_env, JNIEnv* jni_env, jthread thread,
		jclass klass) {
	LoadMethodIDs(jvmti_env, klass);
}


void JNICALL CompiledMethodLoad(jvmtiEnv* jvmti, jmethodID method,
		jint code_size, const void* code_addr,
		jint map_length, const jvmtiAddrLocationMap* map,
		const void* compile_info) {
}


bool Agent::Attach() {
	if (is_attached) {
		return true;
	}

	cpu_profiler = new CPUProfiler();
	lock_profiler = new LockProfiler();

	if (dlsym(RTLD_DEFAULT, "AsyncGetCallTrace") == NULL) {
		LogError("AsyncGetCallTrace not available.");
		return false;
	}

	jvmtiError error;

	jvmtiCapabilities potentialCapabilities = {0};
	jvmti->GetPotentialCapabilities(&potentialCapabilities);

	jvmtiCapabilities capabilities = {0};
	jvmtiEventCallbacks callbacks = {0};

	callbacks.ClassLoad = ClassLoad;
	callbacks.ClassPrepare = ClassPrepare;

	if (potentialCapabilities.can_generate_compiled_method_load_events) {
		capabilities.can_generate_compiled_method_load_events = 1;
		callbacks.CompiledMethodLoad = CompiledMethodLoad;
	}
	else {
		LogInfo("Missing capability: can_generate_compiled_method_load_events");
	}

	if (potentialCapabilities.can_get_source_file_name) {
		capabilities.can_get_source_file_name = 1;
	}
	else {
		LogInfo("Missing capability: can_get_source_file_name");
	}

	if (potentialCapabilities.can_get_line_numbers) {
		capabilities.can_get_line_numbers = 1;
	}
	else {
		LogInfo("Missing capability: can_get_line_numbers");
	}

	if (potentialCapabilities.can_generate_monitor_events) {
		capabilities.can_generate_monitor_events = 1;
		callbacks.MonitorContendedEnter = LockProfiler::MonitorContendedEnter;
		callbacks.MonitorContendedEntered = LockProfiler::MonitorContendedEntered;
	}
	else {
		LogInfo("Missing capability: can_generate_monitor_events");
	}

	if (potentialCapabilities.can_tag_objects) {
		capabilities.can_tag_objects = 1;
	}
	else {
		LogInfo("Missing capability: can_tag_objects");
	}
	
	error = jvmti->AddCapabilities(&capabilities);
	CheckJVMTIError(jvmti, error, "AddCapabilities");

	error = jvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
	CheckJVMTIError(jvmti, error, "SetEventCallbacks");

	jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_LOAD, NULL);
	jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_CLASS_PREPARE, NULL);
	jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_COMPILED_METHOD_LOAD, NULL);
	//error = jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_EXCEPTION, NULL);
	CheckJVMTIError(jvmti, error, "SetEventNotificationMode");

	// Initialize methodIDs by calling GetClassMethods, needed for AsyncGetCallTrace to function.
	LoadAllMethodIDs(jvmti);
	jvmti->GenerateEvents(JVMTI_EVENT_DYNAMIC_CODE_GENERATED);
    jvmti->GenerateEvents(JVMTI_EVENT_COMPILED_METHOD_LOAD);

	is_attached = true;
	LogInfo("The JVMTI agent has been attached successfuly.");

	return true;
}


extern "C" JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
	// the agent is not loaded via agentpath, it's loaded from java agent
	return JNI_ERR;
}


extern "C" JNIEXPORT jint JNICALL Agent_OnAttach(JavaVM *jvm, char *options, void *reserved) {
	return JNI_ERR;
}


extern "C" JNIEXPORT void JNICALL Agent_OnUnload(JavaVM *vm) {
	Agent::instance.LogInfo("Agent_OnUnload");
}


extern "C" JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* jvm, void* reserved) {
	Agent* agent = &Agent::instance;
	
	agent->Init(jvm);

	return JNI_VERSION_1_6;
}


extern "C" JNIEXPORT jboolean JNICALL Java_com_stackimpact_agent_Agent_attach(JNIEnv* env, jobject unused, jboolean debug_mode) {
	Agent* agent = &Agent::instance;
	
	agent->jni = env;
	agent->debug_mode = (bool)debug_mode;

	return agent->Attach();
}
