/*
 * Copyright (C) CNRS, INRIA, Université Bordeaux 1, Télécom SudParis
 * See COPYING in top-level directory.
 */

#include <eztrace-lib/eztrace.h>
#include <eztrace-lib/eztrace_module.h>
#include <starpu.h>
#include <starpu_opencl.h>

#ifdef USE_MPI
#include <starpu_mpi.h>
#endif

#define CURRENT_MODULE starpuv2
DECLARE_CURRENT_MODULE;

/* set to 1 when all the hooks are set.
 * This is usefull in order to avoid recursive calls
 */
static int _starpu_initialized = 0;

int starpu_run_job_id = -1;

static void starpu_register_functions() {
  if(starpu_run_job_id < 0 ) {
    starpu_run_job_id = ezt_otf2_register_function("StarPU job");
  }
}

#define STARPU_ENTER(event_id) do {					\
    if(event_id<0) {							\
      starpu_register_functions();					\
      eztrace_assert(event_id >= 0);					\
    }									\
    EZTRACE_SHOULD_TRACE(OTF2_EvtWriter_Enter(evt_writer,		\
					      NULL,			\
					      ezt_get_timestamp(),	\
					      event_id));		\
  }while(0)

#define STARPU_LEAVE(event_id) do {					\
    if(event_id<0) {							\
      starpu_register_functions();					\
      eztrace_assert(event_id >= 0);					\
    }									\
    EZTRACE_SHOULD_TRACE(OTF2_EvtWriter_Leave(evt_writer,		\
					      NULL,			\
					      ezt_get_timestamp(),	\
					      event_id));		\
  }while(0)


int (*libstarpu_task_bundle_remove) (starpu_task_bundle_t bundle, struct starpu_task* task);

PPTRACE_START_INTERCEPT_FUNCTIONS(starpuv2)
INTERCEPT3("starpu_task_bundle_remove", libstarpu_task_bundle_remove)
PPTRACE_END_INTERCEPT_FUNCTIONS(starpuv2)


#define MAX_STRING_LENGTH 128
#define NB_FUNCTION_MAX 1024
struct starpu_function{
  void* fun_ptr;
  int event_id;
  char name[MAX_STRING_LENGTH];
};

static struct starpu_function functions[NB_FUNCTION_MAX];
static _Atomic int nb_functions = 0;
static pthread_mutex_t lock;

struct starpu_function* get_function(void* fun_ptr, char* task_name) {
  for(int i=0; i<nb_functions; i++) {
    if(functions[i].fun_ptr == fun_ptr)
      return &functions[i];
  }

  pthread_mutex_lock(&lock);

  /* we need to check again in case of a race condition occured */
  for(int i=0; i<nb_functions; i++) {
    if(functions[i].fun_ptr == fun_ptr) {
      pthread_mutex_unlock(&lock);
      return &functions[i];
    }
  }

  int id = nb_functions++;
  if(id > NB_FUNCTION_MAX) {
    eztrace_error("[EZTrace::starpu] Too many functions registered!\n");
  }
  functions[id].fun_ptr = fun_ptr;
  snprintf(functions[id].name, MAX_STRING_LENGTH, "%s", task_name);
  functions[id].event_id = ezt_otf2_register_function(functions[id].name);
  pthread_mutex_unlock(&lock);  
  return &functions[id];
}


void myfunction_cb(struct starpu_prof_tool_info *prof_info,
		   union starpu_prof_tool_event_info *event_info,
		   struct starpu_prof_tool_api_info *api_info) {
	if (NULL == prof_info)
	  return;

	switch (prof_info->event_type)
	{
	case starpu_prof_tool_event_start_cpu_exec:
	  {
	    struct starpu_function* f = get_function(prof_info->fun_ptr, prof_info->task_name);
	    EZT_OTF2_EvtWriter_Enter(evt_writer, NULL, ezt_get_timestamp(), f->event_id);
	    break;
	  }
	case starpu_prof_tool_event_end_cpu_exec:
	  {
	    struct starpu_function* f = get_function(prof_info->fun_ptr, prof_info->task_name);
	    EZT_OTF2_EvtWriter_Leave(evt_writer, NULL, ezt_get_timestamp(), f->event_id);
	    break;
	  }
//	case starpu_prof_tool_event_start_transfer:
//		printf("Start transfer on memnode %ud\n", prof_info->memnode);
//		break;
//	case starpu_prof_tool_event_end_transfer:
//		printf("End transfer on memnode %ud\n", prof_info->memnode);
//		break;
	default:
		printf("Unknown callback %d\n",  prof_info->event_type);
		break;
	}
}

void starpu_prof_tool_library_register(starpu_prof_tool_entry_register_func reg, starpu_prof_tool_entry_register_func unreg) {
  eztrace_log(dbg_lvl_normal, "[EZTrace::starpu] Registering profiling hooks!\n");

  
  enum  starpu_prof_tool_command info = 0;
  //  reg(starpu_prof_tool_event_driver_init, &myfunction_cb, info);
  //  reg(starpu_prof_tool_event_driver_init_start, &myfunction_cb, info);
  //  reg(starpu_prof_tool_event_driver_init_end, &myfunction_cb, info);
  reg(starpu_prof_tool_event_start_cpu_exec, &myfunction_cb, info);
  reg(starpu_prof_tool_event_end_cpu_exec, &myfunction_cb, info);
  //  reg(starpu_prof_tool_event_start_transfer, &myfunction_cb, info);
  //  reg(starpu_prof_tool_event_end_transfer, &myfunction_cb, info);
}

static void init_starpu() {
  INSTRUMENT_FUNCTIONS(starpuv2);
  
  
  if (eztrace_autostart_enabled())
    eztrace_start();
}

static void finalize_starpu() {
  if(_starpu_initialized) {
    _starpu_initialized = 0;
    eztrace_stop();
  }
}

static void _starpu_finalize (void) __attribute__ ((destructor));
static void _starpu_finalize (void) {
  finalize_starpu();
}

static void _starpu_init (void) __attribute__ ((constructor));
/* Initialize the current library */
static void _starpu_init (void) {
  EZT_REGISTER_MODULE(starpuv2, "Module for StarPU functions",
		      init_starpu, finalize_starpu);
}
