/*
* Copyright (c) 2017, Oracle and/or its affiliates. All rights reserved.
* ORACLE PROPRIETARY/CONFIDENTIAL. Use is subject to license terms.
*/

#include <stdio.h>
#include <stdlib.h>
#include <memory.h>

#include "jni.h"
#include "jvmti.h"

static jvmtiEnv *jvmti;

jint Agent_Initialize(JavaVM *vm, char *options, void *reserved);

JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
  return Agent_Initialize(jvm, options, reserved);
}

JNIEXPORT jint JNICALL Agent_OnAttach(JavaVM *jvm, char *options, void *reserved) {
  return Agent_Initialize(jvm, options, reserved);
}

JNIEXPORT jint JNICALL Java_ObjectsWithTagsResurrectionBug_setTag0(JNIEnv *env, jclass cls, jobject o, jlong tag) {
  return (*jvmti)->SetTag(jvmti, o, tag);
}

JNIEXPORT jobjectArray JNICALL Java_ObjectsWithTagsResurrectionBug_getObjectsWithTags(JNIEnv *env, jclass this_cls, jint count, jlongArray array) {
  jlong *tags;
  jint i;
  jvmtiError err;
  jobject *object_results;
  jclass cls;
  jobjectArray object_array;
  
  /* copy input list-of-tags from java into C */
  
  tags = (jlong*)malloc(count * sizeof(jlong));
  (*env)->GetLongArrayRegion(env, array, 0, count, tags);
  
  err = (*jvmti)->GetObjectsWithTags(jvmti, count, tags, &count, &object_results, NULL);
  
  /* free the input argument that we malloced */
  free(tags);
  if (err != JVMTI_ERROR_NONE) {
    fprintf(stderr, "ERROR: GetObjectsWithTags failed: %d\n", err);
    return NULL;
  }

  /* copy output from C arrays into Java arrays */
  
  cls = (*env)->FindClass(env, "java/lang/Object");

  object_array = (*env)->NewObjectArray(env, count, cls, NULL);
  
  for (i=0; i<count; i++) {
    (*env)->SetObjectArrayElement(env, object_array, i, object_results[i]);
  }
  
  (*jvmti)->Deallocate(jvmti, (unsigned char*)object_results);
  
  return object_array;
}


/*
 * Agent_Initialize - add capabilities and enables OBJECT_FREE event
 */
jint Agent_Initialize(JavaVM *vm, char *options, void *reserved) {
  jint rc;
  jvmtiError err;
  jvmtiCapabilities capabilities;

  /* get JVMTI environment */

  rc = (*vm)->GetEnv(vm, (void **)&jvmti, JVMTI_VERSION);
  if (rc != JNI_OK) {
    fprintf(stderr, "Unable to create jvmtiEnv, GetEnv failed, error=%d\n", rc);
    return -1;
  }
  
  /* add annotate object capability */
  err = (*jvmti)->GetCapabilities(jvmti, &capabilities);
  if (err != JVMTI_ERROR_NONE) {
    fprintf(stderr, "GetCapabilities failed, error=%d\n", err);
  }
  capabilities.can_tag_objects = 1;
  err = (*jvmti)->AddCapabilities(jvmti, &capabilities);

  if (err != JVMTI_ERROR_NONE) {
    fprintf(stderr, "AddCapabilities failed, error=%d\n", err);
    return -1;
  }
  
  return 0;
}