/*****************************************************************************
 * system include files
 ****************************************************************************/

#include <assert.h>

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/*****************************************************************************
 * ompt include files
 ****************************************************************************/

#include "ompt-specific.cpp"

/*****************************************************************************
 * macros
 ****************************************************************************/

#define ompt_get_callback_success 1
#define ompt_get_callback_failure 0

#define no_tool_present 0

#define OMPT_API_ROUTINE static

#ifndef OMPT_STR_MATCH
#define OMPT_STR_MATCH(haystack, needle) (!strcasecmp(haystack, needle))
#endif

/*****************************************************************************
 * types
 ****************************************************************************/

typedef struct {
  const char *state_name;
  ompt_state_t state_id;
} ompt_state_info_t;

enum tool_setting_e {
  omp_tool_error,
  omp_tool_unset,
  omp_tool_disabled,
  omp_tool_enabled
};

typedef void (*ompt_initialize_t)(ompt_function_lookup_t ompt_fn_lookup,
                                  const char *version,
                                  unsigned int ompt_version);

/*****************************************************************************
 * global variables
 ****************************************************************************/

int ompt_enabled = 0;

ompt_state_info_t ompt_state_info[] = {
#define ompt_state_macro(state, code) {#state, state},
    FOREACH_OMPT_STATE(ompt_state_macro)
#undef ompt_state_macro
};

ompt_callbacks_t ompt_callbacks;

static ompt_initialize_t ompt_initialize_fn = NULL;

/*****************************************************************************
 * forward declarations
 ****************************************************************************/

static ompt_interface_fn_t ompt_fn_lookup(const char *s);

OMPT_API_ROUTINE ompt_thread_id_t ompt_get_thread_id(void);

/*****************************************************************************
 * initialization and finalization (private operations)
 ****************************************************************************/

/* On Unix-like systems that support weak symbols the following implementation
 * of ompt_tool() will be used in case no tool-supplied implementation of
 * this function is present in the address space of a process.
 *
 * On Windows, the ompt_tool_windows function is used to find the
 * ompt_tool symbol across all modules loaded by a process. If ompt_tool is
 * found, ompt_tool's return value is used to initialize the tool. Otherwise,
 * NULL is returned and OMPT won't be enabled */
#if OMPT_HAVE_WEAK_ATTRIBUTE
_OMP_EXTERN
__attribute__((weak)) ompt_initialize_t ompt_tool() {
#if OMPT_DEBUG
  printf("ompt_tool() is called from the RTL\n");
#endif
  return NULL;
}

#elif OMPT_HAVE_PSAPI

#include <psapi.h>
#pragma comment(lib, "psapi.lib")
#define ompt_tool ompt_tool_windows

// The number of loaded modules to start enumeration with EnumProcessModules()
#define NUM_MODULES 128

static ompt_initialize_t ompt_tool_windows() {
  int i;
  DWORD needed, new_size;
  HMODULE *modules;
  HANDLE process = GetCurrentProcess();
  modules = (HMODULE *)malloc(NUM_MODULES * sizeof(HMODULE));
  ompt_initialize_t (*ompt_tool_p)() = NULL;

#if OMPT_DEBUG
  printf("ompt_tool_windows(): looking for ompt_tool\n");
#endif
  if (!EnumProcessModules(process, modules, NUM_MODULES * sizeof(HMODULE),
                          &needed)) {
    // Regardless of the error reason use the stub initialization function
    free(modules);
    return NULL;
  }
  // Check if NUM_MODULES is enough to list all modules
  new_size = needed / sizeof(HMODULE);
  if (new_size > NUM_MODULES) {
#if OMPT_DEBUG
    printf("ompt_tool_windows(): resize buffer to %d bytes\n", needed);
#endif
    modules = (HMODULE *)realloc(modules, needed);
    // If resizing failed use the stub function.
    if (!EnumProcessModules(process, modules, needed, &needed)) {
      free(modules);
      return NULL;
    }
  }
  for (i = 0; i < new_size; ++i) {
    (FARPROC &)ompt_tool_p = GetProcAddress(modules[i], "ompt_tool");
    if (ompt_tool_p) {
#if OMPT_DEBUG
      TCHAR modName[MAX_PATH];
      if (GetModuleFileName(modules[i], modName, MAX_PATH))
        printf("ompt_tool_windows(): ompt_tool found in module %s\n", modName);
#endif
      free(modules);
      return ompt_tool_p();
    }
#if OMPT_DEBUG
    else {
      TCHAR modName[MAX_PATH];
      if (GetModuleFileName(modules[i], modName, MAX_PATH))
        printf("ompt_tool_windows(): ompt_tool not found in module %s\n",
               modName);
    }
#endif
  }
  free(modules);
  return NULL;
}
#else
#error Either __attribute__((weak)) or psapi.dll are required for OMPT support
#endif // OMPT_HAVE_WEAK_ATTRIBUTE

void ompt_pre_init() {
  //--------------------------------------------------
  // Execute the pre-initialization logic only once.
  //--------------------------------------------------
  static int ompt_pre_initialized = 0;

  if (ompt_pre_initialized)
    return;

  ompt_pre_initialized = 1;

  //--------------------------------------------------
  // Use a tool iff a tool is enabled and available.
  //--------------------------------------------------
  const char *ompt_env_var = getenv("OMP_TOOL");
  tool_setting_e tool_setting = omp_tool_error;

  if (!ompt_env_var || !strcmp(ompt_env_var, ""))
    tool_setting = omp_tool_unset;
  else if (OMPT_STR_MATCH(ompt_env_var, "disabled"))
    tool_setting = omp_tool_disabled;
  else if (OMPT_STR_MATCH(ompt_env_var, "enabled"))
    tool_setting = omp_tool_enabled;

#if OMPT_DEBUG
  printf("ompt_pre_init(): tool_setting = %d\n", tool_setting);
#endif
  switch (tool_setting) {
  case omp_tool_disabled:
    break;

  case omp_tool_unset:
  case omp_tool_enabled:
    ompt_initialize_fn = ompt_tool();
    if (ompt_initialize_fn) {
      ompt_enabled = 1;
    }
    break;

  case omp_tool_error:
    fprintf(stderr, "Warning: OMP_TOOL has invalid value \"%s\".\n"
                    "  legal values are (NULL,\"\",\"disabled\","
                    "\"enabled\").\n",
            ompt_env_var);
    break;
  }
#if OMPT_DEBUG
  printf("ompt_pre_init(): ompt_enabled = %d\n", ompt_enabled);
#endif
}

void ompt_post_init() {
  //--------------------------------------------------
  // Execute the post-initialization logic only once.
  //--------------------------------------------------
  static int ompt_post_initialized = 0;

  if (ompt_post_initialized)
    return;

  ompt_post_initialized = 1;

  //--------------------------------------------------
  // Initialize the tool if so indicated.
  //--------------------------------------------------
  if (ompt_enabled) {
    ompt_initialize_fn(ompt_fn_lookup, ompt_get_runtime_version(),
                       OMPT_VERSION);

    ompt_thread_t *root_thread = ompt_get_thread();

    ompt_set_thread_state(root_thread, ompt_state_overhead);

    if (ompt_callbacks.ompt_callback(ompt_event_thread_begin)) {
      ompt_callbacks.ompt_callback(ompt_event_thread_begin)(
          ompt_thread_initial, ompt_get_thread_id());
    }

    ompt_set_thread_state(root_thread, ompt_state_work_serial);
  }
}

void ompt_fini() {
  if (ompt_enabled) {
    if (ompt_callbacks.ompt_callback(ompt_event_runtime_shutdown)) {
      ompt_callbacks.ompt_callback(ompt_event_runtime_shutdown)();
    }
  }

  ompt_enabled = 0;
}

/*****************************************************************************
 * interface operations
 ****************************************************************************/

/*****************************************************************************
 * state
 ****************************************************************************/

OMPT_API_ROUTINE int ompt_enumerate_state(int current_state, int *next_state,
                                          const char **next_state_name) {
  const static int len = sizeof(ompt_state_info) / sizeof(ompt_state_info_t);
  int i = 0;

  for (i = 0; i < len - 1; i++) {
    if (ompt_state_info[i].state_id == current_state) {
      *next_state = ompt_state_info[i + 1].state_id;
      *next_state_name = ompt_state_info[i + 1].state_name;
      return 1;
    }
  }

  return 0;
}

/*****************************************************************************
 * callbacks
 ****************************************************************************/

OMPT_API_ROUTINE int ompt_set_callback(ompt_event_t evid, ompt_callback_t cb) {
  switch (evid) {

#define ompt_event_macro(event_name, callback_type, event_id)                  \
  case event_name:                                                             \
    if (ompt_event_implementation_status(event_name)) {                        \
      ompt_callbacks.ompt_callback(event_name) = (callback_type)cb;            \
    }                                                                          \
    return ompt_event_implementation_status(event_name);

    FOREACH_OMPT_EVENT(ompt_event_macro)

#undef ompt_event_macro

  default:
    return ompt_set_result_registration_error;
  }
}

OMPT_API_ROUTINE int ompt_get_callback(ompt_event_t evid, ompt_callback_t *cb) {
  switch (evid) {

#define ompt_event_macro(event_name, callback_type, event_id)                  \
  case event_name:                                                             \
    if (ompt_event_implementation_status(event_name)) {                        \
      ompt_callback_t mycb =                                                   \
          (ompt_callback_t)ompt_callbacks.ompt_callback(event_name);           \
      if (mycb) {                                                              \
        *cb = mycb;                                                            \
        return ompt_get_callback_success;                                      \
      }                                                                        \
    }                                                                          \
    return ompt_get_callback_failure;

    FOREACH_OMPT_EVENT(ompt_event_macro)

#undef ompt_event_macro

  default:
    return ompt_get_callback_failure;
  }
}

/*****************************************************************************
 * parallel regions
 ****************************************************************************/

OMPT_API_ROUTINE ompt_parallel_id_t ompt_get_parallel_id(int ancestor_level) {
  return __ompt_get_parallel_id_internal(ancestor_level);
}

OMPT_API_ROUTINE int ompt_get_parallel_team_size(int ancestor_level) {
  return __ompt_get_parallel_team_size_internal(ancestor_level);
}

OMPT_API_ROUTINE void *ompt_get_parallel_function(int ancestor_level) {
  return __ompt_get_parallel_function_internal(ancestor_level);
}

OMPT_API_ROUTINE ompt_state_t ompt_get_state(ompt_wait_id_t *ompt_wait_id) {
  ompt_state_t thread_state = __ompt_get_state_internal(ompt_wait_id);

  if (thread_state == ompt_state_undefined) {
    thread_state = ompt_state_work_serial;
  }

  return thread_state;
}

/*****************************************************************************
 * threads
 ****************************************************************************/

OMPT_API_ROUTINE void *ompt_get_idle_frame() {
  return __ompt_get_idle_frame_internal();
}

/*****************************************************************************
 * tasks
 ****************************************************************************/

OMPT_API_ROUTINE ompt_thread_id_t ompt_get_thread_id(void) {
  return __ompt_get_thread_id_internal();
}

OMPT_API_ROUTINE ompt_task_id_t ompt_get_task_id(int depth) {
  return __ompt_get_task_id_internal(depth);
}

OMPT_API_ROUTINE ompt_frame_t *ompt_get_task_frame(int depth) {
  return __ompt_get_task_frame_internal(depth);
}

OMPT_API_ROUTINE void *ompt_get_task_function(int depth) {
  return __ompt_get_task_function_internal(depth);
}

/*****************************************************************************
 * placeholders
 ****************************************************************************/

// Don't define this as static. The loader may choose to eliminate the symbol
// even though it is needed by tools.
#define OMPT_API_PLACEHOLDER

// Ensure that placeholders don't have mangled names in the symbol table.
#ifdef __cplusplus
extern "C" {
#endif

OMPT_API_PLACEHOLDER void ompt_idle(void) {
  // This function is a placeholder used to represent the calling context of
  // idle OpenMP worker threads. It is not meant to be invoked.
  assert(0);
}

OMPT_API_PLACEHOLDER void ompt_overhead(void) {
  // This function is a placeholder used to represent the OpenMP context of
  // threads working in the OpenMP runtime.  It is not meant to be invoked.
  assert(0);
}

OMPT_API_PLACEHOLDER void ompt_barrier_wait(void) {
  // This function is a placeholder used to represent the OpenMP context of
  // threads waiting for a barrier in the OpenMP runtime. It is not meant
  // to be invoked.
  assert(0);
}

OMPT_API_PLACEHOLDER void ompt_task_wait(void) {
  // This function is a placeholder used to represent the OpenMP context of
  // threads waiting for a task in the OpenMP runtime. It is not meant
  // to be invoked.
  assert(0);
}

OMPT_API_PLACEHOLDER void ompt_mutex_wait(void) {
  // This function is a placeholder used to represent the OpenMP context of
  // threads waiting for a mutex in the OpenMP runtime. It is not meant
  // to be invoked.
  assert(0);
}

#ifdef __cplusplus
};
#endif

/*****************************************************************************
 * compatability
 ****************************************************************************/

OMPT_API_ROUTINE int ompt_get_ompt_version() { return OMPT_VERSION; }

/*****************************************************************************
 * application-facing API
 ****************************************************************************/

/*----------------------------------------------------------------------------
 | control
 ---------------------------------------------------------------------------*/

_OMP_EXTERN void ompt_control(uint64_t command, uint64_t modifier) {
  if (ompt_enabled && ompt_callbacks.ompt_callback(ompt_event_control)) {
    ompt_callbacks.ompt_callback(ompt_event_control)(command, modifier);
  }
}

/*****************************************************************************
 * API inquiry for tool
 ****************************************************************************/

static ompt_interface_fn_t ompt_fn_lookup(const char *s) {

#define ompt_interface_fn(fn)                                                  \
  if (strcmp(s, #fn) == 0)                                                     \
    return (ompt_interface_fn_t)fn;

  FOREACH_OMPT_INQUIRY_FN(ompt_interface_fn)

  FOREACH_OMPT_PLACEHOLDER_FN(ompt_interface_fn)

  return (ompt_interface_fn_t)0;
}