/*
 * Copyright (c) 2016, 2017, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

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

#ifdef __cplusplus
extern "C" {
#endif

#ifndef JNI_ENV_ARG

#ifdef __cplusplus
#define JNI_ENV_ARG(x, y) y
#define JNI_ENV_PTR(x) x
#else
#define JNI_ENV_ARG(x,y) x, y
#define JNI_ENV_PTR(x) (*x)
#endif

#endif

#define PASSED 0
#define FAILED 2

#define MAX_TRACES 400

static const char *EXC_CNAME = "java/lang/Exception";
static jvmtiEnv *jvmti = NULL;

static
jint throw_exc(JNIEnv *env, char *msg) {
    jclass exc_class = JNI_ENV_PTR(env)->FindClass(JNI_ENV_ARG(env, EXC_CNAME));

    if (exc_class == NULL) {
        printf("throw_exc: Error in FindClass(env, %s)\n", EXC_CNAME);
        return -1;
    }
    return JNI_ENV_PTR(env)->ThrowNew(JNI_ENV_ARG(env, exc_class), msg);
}

static jint Agent_Initialize(JavaVM *jvm, 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 JNI_OnLoad(JavaVM *jvm, void *reserved) {
    return JNI_VERSION_1_8;
}

JNIEXPORT void JNICALL OnVMInit(jvmtiEnv *jvmti, JNIEnv *jni_env, jthread thread) {
  (*jvmti)->StartHeapSampling(jvmti, 1<<19, MAX_TRACES);
}

JNIEXPORT void JNICALL OnClassLoad(jvmtiEnv *jvmti_env, JNIEnv *jni_env,
                                              jthread thread, jclass klass) {
  // NOP.
}

JNIEXPORT void JNICALL OnClassPrepare(jvmtiEnv *jvmti_env, JNIEnv *jni_env,
                                      jthread thread, jclass klass) {
  // We need to do this to "prime the pump", as it were -- make sure
  // that all of the methodIDs have been initialized internally, for
  // AsyncGetCallTrace.
  jint method_count;
  jmethodID *methods = 0;
  jvmtiError err = (*jvmti)->GetClassMethods(jvmti, klass, &method_count, &methods);
  if ((err != JVMTI_ERROR_NONE) && (err != JVMTI_ERROR_CLASS_NOT_PREPARED)) {
    // JVMTI_ERROR_CLASS_NOT_PREPARED is okay because some classes may
    // be loaded but not prepared at this point.
    throw_exc(jni_env, "Failed to create method IDs for methods in class\n");
  }
}

static
jint Agent_Initialize(JavaVM *jvm, char *options, void *reserved) {
  jint res;

  res = JNI_ENV_PTR(jvm)->GetEnv(JNI_ENV_ARG(jvm, (void **) &jvmti),
                                 JVMTI_VERSION_9);
  if (res != JNI_OK || jvmti == NULL) {
    printf("    Error: wrong result of a valid call to GetEnv!\n");
    return JNI_ERR;
  }

  jvmtiEventCallbacks callbacks;
  memset(&callbacks, 0, sizeof(callbacks));

  callbacks.VMInit = &OnVMInit;
  callbacks.ClassLoad = &OnClassLoad;
  callbacks.ClassPrepare = &OnClassPrepare;

  jvmtiCapabilities caps;
  memset(&caps, 0, sizeof(caps));
  // Get line numbers and filename for the test.
  caps.can_get_line_numbers = 1;
  caps.can_get_source_file_name = 1;
  int ernum = (*jvmti)->AddCapabilities(jvmti, &caps);

  ernum = (*jvmti)->SetEventCallbacks(jvmti, &callbacks, sizeof(jvmtiEventCallbacks));
  ernum = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE, JVMTI_EVENT_VM_INIT, NULL);
  ernum = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE, JVMTI_EVENT_CLASS_LOAD, NULL);
  ernum = (*jvmti)->SetEventNotificationMode(jvmti, JVMTI_ENABLE, JVMTI_EVENT_CLASS_PREPARE, NULL);

  return JNI_OK;
}

// Given a method and a location, this method gets the line number.
// Kind of expensive, comparatively.
jint GetLineNumber(jvmtiEnv *jvmti, jmethodID method, jlocation location) {
  // The location is -1 if the bci isn't known or -3 for a native method.
  if (location == -1 || location == -3) {
    return -1;
  }

  // Read the line number table.
  jvmtiLineNumberEntry *table_ptr = 0;
  jint line_number_table_entries;
  int jvmti_error = (*jvmti)->GetLineNumberTable(jvmti, method,
                                                 &line_number_table_entries,
                                                 &table_ptr);

  if (JVMTI_ERROR_NONE != jvmti_error) {
    return -1;
  }
  if (line_number_table_entries <= 0) {
    return -1;
  }
  if (line_number_table_entries == 1) {
    return table_ptr[0].line_number;
  }

  // Go through all the line numbers...
  jint last_location = table_ptr[0].start_location;
  int l;
  for (l = 1; l < line_number_table_entries; l++) {
    // ... and if you see one that is in the right place for your
    // location, you've found the line number!
    if ((location < table_ptr[l].start_location) &&
        (location >= last_location)) {
      return table_ptr[l - 1].line_number;
    }
    last_location = table_ptr[l].start_location;
  }

  if (location >= last_location) {
    return table_ptr[line_number_table_entries - 1].line_number;
  } else {
    return -1;
  }
}

typedef struct _ExpectedContentFrame {
  char *name;
  char *signature;
  char *file_name;
  int line_number;
} ExpectedContentFrame;

static jint CheckSampleContent(JNIEnv *env,
                               jvmtiStackTrace *trace,
                               ExpectedContentFrame *expected,
                               int expected_count) {
  int i;

  if (expected_count > trace->frame_count) {
    return FAILED;
  }

  for (i = 0; i < expected_count; i++) {
    // Get basic information out of the trace.
    int bci = trace->frames[i].bci;
    jmethodID methodid = trace->frames[i].method_id;
    char *name = NULL, *signature = NULL, *file_name = NULL;

    if (bci < 0) {
      return FAILED;
    }

    // Transform into usable information.
    int line_number = GetLineNumber(jvmti, methodid, bci);
    (*jvmti)->GetMethodName(jvmti, methodid, &name, &signature, 0);

    jclass declaring_class;
    if (JVMTI_ERROR_NONE !=
        (*jvmti)->GetMethodDeclaringClass(jvmti, methodid, &declaring_class)) {
      return FAILED;
    }

    jvmtiError err = (*jvmti)->GetSourceFileName(jvmti, declaring_class,
                                                 &file_name);
    if (err != JVMTI_ERROR_NONE) {
      return FAILED;
    }

    // Compare now, none should be NULL.
    if (name == NULL) {
      return FAILED;
    }

    if (file_name == NULL) {
      return FAILED;
    }

    if (signature == NULL) {
      return FAILED;
    }

    if (strcmp(name, expected[i].name) ||
        strcmp(signature, expected[i].signature) ||
        strcmp(file_name, expected[i].file_name) ||
        line_number != expected[i].line_number) {
      return FAILED;
    }
  }

  return PASSED;
}

static jint CheckSamples(JNIEnv* env, jvmtiStackTrace* traces, int trace_count) {
  // Only look at the first few frames, underneath we go under main and it is
  // less interesting.

  // TODO: There is an issue with the stacktrace walker and I think it is
  // getting confused, so this is just to show what the test will look like.
  // Basically the line 61 is not yet accuracte, I have a fix for it but need to
  // double check.
  // The full system is not yet functional.
  ExpectedContentFrame expected_contents[] = {
    { "helper", "()I", "HeapMonitorTest.java", 61 },
    { "main", "([Ljava/lang/String;)V", "HeapMonitorTest.java", 69 },
  };

  const int expected_contents_count =
      sizeof(expected_contents) / sizeof(expected_contents[0]);

  // We also expect the code to record correctly the BCI, retrieve the line
  // number, have the right method and the class name of the first frames.
  int i;
  for (i = 0; i < trace_count; i++) {
    jvmtiStackTrace *trace = traces + i;
    int j;
    for (j = 0; j < trace->frame_count; j++) {
      // Get basic information out of the trace.
      int bci = trace->frames[j].bci;
      jmethodID methodid = trace->frames[j].method_id;
      char *name = NULL, *signature = NULL, *file_name = NULL;

      // Transform into usable information.
      int line_number = GetLineNumber(jvmti, methodid, bci);
      (*jvmti)->GetMethodName(jvmti, methodid, &name, &signature, 0);

      jclass declaring_class = 0;
      (*jvmti)->GetMethodDeclaringClass(jvmti, methodid, &declaring_class);

      if (declaring_class) {
        jvmtiError err = (*jvmti)->GetSourceFileName(jvmti, declaring_class, &file_name);
      }
    }
  }

  for (i = 0; i < trace_count; i++) {
    jvmtiStackTrace *trace = traces + i;
    if (CheckSampleContent(env, trace, expected_contents,
                           expected_contents_count) == PASSED) {
      // At least one frame matched what we were looking for.
      return PASSED;
    }
  }

  return FAILED;
}

static jint GetTraces(JNIEnv* env) {
  jvmtiStackTraces fulltraces;

  // Call the JVMTI interface to get the traces.
  (*jvmti)->GetLiveTraces(jvmti, &fulltraces);

  if (CheckSamples(env, fulltraces.stack_traces,
                   fulltraces.trace_count) == FAILED) {
    return FAILED;
  }

  // Release Traces.
  (*jvmti)->ReleaseTraces(jvmti, &fulltraces);

  return PASSED;
}

JNIEXPORT jint JNICALL
Java_MyPackage_HeapMonitorTest_check(JNIEnv *env) {
    jobject loader = NULL;

    if (jvmti == NULL) {
        throw_exc(env, "JVMTI client was not properly loaded!\n");
        return FAILED;
    }

    return GetTraces(env);
}

#ifdef __cplusplus
}
#endif