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

#include "eztrace-lib/eztrace.h"

#include "eztrace-lib/eztrace_otf2.h"
#include "eztrace-lib/eztrace_mpi.h"
#include "eztrace-lib/eztrace_module.h"
#include "eztrace-core/eztrace_attributes.h"
#include "eztrace-core/eztrace_config.h"
#include "eztrace-core/eztrace_spinlock.h"
#include "eztrace-core/eztrace_utils.h"

#include <otf2/otf2.h>
#include <otf2/OTF2_Pthread_Locks.h>
#include <limits.h>
#include <sys/wait.h>
#include <stdatomic.h> 
#include <inttypes.h>
#include <errno.h>

__thread OTF2_EvtWriter* evt_writer=NULL;
 // set to 1 during OTF2 startup. In this case, it is safe to record
 // OTF2 events even if EZTRACE_SAFE is false
static __thread int _ezt_otf2_safe = 0;
extern int eztrace_should_trace;
int ezt_otf2_set_mpi_rank_called = 0;

static int using_mpi = 0;
int ezt_mpi_rank = 0;
static int mpi_comm_size = 1;
static _Atomic int next_string_id = 0;
static _Atomic int next_region_id = 0;
static _Atomic int next_thread_id = 0;
static _Atomic int next_local_thread_id = 0;
static _Atomic int next_attribute_id = 0;
static _Atomic int next_thread_team_id = 1;
static _Atomic int next_mpi_comm_id = 2;

OTF2_CommRef comm_world_ref = 0;
static int max_id = INT_MAX;
int otf2_chunk_size=1;

static uint64_t epoch_start;
static uint64_t epoch_stop;

struct ezt_internal_module* ezt_otf2_module = NULL;

static pthread_mutex_t otf2_lock;
static struct timespec pre_flush_ts;
OTF2_TimeStamp first_timestamp = 0;


static OTF2_FlushType pre_flush( void*            userData __attribute__((unused)),
				 OTF2_FileType    fileType __attribute__((unused)),
				 OTF2_LocationRef location __attribute__((unused)),
				 void*            callerData __attribute__((unused)),
				 bool             final  __attribute__((unused))) {
  clock_gettime(CLOCK_MONOTONIC, &pre_flush_ts);
  eztrace_log(dbg_lvl_verbose, "Start flushing\n");


  // OTF2 currently does not *always* call post_flush after flushing.
  // When closing the trace file (when final is true), the post flush
  // callback will not be called.
  
  // In that case, do not pause eztrace because the thread would not
  // be resumed
  if(!final)
    thread_status = ezt_trace_status_paused;
  return OTF2_FLUSH;
}

#define TIME_DIFF(t1, t2) (((t2).tv_sec - (t1).tv_sec)*1e9 +((t2).tv_nsec - (t1).tv_nsec))
static OTF2_TimeStamp post_flush( void*            userData __attribute__((unused)),
				  OTF2_FileType    fileType __attribute__((unused)),
				  OTF2_LocationRef location __attribute__((unused)) ) {
  /* TODO: what if flush was called at the end ? */
  static struct timespec post_flush_ts;
  clock_gettime(CLOCK_MONOTONIC, &post_flush_ts);
  static int nb_flush=0;

  thread_status = ezt_trace_status_running;
  eztrace_log(dbg_lvl_verbose, "Flush %d ends! (%f us)\n", nb_flush++, TIME_DIFF(pre_flush_ts, post_flush_ts)/1e6);
  return ezt_get_timestamp();
}

static OTF2_FlushCallbacks flush_callbacks =  {
    .otf2_pre_flush  = pre_flush,
    .otf2_post_flush = post_flush
};



struct to_register_string {
  int id;
  int len;
  char* buffer;
};

struct to_register_function {
  int id;
  int function_name;
};

struct to_register_thread {
  int id;
  int string_id;
  int mpi_rank;
};

struct to_register_gpu {
  int id;
  int string_id;
  int mpi_rank;
};

struct to_register_attribute {
  int id;
  int string_id;
  OTF2_Type type;    
};

struct to_register_thread_team {
  int id;
  int string_id;
  int parent_id;
  int nb_threads;
};

struct to_register_thread_team_member {
  int thread_team_id;
  int thread_rank;
  uint64_t thread_id;
  int nb_threads;
};

struct to_register_mpi_comm {
  int comm_id;
  int comm_size;
  uint64_t* members; 		/* mpi ranks */
};

enum to_register_type
  {
   type_string,
   type_function,
   type_thread,
   type_gpu,
   type_attribute,
   type_thread_team,
   type_thread_team_member,
   type_mpi_comm,
   type_invalid,
  };

struct to_register {
  enum to_register_type type;
  
  union {
    struct to_register_string string;
    struct to_register_function function;
    struct to_register_thread thread;
    struct to_register_gpu gpu;
    struct to_register_attribute attribute;
    struct to_register_thread_team thread_team;
    struct to_register_thread_team_member thread_team_member;
    struct to_register_mpi_comm mpi_comm;
  } data;
  struct to_register* next;
};

struct to_register* to_register=NULL;
struct to_register* to_register_tail=NULL;

static int _register_string(int string_id, const char* string) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);
  assert(string_id >= 0);

  pthread_mutex_lock(&otf2_lock);
  OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteString( _ezt_trace.global_def_writer,
							 string_id,
							 string);
  pthread_mutex_unlock(&otf2_lock);
  if (err != OTF2_SUCCESS) {
    eztrace_warn("OTF2_GlobalDefWriter_WriteString failed: %s (%s)\n",
		 OTF2_Error_GetName(err), OTF2_Error_GetDescription(err));

    eztrace_warn("string: '%s'\n",
		 string);
    
    eztrace_warn("To debug: gdb attach %d\n", getpid());
    while(1);
    return -1;
  }
  return 0;
}

static int _register_function(int id, int function_name) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);
  assert(id >= 0);
  assert(function_name >= 0);

  pthread_mutex_lock(&otf2_lock);
  OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteRegion
    (
     _ezt_trace.global_def_writer,
     id /* id */,
     function_name /* region name  */,
     function_name /* alternative name */,
     function_name /* description */,
     OTF2_REGION_ROLE_FUNCTION,
     OTF2_PARADIGM_USER,
     OTF2_REGION_FLAG_NONE,
     0 /* source file */,
     0 /* begin lno */,
     0 /* end lno */ );
  pthread_mutex_unlock(&otf2_lock);
  if (err != OTF2_SUCCESS) {
    return -1;
  }
  return 0;
}

static int _register_thread(int id, int thread_name, int mpi_rank) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);
  assert(id >= 0);
  assert(thread_name >= 0);

  pthread_mutex_lock(&otf2_lock);
  OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteLocation
    ( _ezt_trace.global_def_writer,
      id,
      thread_name,
      OTF2_LOCATION_TYPE_CPU_THREAD,
      2 /* # events */,
      mpi_rank);
  pthread_mutex_unlock(&otf2_lock);
  if (err != OTF2_SUCCESS) {
    return -1;
  }
  return 0;  
}

static int _register_gpu(int id, int gpu_name, int mpi_rank) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);
  assert(id >= 0);
  assert(gpu_name >= 0);

  pthread_mutex_lock(&otf2_lock);
  OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteLocation
    ( _ezt_trace.global_def_writer,
      id,
      gpu_name,
#if OTF2_MAJOR_VERSION >= 3
      OTF2_LOCATION_TYPE_ACCELERATOR_STREAM,
#else
      OTF2_LOCATION_TYPE_GPU,
#endif
      2 /* # events */,
      mpi_rank);
  pthread_mutex_unlock(&otf2_lock);
  if (err != OTF2_SUCCESS) {
    return -1;
  }
  return 0;  
}

static int _register_attribute(int id, int attribute_name, OTF2_Type type) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);
  assert(id >= 0);
  assert(attribute_name >= 0);
  pthread_mutex_lock(&otf2_lock);
  OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteAttribute(_ezt_trace.global_def_writer,
							   id,
							   attribute_name,
							   attribute_name,
							   type);
  pthread_mutex_unlock(&otf2_lock);
  if (err != OTF2_SUCCESS) {
    return -1;
  }
  return 0;
}


/* When registering an OpenMP thread team, we need to know which thread are part of the team
 * This is done by enqueueing thread team until all the theads where registered
 */
struct thread_team {
  int id;
  int string_id;
  int parent_id;
  int nb_threads;
  uint64_t *thread_ids;
  int nb_ready_threads; // nb of threads whose id are known
  struct thread_team *next;
};
static struct thread_team *thread_teams = NULL;
static ezt_spinlock thread_team_lock;

static int _register_thread_team(struct to_register*r) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);

  struct thread_team *team = malloc(sizeof(struct thread_team));
  team->id = r->data.thread_team.id;
  team->string_id = r->data.thread_team.string_id;
  team->parent_id = r->data.thread_team.parent_id;
  team->nb_threads = r->data.thread_team.nb_threads;
  team->thread_ids = malloc(sizeof(uint64_t)*team->nb_threads);
  team->nb_ready_threads = 0;

  /* enqueue the new team */
  ezt_spin_lock(&thread_team_lock);
  team->next = thread_teams;
  thread_teams = team;
  ezt_spin_unlock(&thread_team_lock);

  return 0;
}

static int _register_thread_team_member(struct to_register*r) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);

  ezt_spin_lock(&thread_team_lock);

  struct thread_team *team = thread_teams;
  struct thread_team *prev = NULL;
  /* search for the corresponding therad team */
  while(team) {
    if(team->id == r->data.thread_team_member.thread_team_id) {

      if(team->nb_threads == 0)
	team->nb_threads = r->data.thread_team_member.nb_threads;
      assert(team->nb_threads == r->data.thread_team_member.nb_threads);
      assert(team->nb_threads > r->data.thread_team_member.thread_rank);
      team->thread_ids[r->data.thread_team_member.thread_rank] = r->data.thread_team_member.thread_id;
      team->nb_ready_threads++;
      if(team->nb_ready_threads == team->nb_threads) {
	/* all the thread were registered. We can now record the OTF2 event */

	/* remove the team from the list of unterminated thread teams */
	if(prev)	prev->next = team->next;
	else thread_teams = team->next;

	ezt_spin_unlock(&thread_team_lock);

	pthread_mutex_lock(&otf2_lock);
	OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteGroup(_ezt_trace.global_def_writer,
							     team->id,
							     team->string_id,
							     OTF2_GROUP_TYPE_COMM_LOCATIONS,
							     OTF2_PARADIGM_OPENMP,
							     OTF2_GROUP_FLAG_NONE,
							     team->nb_threads,
							     team->thread_ids);
	if(err != OTF2_SUCCESS) {
	  eztrace_error( "OTF2_GlobalDefWriter_WriteGroup failed: %s\n",
			 OTF2_Error_GetDescription(err));
	}
	
	err =  OTF2_GlobalDefWriter_WriteComm(_ezt_trace.global_def_writer,
					      team->id,
					      team->string_id,
					      team->id,
					      team->parent_id
#if OTF2_MAJOR_VERSION >= 3
					      , OTF2_COMM_FLAG_NONE
#endif
					      );

	if(err != OTF2_SUCCESS) {
	  eztrace_error("OTF2_GlobalDefWriter_WriteComm failed: %s\n",
			OTF2_Error_GetDescription(err));
	}
	pthread_mutex_unlock(&otf2_lock);
	free(team->thread_ids);
	free(team);
	return 0;
      }

      ezt_spin_unlock(&thread_team_lock);
      return 0;
    }
    team = team->next;
  }
  ezt_spin_unlock(&thread_team_lock);

  /* not found */
  return -1;
}


static int _register_mpi_comm(int comm_id, int comm_size, uint64_t*members) {
  if((! _ezt_otf2_safe) && (! EZTRACE_SAFE))  return -1;
  assert(ezt_mpi_rank == 0);
  assert(comm_size >= 0);

  int retval = -1;
  char comm_name[1024];
  char comm_group_name[1024];
  snprintf(comm_group_name, 1024, "MPI_COMM_%d_group", comm_id);
  snprintf(comm_name, 1024, "MPI_COMM_%d", comm_id);
  int comm_str = ezt_otf2_register_string(comm_name);
  int comm_group_str = ezt_otf2_register_string(comm_group_name);

  pthread_mutex_lock(&otf2_lock);

  OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteGroup( _ezt_trace.global_def_writer,
							comm_id,
							comm_group_str,
							OTF2_GROUP_TYPE_COMM_GROUP,
							OTF2_PARADIGM_MPI,
							OTF2_GROUP_FLAG_NONE,
							comm_size,
							members );
  if (err != OTF2_SUCCESS) {
    eztrace_warn("OTF2_GlobalDefWriter_WriteGroup failed: %s (%s)\n", 
		 OTF2_Error_GetName(err), OTF2_Error_GetDescription(err));
    goto out;
  }

  OTF2_GlobalDefWriter_WriteComm( _ezt_trace.global_def_writer,
				  comm_id,
				  comm_str,
				  comm_id,
				  comm_world_ref
#if OTF2_MAJOR_VERSION >= 3
				  , OTF2_COMM_FLAG_NONE
#endif
				  );
  if (err != OTF2_SUCCESS) {
    eztrace_warn("OTF2_GlobalDefWriter_WriteComm failed: %s (%s)\n",
		 OTF2_Error_GetName(err), OTF2_Error_GetDescription(err));
    goto out;
  }

  retval = 0;
  pthread_mutex_unlock(&otf2_lock);
 out:
  return retval;
}

static int _register(struct to_register*r) {
  switch(r->type) {
  case type_string:
    return _register_string(r->data.string.id, r->data.string.buffer);
  case type_function:
    return _register_function(r->data.function.id, r->data.function.function_name);
  case type_thread:
    return _register_thread(r->data.thread.id, r->data.thread.string_id, r->data.thread.mpi_rank);
  case type_gpu:
    return _register_gpu(r->data.gpu.id, r->data.gpu.string_id, r->data.gpu.mpi_rank);
  case type_attribute:
    return _register_attribute(r->data.attribute.id, r->data.attribute.string_id, r->data.attribute.type);
  case type_thread_team:
    return _register_thread_team(r);
  case type_thread_team_member:
    return _register_thread_team_member(r);
  case type_mpi_comm:
    return _register_mpi_comm(r->data.mpi_comm.comm_id, r->data.mpi_comm.comm_size, r->data.mpi_comm.members);
  default:
    eztrace_warn("invalid register type: %d\n", r->type);
  }
  return -1;
}


static int _postpone_registration(struct to_register* r) {
  if(ezt_mpi_rank > 0  || _register(r) < 0) {
    r->next = NULL;

    /* Use a spinlock to make enqueuing thread-safe */
    static ezt_spinlock spinlock;
    ezt_spin_lock(&spinlock);

    if(to_register) {
      assert(to_register_tail);
      to_register_tail->next = r;
      to_register_tail = r;
    } else {
      to_register = r;
      to_register_tail = r;
    }
    
    ezt_spin_unlock(&spinlock);
  }

  return 0;
}

static void _register_postponed() {
  if(mpi_comm_size==1)
    /* serial program. Nothing to do */
    return;

  int type_tag = 17;
  int data_tag = 18;

  _ezt_otf2_safe = 1;

  if(ezt_mpi_rank == 0) {
    struct to_register* r = to_register;
    while(r) {
      struct to_register* next = r->next;
      _register(r); 
      free(r);
      r = next;
    }

    for(int i=1; i<mpi_comm_size; i++) {
      struct to_register r;
      
      EZT_MPI_Recv(&r.type, sizeof(r.type), i, type_tag);
      while(r.type != type_invalid) {
	switch(r.type) {
	case type_string: {
	  EZT_MPI_Recv(&r.data.string.id, 2*sizeof(int), i, data_tag);
	  r.data.string.buffer = malloc(r.data.string.len);
	  EZT_MPI_Recv(r.data.string.buffer, r.data.string.len, i, data_tag);
	  _register(&r);
	  free(r.data.string.buffer);
	  r.data.string.buffer = NULL;
	  break;
	}
	case type_function:
	  EZT_MPI_Recv(&r.data.function, sizeof(r.data.function), i, data_tag);
	  _register(&r);
	  break;
	case type_thread:
	  EZT_MPI_Recv(&r.data.thread, sizeof(r.data.thread), i, data_tag);
	  _register(&r);
	  break;

	case type_gpu:
	  EZT_MPI_Recv(&r.data.gpu, sizeof(r.data.gpu), i, data_tag);
	  _register(&r);
	  break;

	case type_attribute:
	  EZT_MPI_Recv(&r.data.attribute, sizeof(r.data.attribute), i, data_tag);
	  _register(&r);
	  break;

	case type_thread_team:
	  EZT_MPI_Recv(&r.data.thread_team, sizeof(r.data.thread_team), i, data_tag);
	  _register(&r);
	  break;
	  
	case type_thread_team_member:
	  EZT_MPI_Recv(&r.data.thread_team_member, sizeof(r.data.thread_team_member), i, data_tag);
	  _register(&r);
	  break;

	case type_mpi_comm:
	  EZT_MPI_Recv(&r.data.mpi_comm.comm_id, 2*sizeof(int), i, data_tag);
	  int buffer_size = r.data.mpi_comm.comm_size*sizeof(uint64_t);
	  r.data.mpi_comm.members = malloc(buffer_size);
	  EZT_MPI_Recv(r.data.mpi_comm.members, buffer_size, i, data_tag);
	  _register(&r);

	  break;
	default:
	  eztrace_error("Invalid register type: %d\n", r.type);
	}
	EZT_MPI_Recv(&r.type, sizeof(r.type), i, type_tag);
      }
    }
  
  } else {
    struct to_register* r = to_register;
    while(r) {
      
      EZT_MPI_Send(&r->type, sizeof(r->type), 0, type_tag);
      switch(r->type) {
      case type_string:
	EZT_MPI_Send(&r->data.string.id, 2*sizeof(int), 0, data_tag);
	EZT_MPI_Send(r->data.string.buffer, r->data.string.len, 0, data_tag);
	free(r->data.string.buffer);
	break;
      case type_function:
	EZT_MPI_Send(&r->data.function, sizeof(r->data.function), 0, data_tag);
	break;
      case type_thread:
	EZT_MPI_Send(&r->data.thread, sizeof(r->data.thread), 0, data_tag);
	break;
      case type_gpu:
	EZT_MPI_Send(&r->data.gpu, sizeof(r->data.gpu), 0, data_tag);
	break;
      case type_attribute:
	EZT_MPI_Send(&r->data.attribute, sizeof(r->data.attribute), 0, data_tag);
	break;

      case type_thread_team:
	EZT_MPI_Send(&r->data.thread_team, sizeof(r->data.thread_team), 0, data_tag);
	break;

      case type_thread_team_member:
	EZT_MPI_Send(&r->data.thread_team_member, sizeof(r->data.thread_team_member), 0, data_tag);
	break;
      case type_mpi_comm:
	EZT_MPI_Send(&r->data.mpi_comm.comm_id, sizeof(int)*2, 0, data_tag);
	int buffer_size = r->data.mpi_comm.comm_size*sizeof(uint64_t);
	EZT_MPI_Send(r->data.mpi_comm.members, buffer_size, 0, data_tag);
	break;

      default:
	eztrace_error("Invalid register type: %d\n", r->type);
      }
      struct to_register* next = r->next;
      free(r);
      r = next;
    }
    /* tell rank 0 that we're done */
    enum to_register_type t = type_invalid;
    EZT_MPI_Send(&t, sizeof(t), 0, type_tag);
  }
  _ezt_otf2_safe = 0;
}

int ezt_otf2_register_mpi_comm(int comm_size, uint64_t *members) {
  //  if(! eztrace_should_trace)
  //    return -1;

  /* TODO: record the parent communicator */

  int comm_id = next_mpi_comm_id++;
  struct to_register *r = malloc(sizeof(struct to_register));
  r->type = type_mpi_comm;
  r->data.mpi_comm.comm_id = comm_id;
  r->data.mpi_comm.comm_size = comm_size;
  size_t buffer_size = comm_size * sizeof(uint64_t);
  r->data.mpi_comm.members = malloc(buffer_size);
  for(int i=0; i< comm_size; i++) {
    r->data.mpi_comm.members[i] = members[i];
  }

  if(_postpone_registration(r) < 0)
    return -1;
  return comm_id;
}

int ezt_otf2_register_thread_team(char *name, int nb_threads) {
  if(! eztrace_should_trace)
    return -1;

  int thread_team_id = next_thread_team_id++;
  
  int string_id = ezt_otf2_register_string(name);
  int parent_id = ezt_mpi_rank;

  struct to_register *r = malloc(sizeof(struct to_register));
  r->type = type_thread_team;
  r->data.thread_team.id = thread_team_id;
  r->data.thread_team.string_id = string_id;
  r->data.thread_team.parent_id = parent_id;
  r->data.thread_team.nb_threads = nb_threads;

  if(_postpone_registration(r) < 0)
    return -1;
  return thread_team_id;
}

int ezt_otf2_register_thread_team_member(int team_id, int my_rank, int nb_threads) {
  if(! eztrace_should_trace)
    return -1;

  uint64_t thread_id = thread_rank;
  struct to_register *r = malloc(sizeof(struct to_register));
  r->type = type_thread_team_member;
  r->data.thread_team_member.thread_team_id = team_id;
  r->data.thread_team_member.thread_rank = my_rank;
  r->data.thread_team_member.thread_id = thread_id;
  r->data.thread_team_member.nb_threads = nb_threads;

  if(_postpone_registration(r) < 0)
    return -1;

  return 0;  
}

int ezt_otf2_register_string(const char* string) {
  if(! eztrace_should_trace)
    return -1;

  int string_id = next_string_id++;

  if(ezt_mpi_rank > 0) {
    struct to_register *r = malloc(sizeof(struct to_register));
    r->type = type_string;
    r->data.string.id = string_id;
    r->data.string.len = strlen(string)+1;
    r->data.string.buffer = malloc(sizeof(char)*r->data.string.len);
    assert(r->data.string.buffer);
    strcpy(r->data.string.buffer, string);
    if(_postpone_registration(r) < 0)
      return -1;
  } else {
    if( _register_string(string_id, string) < 0)
      return -1;
  }
  return string_id;
}


int ezt_otf2_register_attribute(char* name, OTF2_Type type) {
  if(! eztrace_should_trace)
    return -1;

  int attr_id = next_attribute_id++;
  int name_id = ezt_otf2_register_string(name);
  if(ezt_mpi_rank > 0) {
    struct to_register *r = malloc(sizeof(struct to_register));
    r->type = type_attribute;
    r->data.attribute.id = attr_id;
    r->data.attribute.string_id = name_id;
    r->data.attribute.type = type;
    if(_postpone_registration(r) < 0)
      return -1;
  } else {
    if(_register_attribute(attr_id, name_id, type) < 0)
      return -1;
  }
  return attr_id;
}

int ezt_otf2_register_function(const char* function_name) {
  if(! eztrace_should_trace)
    return -1;

  int region_id=next_region_id++;
  int function_name_id = ezt_otf2_register_string(function_name);
  if(ezt_mpi_rank > 0) {
    struct to_register *r = malloc(sizeof(struct to_register));
    r->type = type_function;
    r->data.function.id = region_id;
    r->data.function.function_name = function_name_id;
    if(_postpone_registration(r) < 0)
      return -1;
  } else {
    if(_register_function(region_id, function_name_id) < 0)
      return -1;
  }

  return region_id;
}

int ezt_otf2_register_thread(int local_thread_id) {
  int thread_id = next_thread_id++;
  //  int local_thread_id = next_local_thread_id++;
  char thread_name[128];
  snprintf(thread_name, 128, "P#%dT#%d", ezt_mpi_rank, local_thread_id);
  OTF2_StringRef name = ezt_otf2_register_string(thread_name);

  /* The first thread of each process is already registered */
  if(thread_id % otf2_chunk_size == 0) {
    return thread_id;
  }

  if(ezt_mpi_rank > 0) {
    struct to_register *r = malloc(sizeof(struct to_register));
    r->type = type_thread;
    r->data.thread.id = thread_id;
    r->data.thread.string_id = name;
    r->data.thread.mpi_rank = ezt_mpi_rank;
    if(_postpone_registration(r) < 0)
      return -1;
  } else {
    if(_register_thread(thread_id, name, ezt_mpi_rank) < 0)
      return -1;
  }
  
  _ezt_trace.def_writer = OTF2_Archive_GetDefWriter( _ezt_trace.archive,
						     thread_id );
  return thread_id;  
}

int ezt_otf2_register_gpu(int local_gpu_id) {
  int gpu_id = next_thread_id++;
  //  int local_thread_id = next_local_thread_id++;
  char gpu_name[128];
  snprintf(gpu_name, 128, "P#%dGPU#%d", ezt_mpi_rank, local_gpu_id);
  OTF2_StringRef name = ezt_otf2_register_string(gpu_name);

  if(ezt_mpi_rank > 0) {
    struct to_register *r = malloc(sizeof(struct to_register));
    r->type = type_gpu;
    r->data.gpu.id = gpu_id;
    r->data.gpu.string_id = name;
    r->data.gpu.mpi_rank = ezt_mpi_rank;
    if(_postpone_registration(r) < 0)
      return -1;
  } else {
    if(_register_gpu(gpu_id, name, ezt_mpi_rank) < 0)
      return -1;
  }
  
  return gpu_id;  
}

void ezt_otf2_set_mpi_rank(int rank, int comm_size) {
  static int already_set = 0;
  if(already_set)
    return;
  already_set = 1;
  ezt_mpi_rank = rank;
  mpi_comm_size = comm_size;

  otf2_chunk_size = INT_MAX/comm_size;
  max_id = otf2_chunk_size*(ezt_mpi_rank+1);

  next_string_id = otf2_chunk_size * ezt_mpi_rank;
  next_region_id = otf2_chunk_size * ezt_mpi_rank;
  next_thread_id = otf2_chunk_size * ezt_mpi_rank;
  next_attribute_id = otf2_chunk_size * ezt_mpi_rank;
  if(ezt_mpi_rank > 0) {
    next_mpi_comm_id = otf2_chunk_size * ezt_mpi_rank;
  }

  ezt_otf2_set_mpi_rank_called = 1;
}

static void _eztrace_set_buffer_size() {
  char* res = getenv("EZTRACE_BUFFER_SIZE");
  if (res) {
    _ezt_trace.buffer_size = atoi(res);
  }
}

extern char**environ;
char ld_preload_value[4096];

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

  char* ld_preload = getenv("LD_PRELOAD");
  if(!ld_preload) {
    ld_preload_value[0]='\0';
    return;
  }

  /* save the value of ld_preload so that we can set it back later */
  strncpy(ld_preload_value, ld_preload, 4096);
  int ret = unsetenv("LD_PRELOAD");
  if(ret != 0 ){
    eztrace_error( "unsetenv failed ! %s\n", strerror(errno));
  }

  /* also change the environ variable since exec* function
   * rely on it.
   */
  for (int i=0; environ[i]; i++) {
    if (strstr(environ[i],"LD_PRELOAD=")) {
      eztrace_log(dbg_lvl_verbose, "hacking out LD_PRELOAD from environ[%d]\n",i);
      environ[i][0] = '\0';
    }
  }
  char*plop=getenv("LD_PRELOAD");
  if(plop) {
    eztrace_warn("Warning: cannot unset LD_PRELOAD\nThis is likely to cause problems later.\n");
  }
}

/* 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() {
  if(strlen(ld_preload_value)>0) {
    setenv("LD_PRELOAD", ld_preload_value, 1);
  }
}

int ezt_otf2_initialize_thread(int thread_rank) {
  _ezt_otf2_safe = 1;
  int ret = ezt_otf2_register_thread(thread_rank);
  _ezt_otf2_safe = 0;
  return ret;
}

void ezt_otf2_init() {
  eztrace_log(dbg_lvl_debug, "eztrace_otf2_init starts\n");
  todo_set_status("ezt_otf2", init_started);

  if(ezt_is_mpi_mode()) {
    enum todo_status ezt_mpi_status = todo_get_status("ezt_mpi");
    if(ezt_mpi_status != init_complete) {
      eztrace_log(dbg_lvl_debug, "eztrace_otf2_init needs to wait until MPI is initialized\n");
      add_todo_dependency("ezt_otf2", "ezt_mpi", init_complete);
      /* ezt_mpi initialization is pending */
      todo_set_status("ezt_otf2", init_stopped);
      return;
    }

    using_mpi = 1;
  } else {
    /* not using mpi == mpi with 1 rank */
    ezt_otf2_set_mpi_rank(0, 1);
  }
  
  _ezt_trace.buffer_size = (16 * 1024 * 1024); // 160MB
  _eztrace_set_buffer_size();

  eztrace_set_filename();

  const char* trace_directory = _ezt_trace.filename;
  if(ezt_mpi_rank == 0) {
    /* Delete the trace directory so that otf2 does not print warnings */
    unset_ld_preload();

    int child_pid = fork();
    if(child_pid == 0) {
      _ezt_trace.status = ezt_trace_status_finalized;
      execlp("/bin/rm", "/bin/rm", "-rf", trace_directory, NULL);
      eztrace_error("execlp failed\n");
      exit(EXIT_FAILURE);
    } else {
      int status = 0;
      if(waitpid(child_pid, &status, 0) != child_pid)
	eztrace_warn("waitpid failed\n");
      if(! WIFEXITED(status)){
	eztrace_warn("Warning: cannot delete %s at startup (%d)\n", trace_directory, WEXITSTATUS(status));
	if(WIFSIGNALED(status)) {
	  eztrace_warn("\tterminated with signal %d\n", WTERMSIG(status));
	}
	eztrace_abort();
      }
    }
    reset_ld_preload();
  }
  if(using_mpi) {
    /* make sure rank 0 deleted the trace directory befor entering OTF2_Archive_Open */
    EZT_MPI_Barrier();
  }

  fflush(stdout);
  fflush(stderr);

  uint64_t size = _ezt_trace.buffer_size;
  _ezt_trace.archive = OTF2_Archive_Open( trace_directory,
					  "eztrace_log",
					  OTF2_FILEMODE_WRITE,
					  size /* event chunk size */,
					  size /* def chunk size */,
					  OTF2_SUBSTRATE_POSIX,
					  OTF2_COMPRESSION_NONE );
  if(!_ezt_trace.archive) {
    eztrace_error("OTF2_Archive_Open(%s) failed\n", trace_directory);
  }
  OTF2_Pthread_Archive_SetLockingCallbacks(_ezt_trace.archive, NULL);
  OTF2_Archive_SetFlushCallbacks( _ezt_trace.archive, &flush_callbacks, NULL );
  if(!using_mpi) {
    OTF2_Archive_SetSerialCollectiveCallbacks( _ezt_trace.archive );
  } else {
    EZT_MPI_SetMPICollectiveCallbacks( _ezt_trace.archive );
  }

  OTF2_Archive_OpenEvtFiles( _ezt_trace.archive );

  _ezt_otf2_safe = 1;
  if(ezt_mpi_rank == 0) {
    eztrace_should_trace=1;
    _ezt_trace.global_def_writer = OTF2_Archive_GetGlobalDefWriter( _ezt_trace.archive );

    char cmd[1024];
    _eztrace_get_current_program_name(cmd);
    int program_name = ezt_otf2_register_string(cmd);
    
    OTF2_GlobalDefWriter_WriteSystemTreeNode( _ezt_trace.global_def_writer,
					      0 /* id */,
					      program_name /* name */,
					      program_name /* class */,
					      OTF2_UNDEFINED_SYSTEM_TREE_NODE /* parent */ );

    uint64_t locations[mpi_comm_size];
    uint64_t members[mpi_comm_size];
    int local_thread_id = next_local_thread_id++;

    OTF2_Archive_OpenDefFiles( _ezt_trace.archive );
    for(int i=0; i<mpi_comm_size; i++) {
      char process_name[128];
      snprintf(process_name, 128, "P#%d", i);
      int process_string = ezt_otf2_register_string(process_name);
      assert(process_string>=0);

      OTF2_GlobalDefWriter_WriteLocationGroup(  _ezt_trace.global_def_writer,
						i /* id */,
						process_string /* name */,
						OTF2_LOCATION_GROUP_TYPE_PROCESS,
						0 /* system tree */
#if OTF2_MAJOR_VERSION >= 3
						, OTF2_UNDEFINED_LOCATION_GROUP
#endif
						);
      int thread_id = MPI2OTF(i);
      locations[i] = thread_id;
      members[i] = i;
      char thread_name[128];
      snprintf(thread_name, 128, "P#%dT#%d", i, local_thread_id);
      OTF2_StringRef name = ezt_otf2_register_string(thread_name);

      _register_thread(thread_id, name, i);

      OTF2_DefWriter* def_writer = OTF2_Archive_GetDefWriter( _ezt_trace.archive,
							      thread_id );
      OTF2_Archive_CloseDefWriter( _ezt_trace.archive, def_writer );

    }
    OTF2_Archive_CloseDefFiles( _ezt_trace.archive );

    int comm_world_str = ezt_otf2_register_string("MPI_COMM_WORLD");
    int comm_world_locations_str = ezt_otf2_register_string("All processes");
    int comm_world_group_str = ezt_otf2_register_string("Main process group");
    OTF2_GroupRef comm_world_locations = 0;
    OTF2_GroupRef comm_world_group = 1;
    OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteGroup(_ezt_trace.global_def_writer,
							 comm_world_locations,
							 comm_world_locations_str,
							 OTF2_GROUP_TYPE_COMM_LOCATIONS,
							 OTF2_PARADIGM_MPI,
							 OTF2_GROUP_FLAG_NONE,
							 mpi_comm_size,
							 locations);
    if (err != OTF2_SUCCESS) {
      eztrace_warn("OTF2_GlobalDefWriter_WriteGroup failed: %s (%s)\n",
		   OTF2_Error_GetName(err), OTF2_Error_GetDescription(err));
      return ;
    }

    err = OTF2_GlobalDefWriter_WriteGroup( _ezt_trace.global_def_writer,
					   comm_world_group,
					   comm_world_group_str,
					   OTF2_GROUP_TYPE_COMM_GROUP,
					   OTF2_PARADIGM_MPI,
					   OTF2_GROUP_FLAG_NONE,
					   mpi_comm_size,
					   members);
    if (err != OTF2_SUCCESS) {
      eztrace_warn("OTF2_GlobalDefWriter_WriteGroup failed: %s (%s)\n", 
		   OTF2_Error_GetName(err), OTF2_Error_GetDescription(err));
      return ;
    }

    OTF2_GlobalDefWriter_WriteComm( _ezt_trace.global_def_writer,
				    comm_world_ref,
				    comm_world_str,
				    comm_world_group,
				    OTF2_UNDEFINED_COMM
#if OTF2_MAJOR_VERSION >= 3
				    , OTF2_COMM_FLAG_NONE
#endif
);
    if (err != OTF2_SUCCESS) {
      eztrace_warn("OTF2_GlobalDefWriter_WriteComm failed: %s (%s)\n",
		  OTF2_Error_GetName(err), OTF2_Error_GetDescription(err));
      return ;
    }

  }

  _ezt_otf2_safe = 0;
  epoch_start = ezt_get_timestamp();
  todo_set_status("ezt_otf2", init_complete);
  eztrace_log(dbg_lvl_debug, "eztrace_otf2_init ends\n");
  eztrace_can_trace = 1;
  return;
}

static void _synchronize_epochs() {
  epoch_stop = ezt_get_timestamp();
  uint64_t global_epoch_start;
  uint64_t global_epoch_stop;

  if(mpi_comm_size > 1) {

    EZT_MPI_Reduce( &epoch_start,
		    &global_epoch_start,
		    1, EZT_MPI_UINT64_T, EZT_MPI_MIN,
		    0);
    EZT_MPI_Reduce( &epoch_stop,
		    &global_epoch_stop,
		    1, EZT_MPI_UINT64_T, EZT_MPI_MIN,
		    0);
  } else {
    global_epoch_start = epoch_start;
    global_epoch_stop = epoch_stop;
  }

  if(ezt_mpi_rank == 0) {
    OTF2_ErrorCode err = OTF2_GlobalDefWriter_WriteClockProperties( _ezt_trace.global_def_writer,
								    1e9 /* 1 tick per second */,
								    global_epoch_start,
								    global_epoch_stop - global_epoch_start + 1
#if OTF2_MAJOR_VERSION >= 3
								    , OTF2_UNDEFINED_TIMESTAMP
#endif
								    );
    if( err != OTF2_SUCCESS) {
      eztrace_error("OTF2_GlobalDefWriter_WriteClockProperties failed\n");
    }

  }
}


struct ezt_finalize_callback {
  ezt_finalize_callback_t callback;
  OTF2_EvtWriter* evt_writer;
  enum ezt_trace_status *thread_status;
  uint64_t thread_id;
};

struct ezt_finalize_callback *ezt_finalize_callbacks = NULL;
_Atomic int ezt_finalize_nb_callbacks = 0;
static ezt_spinlock at_finalize_lock;

/* 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) {

  ezt_spin_lock(&at_finalize_lock);

  int index = ezt_finalize_nb_callbacks++;
  ezt_finalize_callbacks = realloc(ezt_finalize_callbacks,
				   ezt_finalize_nb_callbacks*sizeof(struct ezt_finalize_callback));

  ezt_finalize_callbacks[index].callback = callback;
  ezt_finalize_callbacks[index].evt_writer = evt_writer;
  ezt_finalize_callbacks[index].thread_status = thread_status;
  ezt_finalize_callbacks[index].thread_id = thread_id;

  ezt_spin_unlock(&at_finalize_lock);
}

/* in some cases a thread may exit and free its TLS thread_status,
 * which will cause ezt_at_finalize access memory illegally.
 * To avoid this, we cancel at_finalize when the thread destruction is detected
 * 
 * This function assumes that at_finalize_lock is locked
 */
void ezt_at_finalize_cancel_locked(uint64_t thread_id) {
  for(int i=0; i<ezt_finalize_nb_callbacks; i++) {
    if(ezt_finalize_callbacks[i].thread_id == thread_id) {
      ezt_finalize_callbacks[i].callback = NULL;
      ezt_finalize_callbacks[i].thread_status = NULL;
    }
  }
}

void ezt_at_finalize_cancel(uint64_t thread_id) {
  ezt_spin_lock(&at_finalize_lock);
  ezt_at_finalize_cancel_locked(thread_id);
  ezt_spin_unlock(&at_finalize_lock);
}

static void _ezt_at_finalize_run_callback(int i) {
  if(ezt_finalize_callbacks[i].thread_status &&
     *ezt_finalize_callbacks[i].thread_status != ezt_trace_status_finalized) {
    ezt_finalize_callbacks[i].callback(ezt_finalize_callbacks[i].evt_writer,
				       ezt_finalize_callbacks[i].thread_status,
				       ezt_finalize_callbacks[i].thread_id);
    ezt_finalize_callbacks[i].callback = NULL;
    ezt_finalize_callbacks[i].thread_status = NULL;
  }
}

void ezt_at_finalize_run(uint64_t thread_id) {
  ezt_spin_lock(&at_finalize_lock);
  for(int i=0; i<ezt_finalize_nb_callbacks; i++) {
    if(ezt_finalize_callbacks[i].thread_id == thread_id) {
      _ezt_at_finalize_run_callback(i);
    }
  }
  ezt_spin_unlock(&at_finalize_lock);
}

/* execute all the at_finalize callbacks */
void ezt_at_finalize_execute() {
  ezt_spin_lock(&at_finalize_lock);
  for(int i=0; i<ezt_finalize_nb_callbacks; i++) {
    _ezt_at_finalize_run_callback(i);
  }
  ezt_spin_unlock(&at_finalize_lock);
}

void ezt_flush_thread_events(OTF2_EvtWriter* e_writer,
			     uint64_t thread_id) {
  static ezt_spinlock archive_lock;
    ezt_spin_lock(&archive_lock);

    OTF2_Archive_CloseEvtWriter( _ezt_trace.archive, e_writer );

    OTF2_Archive_OpenDefFiles( _ezt_trace.archive );
    OTF2_DefWriter* def_writer = OTF2_Archive_GetDefWriter( _ezt_trace.archive,
                                                            thread_id );

    OTF2_Archive_CloseDefWriter( _ezt_trace.archive, def_writer );
    ezt_spin_unlock(&archive_lock);
}

int ezt_otf2_finalize() {
  static _Atomic int finalized = 0;
  if(++finalized > 1)
    return -1;

  ezt_at_finalize_run(thread_rank);

  ezt_at_finalize_execute();

  ezt_spin_lock(&at_finalize_lock);
  for(int i=0; i<ezt_finalize_nb_callbacks; i++) {
    if(ezt_finalize_callbacks[i].evt_writer) {
      ezt_flush_thread_events(ezt_finalize_callbacks[i].evt_writer,
			      ezt_finalize_callbacks[i].thread_id);
      ezt_finalize_callbacks[i].evt_writer = NULL;
      ezt_finalize_callbacks[i].thread_id = 0;
    }
  }
  ezt_spin_unlock(&at_finalize_lock);


  if(eztrace_should_trace) {
    OTF2_Archive_CloseEvtFiles( _ezt_trace.archive );
  }

  if(ezt_mpi_rank == 0 ) {
    _ezt_trace.global_def_writer = OTF2_Archive_GetGlobalDefWriter( _ezt_trace.archive );
    if( ! _ezt_trace.global_def_writer) {
      eztrace_error("OTF2_Archive_GetGlobalDefWriter\n");
    }
  }

  _synchronize_epochs();
  _register_postponed();

  OTF2_Archive_CloseDefFiles( _ezt_trace.archive );

  if(mpi_comm_size == 1 || ezt_mpi_rank == 0) {
    OTF2_Archive_CloseGlobalDefWriter( _ezt_trace.archive,
				       _ezt_trace.global_def_writer );
  }


  OTF2_Archive_Close( _ezt_trace.archive );


  return 0;
}

void eztrace_otf2_constructor() __attribute__((constructor));
void eztrace_otf2_constructor() {
  eztrace_log(dbg_lvl_debug, "eztrace_otf2 constructor starts\n");
  if(eztrace_should_trace)
    enqueue_todo("ezt_otf2", ezt_otf2_init, NULL, status_invalid);
  eztrace_log(dbg_lvl_debug, "eztrace_otf2 constructor ends\n");
}
