/* -*- c-file-style: "GNU" -*- */
/*
 * Copyright (C) Telecom SudParis
 * See COPYING in top-level directory.
 */

#ifndef _REENTRANT
#define _REENTRANT
#endif

#include <eztrace-core/eztrace_config.h>
#include <eztrace-core/eztrace_htable.h>
#include <eztrace-lib/eztrace.h>
#include <eztrace-instrumentation/pptrace.h>
#include <eztrace-lib/eztrace_module.h>
#include <semaphore.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <otf2/OTF2_AttributeList.h>
#include <omp.h>
#include <omp-tools.h>

static volatile int _ompt_initialized = 0;

static ompt_get_thread_data_t ompt_get_thread_data;
static ompt_get_unique_id_t ompt_get_unique_id;

#define CURRENT_MODULE ompt
DECLARE_CURRENT_MODULE;

static int openmp_parallel_id = -1;
static int openmp_implicit_task_id = -1;
static int openmp_acquire_mutex_id = -1;
static int openmp_task_id = -1;
static int openmp_for_id = -1;
static int openmp_implicit_barrier_id = -1;
static int openmp_barrier_id = -1;
static int loop_ptr_id = -1;
static struct ezt_hashtable mutex_map;


static void _init_ompt_otf2() {
  /* sometimes, ezt_otf2_register_function returns 1 (because eztrace is being
   * initialized). Thus, we may have to call this function several times untils its
   * succeeds.
   */
  static _Atomic int initialized = 0;
  if(initialized) return;

  if(openmp_parallel_id < 0)
    openmp_parallel_id = ezt_otf2_register_function("OpenMP parallel");

  if(openmp_implicit_task_id < 0)
    openmp_implicit_task_id = ezt_otf2_register_function("OpenMP implicit task");

  if(openmp_acquire_mutex_id < 0)
    openmp_acquire_mutex_id = ezt_otf2_register_function("OpenMP acquire mutex");

  if(openmp_task_id < 0)
    openmp_task_id = ezt_otf2_register_function("OpenMP task");

  if(openmp_for_id < 0)
    openmp_for_id = ezt_otf2_register_function("OpenMP loop");

  if(openmp_implicit_barrier_id < 0)
    openmp_implicit_barrier_id = ezt_otf2_register_function("OpenMP implicit barrier");

  if(openmp_barrier_id < 0)
    openmp_barrier_id = ezt_otf2_register_function("OpenMP barrier");

  if(loop_ptr_id < 0)
    loop_ptr_id = ezt_otf2_register_attribute("Loop pointer", OTF2_TYPE_INT64);

  if( openmp_parallel_id != -1 &&
      openmp_implicit_task_id != -1 &&
      openmp_acquire_mutex_id != -1 &&
      openmp_task_id != -1 &&
      openmp_for_id != -1 &&
      openmp_implicit_barrier_id != -1 &&
      openmp_barrier_id != -1 &&
      loop_ptr_id != -1)
    initialized = 1;
}

static void init_ompt() {
  INSTRUMENT_FUNCTIONS(ompt);

  ezt_hashtable_init(&mutex_map, 1024);

  if (eztrace_autostart_enabled())
    eztrace_start();

  _init_ompt_otf2();

  _ompt_initialized = 1;
}


static void finalize_ompt() {
  _ompt_initialized = 0;
  ezt_hashtable_finalize(&mutex_map);
  eztrace_stop();
}

static void on_ompt_callback_work(ompt_work_t wstype, 
				  ompt_scope_endpoint_t endpoint, 
				  ompt_data_t *parallel_data, 
				  ompt_data_t *task_data, 
				  uint64_t count, 
				  const void *codeptr_ra ) {
  if (EZTRACE_SAFE) {
    if(wstype==ompt_work_loop) {
      _init_ompt_otf2();
      /* TODO: get the source code location that corresponds to codeptr_ra */
      OTF2_AttributeList* attribute_list = OTF2_AttributeList_New();
      OTF2_AttributeList_AddAttribute_uint64(attribute_list, loop_ptr_id, (uint64_t)codeptr_ra);

      if( endpoint==ompt_scope_begin) {
	EZT_OTF2_EvtWriter_Enter(evt_writer, attribute_list, ezt_get_timestamp(),
				 openmp_for_id);
      } else {
	EZT_OTF2_EvtWriter_Leave(evt_writer, attribute_list, ezt_get_timestamp(),
				 openmp_for_id);
      }
    }
    /* shall we do something for the other types of work ? 
       ompt_work_sections
       ompt_work_single_executor
       ompt_work_single_other
       ompt_work_workshare
       ompt_work_distribute
       ompt_work_taskloop
       ompt_work_scope
    */


  }
}

static void on_ompt_callback_sync_region( ompt_sync_region_t kind, 
					  ompt_scope_endpoint_t endpoint, 
					  ompt_data_t *parallel_data, 
					  ompt_data_t *task_data, 
					  const void *codeptr_ra) {
  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    switch(kind) {
    case ompt_sync_region_barrier:
    case ompt_sync_region_barrier_explicit:
    case ompt_sync_region_barrier_implementation:
#if HAVE_OMPT_51
    case ompt_sync_region_barrier_teams:
#endif
      if( endpoint==ompt_scope_begin) {
	EZT_OTF2_EvtWriter_Enter(evt_writer, NULL, ezt_get_timestamp(),
					    openmp_barrier_id);
      } else {
	EZT_OTF2_EvtWriter_Leave(evt_writer, NULL, ezt_get_timestamp(),
					    openmp_barrier_id);
      }

      break;

    case ompt_sync_region_barrier_implicit:
#if HAVE_OMPT_51
    case ompt_sync_region_barrier_implicit_workshare:
    case ompt_sync_region_barrier_implicit_parallel:
#endif
      if( endpoint==ompt_scope_begin) {
	EZT_OTF2_EvtWriter_Enter(evt_writer, NULL, ezt_get_timestamp(),
				 openmp_implicit_barrier_id);
      } else {
	EZT_OTF2_EvtWriter_Leave(evt_writer, NULL, ezt_get_timestamp(),
				 openmp_implicit_barrier_id);
      }
      break;

    case ompt_sync_region_taskwait:
    case ompt_sync_region_taskgroup:
    case ompt_sync_region_reduction:
      /* unsupported */
      break;
    }
  }
}


static void
on_ompt_callback_parallel_begin(ompt_data_t *encountering_task_data,
				const ompt_frame_t *encountering_task_frame,
				ompt_data_t *parallel_data,
				unsigned int requested_parallelism,
				int flags,
				const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;

  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadTeamBegin(evt_writer, NULL, ezt_get_timestamp(),
						  OTF2_UNDEFINED_COMM));
    EZT_OTF2_EvtWriter_Enter(evt_writer, NULL, ezt_get_timestamp(), openmp_parallel_id);
  }
}


static void
on_ompt_callback_parallel_end(ompt_data_t *parallel_data,
			      ompt_data_t *encountering_task_data,
			      int flags,
			      const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;

  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    EZT_OTF2_EvtWriter_Leave(evt_writer, NULL, ezt_get_timestamp(), openmp_parallel_id);
    EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadTeamEnd(evt_writer, NULL, ezt_get_timestamp(),
						OTF2_UNDEFINED_COMM));
  }
}


static void
on_ompt_callback_thread_begin(ompt_thread_t thread_type,
			      ompt_data_t *thread_data) {
  /* OpenMP may spawn thread without calling pthread_create, so we may
     have to initialize the current thread */
  if(thread_status == ezt_trace_status_uninitialized) ezt_init_thread();

  thread_data->value = thread_rank;
}


static void
on_ompt_callback_thread_end(ompt_data_t *thread_data) {
  /* OpenMP may destroy threads without calling pthread_exit, so we
     have to finalize them here*/
  ezt_finalize_thread();
}


static void
on_ompt_callback_task_create(ompt_data_t *encountering_task_data,
			     const ompt_frame_t *encountering_task_frame,
			     ompt_data_t *new_task_data,
			     int flags,
			     int has_dependences,
			     const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;
  new_task_data->value = ompt_get_unique_id();

  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadTaskCreate(evt_writer, NULL, ezt_get_timestamp(),
						   OTF2_UNDEFINED_COMM, tid, new_task_data->value));
  }
}


static void
on_ompt_callback_task_schedule(ompt_data_t *prior_task_data,
			       ompt_task_status_t prior_task_status,
			       ompt_data_t *next_task_data) {
  uint64_t tid = ompt_get_thread_data()->value;
  uint32_t ptid = prior_task_data->value;
  uint32_t ntid = next_task_data->value;

  /* TODO: Add attributes for prior task id and next task id
     The following section has been commented out because of an assertion error in _register_attribute (eztrace_otf2.c:194)
     To reproduce the error, uncomment the commented lines and test the module with any application from https://github.com/bsc-pm/bots
  */


  /*
  OTF2_AttributeList* attr_list = OTF2_AttributeList_New();

  // Create attribute for prior task id
  OTF2_AttributeValue ptid_attr_value;
  OTF2_AttributeRef ptid_attr_id = ezt_otf2_register_attribute("Prior task id", OTF2_TYPE_UINT32);
  ptid_attr_value.uint32 = ptid;
  // Create attribute for next task id
  OTF2_AttributeValue ntid_attr_value;
  OTF2_AttributeRef ntid_attr_id = ezt_otf2_register_attribute("Next task id", OTF2_TYPE_UINT32);
  ntid_attr_value.uint32 = next_task_data->value;
  */

  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
      // Add prior task id to attribute list
      // OTF2_AttributeList_AddAttribute(attr_list, ptid_attr_id, OTF2_TYPE_UINT32, ptid_attr_value);

    EZT_OTF2_EvtWriter_Leave(evt_writer, NULL, ezt_get_timestamp(), openmp_task_id);

    // Add prior and next task ids to attribute list
    // OTF2_AttributeList_AddAttribute(attr_list, ptid_attr_id, OTF2_TYPE_UINT32, ptid_attr_value);
    // OTF2_AttributeList_AddAttribute(attr_list, ntid_attr_id, OTF2_TYPE_UINT32, ntid_attr_value);

    EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadTaskSwitch(evt_writer, NULL, ezt_get_timestamp(),
						   OTF2_UNDEFINED_COMM, tid, ntid));

    // Add next task id to attribute list
    // OTF2_AttributeList_AddAttribute(attr_list, ntid_attr_id, OTF2_TYPE_UINT32, ntid_attr_value);

    EZT_OTF2_EvtWriter_Enter(evt_writer, NULL, ezt_get_timestamp(), openmp_task_id);
  }
}


static void
on_ompt_callback_implicit_task(ompt_scope_endpoint_t endpoint,
			       ompt_data_t *parallel_data,
			       ompt_data_t *task_data,
			       unsigned int actual_parallelism,
			       unsigned int index,
			       int flags) {
  uint64_t tid = ompt_get_thread_data()->value;

  /* TODO: in some cases, implicit tasks may end at the thread
   * destruction (when it's too late for recording an event).
   *
   * This leads to "unbalanced" traces where some of the ENTER events
   * are never matched with a LEAVE event.
   */
  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    if (endpoint == ompt_scope_begin) {
      EZT_OTF2_EvtWriter_Enter(evt_writer, NULL, ezt_get_timestamp(),
			       openmp_implicit_task_id);
    } else if (endpoint == ompt_scope_end) {
      EZT_OTF2_EvtWriter_Leave(evt_writer, NULL, ezt_get_timestamp(),
			       openmp_implicit_task_id);
    }
  }
}


static _Atomic uint32_t next_mutex_id = 0;

struct ezt_omp_mutex_info {
  uint32_t acquisition_order;
  uint32_t mutex_id;
  void* addr;
};


static void
on_ompt_callback_lock_init(ompt_mutex_t kind,
			   unsigned int hint,
			   unsigned int impl,
			   ompt_wait_id_t wait_id,
			   const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;

  struct ezt_omp_mutex_info* l = malloc(sizeof(struct ezt_omp_mutex_info));
  l->acquisition_order = 0;
  l->mutex_id = next_mutex_id++;
  l->addr = (void*)wait_id;

  ezt_hashtable_insert(&mutex_map, hash_function_int64(wait_id), (void*)l);
  assert(ezt_hashtable_get(&mutex_map, hash_function_int64(wait_id)) == l);
}


static void
on_ompt_callback_lock_destroy(ompt_mutex_t kind,
			      unsigned int hint,
			      unsigned int impl,
			      ompt_wait_id_t wait_id,
			      const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;
}

static void
on_ompt_callback_mutex_acquire(ompt_mutex_t kind,
			       unsigned int hint,
			       unsigned int impl,
			       ompt_wait_id_t wait_id,
			       const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;

  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    struct ezt_omp_mutex_info* l = ezt_hashtable_get(&mutex_map, hash_function_int64(wait_id));

    // Add mutex to mutex map in case it hasn't been added by `ompt_lock_init` (e.g. for omp_mutex_critical)
    if (l == NULL) {
      struct ezt_omp_mutex_info* l_new = malloc(sizeof(struct ezt_omp_mutex_info));
      l_new->acquisition_order = 0;
      l_new->mutex_id = next_mutex_id++;
      l_new->addr = (void*)wait_id;
      ezt_hashtable_insert(&mutex_map, hash_function_int64(wait_id), (void*)l_new);
      l = ezt_hashtable_get(&mutex_map, hash_function_int64(wait_id));
    }
    assert(l);

    l->acquisition_order++;

    EZT_OTF2_EvtWriter_Enter(evt_writer, NULL, ezt_get_timestamp(), openmp_acquire_mutex_id);

    EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadAcquireLock(evt_writer, NULL, ezt_get_timestamp(),
						    OTF2_PARADIGM_OPENMP, l->mutex_id,
						    l->acquisition_order));
  }
}


static void
on_ompt_callback_mutex_acquired(ompt_mutex_t kind,
				ompt_wait_id_t wait_id,
				const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;

  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    EZT_OTF2_EvtWriter_Leave(evt_writer, NULL, ezt_get_timestamp(),
			     openmp_acquire_mutex_id);
  }
}


static void
on_ompt_callback_mutex_released(ompt_mutex_t kind,
				ompt_wait_id_t wait_id,
				const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;

  if (EZTRACE_SAFE) {
    _init_ompt_otf2();
    struct ezt_omp_mutex_info* l = ezt_hashtable_get(&mutex_map, hash_function_int64(wait_id));
      assert(l);

      EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadReleaseLock(evt_writer, NULL, ezt_get_timestamp(),
						      OTF2_PARADIGM_OPENMP, l->mutex_id,
						      l->acquisition_order));
  }
}


static void
on_ompt_callback_sync_region_wait(ompt_sync_region_t kind,
				  ompt_scope_endpoint_t endpoint,
				  ompt_data_t *parallel_data,
				  ompt_data_t *task_data,
				  const void *codeptr_ra) {
  uint64_t tid = ompt_get_thread_data()->value;
}


/* No function to intercept since we rely on the OMPT interface */
PPTRACE_START_INTERCEPT_FUNCTIONS(ompt)
PPTRACE_END_INTERCEPT_FUNCTIONS(ompt)


#define register_callback(name, type)				\
  do{                                                           \
    type f_##name = &on_##name;                                 \
    if (ompt_set_callback(name, (ompt_callback_t)f_##name) ==   \
	ompt_set_never)                                         \
      printf("0: Could not register callback '" #name "'\n");   \
  }while(0)


int ompt_initialize(ompt_function_lookup_t lookup,
		    int initial_device_num,
		    ompt_data_t* data) {
  *(double*)(data->ptr)=omp_get_wtime();

  ompt_set_callback_t ompt_set_callback = (ompt_set_callback_t) lookup("ompt_set_callback");
  ompt_get_thread_data = (ompt_get_thread_data_t) lookup("ompt_get_thread_data");
  ompt_get_unique_id = (ompt_get_unique_id_t) lookup("ompt_get_unique_id");

  register_callback(ompt_callback_parallel_begin, ompt_callback_parallel_begin_t);
  register_callback(ompt_callback_parallel_end, ompt_callback_parallel_end_t);
  register_callback(ompt_callback_thread_begin, ompt_callback_thread_begin_t);
  register_callback(ompt_callback_thread_end, ompt_callback_thread_end_t);
  register_callback(ompt_callback_task_create, ompt_callback_task_create_t);
  register_callback(ompt_callback_task_schedule, ompt_callback_task_schedule_t);
  register_callback(ompt_callback_implicit_task, ompt_callback_implicit_task_t);
  register_callback(ompt_callback_sync_region_wait, ompt_callback_sync_region_t);
  register_callback(ompt_callback_mutex_acquire, ompt_callback_mutex_acquire_t);
  register_callback(ompt_callback_lock_destroy, ompt_callback_mutex_acquire_t);
  register_callback(ompt_callback_lock_init, ompt_callback_mutex_acquire_t);
  register_callback(ompt_callback_mutex_acquired, ompt_callback_mutex_t);
  register_callback(ompt_callback_mutex_released, ompt_callback_mutex_t);

  register_callback(ompt_callback_work, ompt_callback_work_t);
  register_callback(ompt_callback_sync_region, ompt_callback_sync_region_t);

  return 1; //success
}


void ompt_finalize(ompt_data_t* data) {
  /* Nothing to do */
}

static void _ompt_init(void) __attribute__((constructor));

static double t=0;
ompt_start_tool_result_t* ompt_start_tool(unsigned int omp_version,
					  const char *runtime_version) {

  _ompt_init();
  t = omp_get_wtime();

  static ompt_start_tool_result_t ompt_start_tool_result = {&ompt_initialize,
							    &ompt_finalize,
							    {.ptr=&t}};
  return &ompt_start_tool_result;
}


static void _ompt_init(void) {
  static int init_called=0;
  if(init_called)
    return;
  init_called=1;

  eztrace_log(dbg_lvl_debug, "eztrace_ompt constructor starts\n");
  EZT_REGISTER_MODULE(ompt, "Module for OpenMP directives, using OMPT",
		      init_ompt, finalize_ompt);
  eztrace_log(dbg_lvl_debug, "eztrace_ompt constructor ends\n");
}

static void _ompt_destructor() __attribute__((destructor));
static void _ompt_destructor() {
  /* The ompt module may be loaded by libompt directly instead of using LD_PRELOAD
   * In this case, the ompt module may be destroyed before eztrace terminaison. This
   * result in eztrace accessing memory regions that were free (eg the "ompt" structure)
   *
   * To avoid that, we add a destructor to make sure eztrace terminates before ompt
   */
  eztrace_stop();
}
