#include <jvmti.h>
#include <stdio.h>

static int events;

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

// Continuously generate CompiledMethodLoad events for all currently compiled methods 
void JNICALL GenerateEventsThread(jvmtiEnv* jvmti, JNIEnv* jni, void* arg) {
    jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_COMPILED_METHOD_LOAD, NULL);

    while (true) {
        events = 0;
        jvmti->GenerateEvents(JVMTI_EVENT_COMPILED_METHOD_LOAD);
        printf("Generated %d events\n", events);
    }
}

// As soon as VM starts, run a separate Agent thread that will generate CompiledMethodLoad events
void JNICALL VMInit(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread) {
    jclass thread_class = jni->FindClass("java/lang/Thread");
    jmethodID thread_constructor = jni->GetMethodID(thread_class, "<init>", "()V");
    jthread agent_thread = jni->NewObject(thread_class, thread_constructor);
    
    jvmti->RunAgentThread(agent_thread, GenerateEventsThread, NULL, JVMTI_THREAD_NORM_PRIORITY);
}

jint Agent_OnLoad(JavaVM* vm, char* options, void* reserved) {
    jvmtiEnv* jvmti;
    vm->GetEnv((void**)&jvmti, JVMTI_VERSION_1_0);

    jvmtiCapabilities capabilities = {0};
    capabilities.can_generate_compiled_method_load_events = 1;
    jvmti->AddCapabilities(&capabilities);

    jvmtiEventCallbacks callbacks = {0};
    callbacks.VMInit = VMInit;
    callbacks.CompiledMethodLoad = CompiledMethodLoad;
    jvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
    jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_INIT, NULL);

    return 0;
}