/* -*- c-file-style: "GNU" -*- */
/*
 * Copyright (C) CNRS, INRIA, Universite Bordeaux 1, Telecom SudParis
 * See COPYING in top-level directory.
 */

#ifndef __EZTRACE_H__
#define __EZTRACE_H__

#include <dlfcn.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <unistd.h>
#include <inttypes.h>

#include <eztrace-lib/eztrace_otf2.h>
#include <eztrace-lib/eztrace_internals.h>
#include <eztrace-core/eztrace_macros.h>
#include <eztrace-core/eztrace_config.h>
#include <eztrace-core/eztrace_sampling.h>
#include <eztrace-core/eztrace_types.h>


// #define EZTRACE_API_VERSION 0x00000900 // up to eztrace-1.x

// With eztrace 2.0, the API changed
#define EZTRACE_API_VERSION 0x00002000

/* start the recording of events */
void eztrace_start();
/* stop recording events and write the trace to the disk */
void eztrace_stop();
/* start eztrace but only record some event (thread creation, mpi_init, mpi communicators creation..)*/
void eztrace_silent_start();

/* Stop the recording of events  */
void eztrace_pause();
/* Restart the recording of events*/
void eztrace_resume();


/* unset LD_PRELOAD
 * this makes sure that forked processes will not be instrumented with eztrace
 */
void unset_ld_preload();

/* set LD_PRELOAD so that future forked processes are instrumented
 * you need to call unset_ld_preload before calling this function
 */
void reset_ld_preload();

/* change the trace filename */
void eztrace_set_filename();

extern int eztrace_can_trace;

/* record events (code=EZTRACE_CALLING_FUNCTION) containing the backtrace */
void eztrace_record_backtrace(int backtrace_depth);
void record_backtrace();
void eztrace_abort();

/* get information on frame frameid
 * the collected information is copied into output_str
 * buffer_len is the maximum number of bytes written to output_str
 */
void eztrace_get_stack_frame(int frameid, char* output_str, int buffer_len);

/* initialize a thread */
void ezt_init_thread();
/* finalize a thread */
void ezt_finalize_thread();
void ezt_pthread_first_event();

/* return 1 if eztrace is in MPI mode */
int ezt_is_mpi_mode();

/* return the number threads that were created */
int ezt_get_nb_threads();

struct ezt_thread_info {
  int thread_rank;
  pid_t tid;
  char thread_name[50];
  int otf2_thread_id;
};
struct ezt_thread_info* get_thread_info_by_pid(pid_t tid);
struct ezt_thread_info* get_thread_info(int thread_id);


typedef void (*eztrace_function_t)(void);

/* register a callback to be called when the application calls
 * eztrace_start().
 * This function is only usefull when AUTOSTART is disabled
 */
void eztrace_register_init_routine(eztrace_function_t init_func);

typedef void (*eztrace_atexit_function_t)(void*);

/*
 * register a function to be called before eztrace_stop.
 */
void eztrace_atexit(eztrace_atexit_function_t f, void* param);

#define RECORD_HW_COUNTERS()			\
  {						\
    ezt_sampling_check_callbacks();		\
  }

enum ezt_trace_status {
		       ezt_trace_status_uninitialized,
		       ezt_trace_status_running,
		       ezt_trace_status_paused,
		       ezt_trace_status_stopped,
		       ezt_trace_status_being_finalized, /* atexit functions are being called */
		       ezt_trace_status_finalized,       /* event recording has stopped */
};

typedef void (*ezt_finalize_callback_t)(OTF2_EvtWriter*, enum ezt_trace_status *, uint64_t);
/* register a callback to be called during ezt_otf2_finalize before closing the OTF2 archive
 * this can be used for finalizing a thread
 */
void ezt_at_finalize(ezt_finalize_callback_t callback,
		     OTF2_EvtWriter* evt_writer,
		     enum ezt_trace_status *thread_status,
		     uint64_t thread_id);

/* unregister a callback */
void ezt_at_finalize_cancel(uint64_t thread_id);

/* unregister a callback 
 * This function assumes that at_finalize_lock is locked
 */
void ezt_at_finalize_cancel_locked(uint64_t thread_id);

/* run the finalization callback for thread thread_id */
void ezt_at_finalize_run(uint64_t thread_id);

#define EZT_NB_MODULES_MAX 32

enum eztrace_debug_level {
  dbg_lvl_error,
  dbg_lvl_quiet,
  dbg_lvl_normal,
  dbg_lvl_verbose,
  dbg_lvl_debug,
  dbg_lvl_max,
};

struct _ezt_write_trace {
  OTF2_Archive* archive;
  OTF2_GlobalDefWriter* global_def_writer;
  OTF2_DefWriter* def_writer;
  enum ezt_trace_status status;
  enum eztrace_debug_level debug_level;
  int timestamp_config;
  size_t buffer_size;
  char* filename;
  /* number of init functions to be called when the eztrace_start
   * is called.
   */
  int nb_module;
  eztrace_function_t init_routines[EZT_NB_MODULES_MAX];
};
extern struct _ezt_write_trace _ezt_trace;
extern int* _ez_timestamp_config; /* points to _ezt_trace->timestamp_config */
extern __thread OTF2_EvtWriter* evt_writer;
extern __thread enum ezt_trace_status thread_status;
extern __thread uint64_t thread_rank;
extern __thread uint64_t otf2_thread_id;
extern int ezt_mpi_rank;

extern int system_tree_node_ref;
extern int eztrace_should_trace;
  
/* is it safe to record events ? */
#define EZTRACE_SAFE ((_ezt_trace.status == ezt_trace_status_running ||	\
		      _ezt_trace.status == ezt_trace_status_being_finalized) && \
		      thread_status == ezt_trace_status_running &&     \
		      eztrace_should_trace)

#define EZTRACE_SHOULD_TRACE(x) do { if(EZTRACE_SAFE) x; } while(0)


#define EZT_OTF2_EvtWriter_Enter(writer, attr, time, region) do {	\
    if(region<0) {							\
      fprintf(stderr, "error in %s:%d region=%d\n", __FILE__, __LINE__, region); \
      eztrace_abort();								\
    }									\
    if(EZTRACE_SAFE){							\
      EZT_OTF2_CHECK(OTF2_EvtWriter_Enter(writer, attr, time, region));	\
    }									\
  }while(0)

#define EZT_OTF2_EvtWriter_Leave(writer, attr, time, region) do {	\
    if(region<0) {							\
      fprintf(stderr, "error in %s:%d region=%d\n", __FILE__, __LINE__, region); \
      eztrace_abort();								\
    }									\
    if(EZTRACE_SAFE){							\
      EZT_OTF2_CHECK(OTF2_EvtWriter_Leave(writer, attr, time, region));	\
    }									\
  }while(0)

int _eztrace_fd();
#define eztrace_log(_debug_level_, format, ...)				\
  do {									\
    if (_ezt_trace.debug_level >= _debug_level_) {			\
      dprintf(_eztrace_fd(), "[P%dT%" PRIu64 "] " format, ezt_mpi_rank, thread_rank, ##__VA_ARGS__); \
    }									\
} while(0)

#define eztrace_warn(format, ...) \
  do {									\
  eztrace_log(dbg_lvl_normal, "EZTrace warning in %s (%s:%d): " format,	\
	      __func__, __FILE__, __LINE__,				\
	      ##__VA_ARGS__);						\
  } while(0)

#define eztrace_error(format, ...)					\
  do {									\
    eztrace_log(dbg_lvl_error, "EZTrace error in %s (%s:%d): " format,	\
		__func__, __FILE__, __LINE__,				\
		##__VA_ARGS__);						\
    eztrace_abort();								\
  } while(0)

#ifdef NDEBUG
#define eztrace_assert(cond) do { (void)sizeof(cond);} while (0)
#else
#define eztrace_assert(cond)			\
  do {						\
    if(!(cond))					\
      eztrace_error("Assertion failed");	\
  } while(0)
#endif

#define FUNCTION_ENTRY_(fname)						\
  static __thread int recursion_shield=0;				\
  eztrace_log(dbg_lvl_verbose, "Entering [%s]\n", fname);		\
  static struct ezt_instrumented_function* function=NULL;		\
									\
  recursion_shield++;							\
  if(recursion_shield == 1 &&						\
     eztrace_can_trace &&						\
     _ezt_trace.status == ezt_trace_status_running &&			\
     thread_status == ezt_trace_status_running) {			\
    EZTRACE_PROTECT {							\
      EZTRACE_PROTECT_ON();						\
      if(!function) {							\
	function=find_instrumented_function(fname, INSTRUMENTED_FUNCTIONS); \
      }									\
      if(function->event_id < 0) {					\
	instrument_function(function);					\
      }									\
      assert(function->event_id >= 0);					\
      EZT_OTF2_EvtWriter_Enter( evt_writer,				\
				NULL,					\
				ezt_get_timestamp(),			\
				function->event_id );			\
      EZTRACE_PROTECT_OFF();						\
    }									\
  }


#define FUNCTION_ENTRY	FUNCTION_ENTRY_(__func__)

#define FUNCTION_ENTRY_WITH_ARGS(...)					\
  static __thread int recursion_shield=0;				\
  eztrace_log(dbg_lvl_verbose, "Entering [%s]\n", __func__);		\
  RECORD_HW_COUNTERS();							\
  static struct ezt_instrumented_function* function=NULL;		\
  static int entry_attr_id[__VA_NARGS__(__VA_ARGS__)];			\
  recursion_shield++;							\
									\
  if(recursion_shield == 1 &&						\
     eztrace_can_trace &&						\
     _ezt_trace.status == ezt_trace_status_running &&			\
     thread_status == ezt_trace_status_running) {			\
    EZTRACE_PROTECT {							\
      EZTRACE_PROTECT_ON();						\
      if(!function) {							\
	function=find_instrumented_function(__func__, INSTRUMENTED_FUNCTIONS); \
      }									\
      if(function->event_id < 0) {					\
	instrument_function(function);					\
      }									\
      assert(function->event_id>=0);					\
      static int first_time = 1;					\
      if(first_time) {							\
	FOR_EACH_I(REGISTER_ARG_ENTRY, __VA_ARGS__);			\
	first_time = 0;							\
      }									\
      OTF2_AttributeList* attribute_list = OTF2_AttributeList_New();	\
      FOR_EACH_I(ADD_ATTR_ENTRY, __VA_ARGS__);				\
  									\
      EZT_OTF2_EvtWriter_Enter( evt_writer,				\
				attribute_list,				\
				ezt_get_timestamp(),			\
				function->event_id );			\
      OTF2_AttributeList_Delete(attribute_list);			\
      EZTRACE_PROTECT_OFF();						\
    }									\
  }

#define FUNCTION_EXIT_(fname)					\
  eztrace_log(dbg_lvl_verbose, "Leaving [%s]\n", fname);	\
  recursion_shield--;						\
  if(recursion_shield == 0 &&					\
     eztrace_can_trace &&					\
     _ezt_trace.status == ezt_trace_status_running &&		\
     thread_status == ezt_trace_status_running) {		\
  EZTRACE_PROTECT {						\
    EZTRACE_PROTECT_ON();					\
    assert(function);						\
    assert(function->event_id >= 0);				\
    EZT_OTF2_EvtWriter_Leave( evt_writer,			\
			      NULL,				\
			      ezt_get_timestamp(),		\
			      function->event_id );		\
    EZTRACE_PROTECT_OFF();					\
  }								\
  }

#define FUNCTION_EXIT FUNCTION_EXIT_(__func__)

#define FUNCTION_EXIT_WITH_ARGS(...)   do {				\
    eztrace_log(dbg_lvl_verbose, "Leaving [%s]\n", __func__);		\
    static int exit_attr_id[__VA_NARGS__(__VA_ARGS__)];			\
    static int exit_attr_initialized = 0;				\
    if(recursion_shield == 1 &&						\
       eztrace_can_trace &&						\
       _ezt_trace.status == ezt_trace_status_running &&			\
       thread_status == ezt_trace_status_running) {			\
      EZTRACE_PROTECT {							\
	EZTRACE_PROTECT_ON();						\
	if(exit_attr_initialized == 0) {				\
	  exit_attr_initialized = 1;					\
	  FOR_EACH_I(REGISTER_ARG_EXIT, __VA_ARGS__);			\
	}								\
	assert(function);						\
	assert(function->event_id >= 0);				\
	OTF2_AttributeList* attribute_list = OTF2_AttributeList_New();	\
	FOR_EACH_I(ADD_ATTR_EXIT, __VA_ARGS__);				\
	EZT_OTF2_EvtWriter_Leave( evt_writer,				\
				  attribute_list,			\
				  ezt_get_timestamp(),			\
				  function->event_id );			\
	OTF2_AttributeList_Delete(attribute_list);			\
	EZTRACE_PROTECT_OFF();						\
      }									\
    }									\
    recursion_shield--;							\
  } while(0)

/* check whether dlsym returned successfully */
#define TREAT_ERROR()				\
  do {						\
    char* error;				\
    if ((error = dlerror()) != NULL) {		\
      fputs(error, stderr);			\
      eztrace_abort();				\
    }						\
  } while (0)

/* intercept function func and store its previous value into var */
#define INTERCEPT(func, var)					\
  do {								\
    if (var)							\
      break;							\
    void* _handle = RTLD_NEXT;					\
    var = (typeof(var))(uintptr_t)dlsym(_handle, func);	\
    TREAT_ERROR();						\
  } while (0)

/* record an event (code=EZTRACE_CALLING_FUNCTION) with the calling function name */
void record_backtrace();

/* return 1 if an event is being recorded. */
int recursion_shield_on();

void set_recursion_shield_on();
void set_recursion_shield_off();

/* avoid infinite recursion */
#define EZTRACE_PROTECT if (!recursion_shield_on())

/* Set the recursion shield */
#define EZTRACE_PROTECT_ON() set_recursion_shield_on()
#define EZTRACE_PROTECT_OFF() set_recursion_shield_off()

/* pptrace stuff */

#define DECLARE_MODULE(module_name)					\
  extern struct ezt_instrumented_function PPTRACE_SYMBOL_LIST(module_name)[]

#define DECLARE_CURRENT_MODULE DECLARE_MODULE(CURRENT_MODULE)

//#define CURRENT_MODULE lib
#define INSTRUMENTED_FUNCTIONS PPTRACE_CURRENT_SYMBOL_LIST

#define PPTRACE_START_INTERCEPT_FUNCTIONS(module_name) struct ezt_instrumented_function PPTRACE_SYMBOL_LIST(module_name)[] = {

  
#define PPTRACE_END_INTERCEPT_FUNCTIONS(module_name)			\
  FUNCTION_NONE								\
  }									\
    ;									\
  extern struct ezt_instrumented_function* PPTRACE_SYMBOL_EXTERNAL(module_name) __attribute__((alias(PPTRACE_SYMBOL_ALIAS(module_name))));

#define INTERCEPT3(func, var) {			\
    .function_name=func,			\
      .callback=&(var),				\
      .event_id = -1,				\
      },

#define FUNCTION_NONE  {			\
    .function_name="",				\
      .callback=NULL,				\
      .event_id = -1,				\
      },

static void instrument_function(struct ezt_instrumented_function* f) __attribute__((unused));
static void instrument_functions(struct ezt_instrumented_function* functions) __attribute__((unused));
static struct ezt_instrumented_function* find_instrumented_function(const char* fname, struct ezt_instrumented_function* functions) __attribute__((unused));

static void instrument_function(struct ezt_instrumented_function* f) {
  
  if(f->event_id >= 0) {
    /* this function has already been initialized */
    return;
  }

  assert(f->callback != NULL);

  static __thread int recursion_shield = 0;
  recursion_shield++;
  if(recursion_shield == 1) {
    if(*(void**)f->callback == NULL) {
      eztrace_log(dbg_lvl_debug, "Instrumenting %s using dlsym\n", f->function_name);
      /* binary instrumentation did not find the symbol. */
      void* ptr = dlsym(RTLD_NEXT, f->function_name);
      if(!ptr) {
	eztrace_log(dbg_lvl_debug, "Warning: cannot find symbol %s: %s\n", f->function_name, dlerror());
      } else {
	memcpy(f->callback, &ptr, sizeof(void*));
      }
    } else {
      eztrace_log(dbg_lvl_debug, "No need to instrument %s because of binary instrumentation\n", f->function_name);
    }

    if(todo_get_status("eztrace") == init_complete &&
       todo_get_status("ezt_otf2") == init_complete &&
       _ezt_trace.status < ezt_trace_status_being_finalized) {
      f->event_id = ezt_otf2_register_function(f->function_name);
    }
  }
  recursion_shield--;
}

static struct ezt_instrumented_function* find_instrumented_function(const char* fname, struct ezt_instrumented_function* functions) {
  struct ezt_instrumented_function*f=NULL;
  for(f = functions;
      strcmp(f->function_name, "") != 0;
      f++) {
    if(strcmp(f->function_name, fname) == 0) {
      return f;
    }
  }
  return NULL;
}

static void instrument_functions(struct ezt_instrumented_function* functions) {
  struct ezt_instrumented_function*f=NULL;

  for(f = functions;
      strcmp(f->function_name, "") != 0;
      f++) {
    instrument_function(f);
  }
}

#define INSTRUMENT_ALL_FUNCTIONS() do {			\
    EZTRACE_PROTECT {					\
      EZTRACE_PROTECT_ON();				\
      instrument_functions(INSTRUMENTED_FUNCTIONS);	\
      EZTRACE_PROTECT_OFF();				\
    }							\
  }while(0);

#define INSTRUMENT_FUNCTIONS(module) do {			\
    EZTRACE_PROTECT {						\
      EZTRACE_PROTECT_ON();					\
      instrument_functions(PPTRACE_SYMBOL_LIST(module));	\
      EZTRACE_PROTECT_OFF();					\
    }								\
  } while(0)


/* instrument one function */
#define INTERCEPT_FUNCTION(fname, cb) do {		\
    if(!cb) {						\
      struct ezt_instrumented_function*f=NULL;		\
      for(f = INSTRUMENTED_FUNCTIONS;			\
	  strcmp(f->function_name, fname) != 0 &&	\
	    strcmp(f->function_name, "") != 0;		\
	  f++) {					\
      }							\
      instrument_function(f);				\
    }							\
  } while(0)

#endif /* __EZTRACE_H__ */
