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

#ifndef _REENTRANT
#define _REENTRANT
#endif

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

#include <pthread.h>
#include <semaphore.h>
#include <stdlib.h>
#include <unistd.h>
#include <sched.h>
#include <sys/syscall.h>
#include <stdatomic.h>

#define CURRENT_MODULE eztrace_core
DECLARE_CURRENT_MODULE;

static volatile int _pthread_core_initialized = 0;
extern int eztrace_should_trace;
int (*libpthread_create)(pthread_t* thread, const pthread_attr_t* attr,
                         void* (*start_routine)(void*), void* arg);
int (*libpthread_join)(pthread_t th, void** thread_return);
void (*libpthread_exit)(void* thread_return);

static int pthread_create_id=-1;
static int pthread_join_id=-1;
static int pthread_exit_id=-1;

/* this function shall be called by each new thread */
void eztrace_set_alarm();

/* Thread creation/destruction callbacks */

/* Internal structure used for transmitting the function and argument
 * during pthread_create.
 */
struct _pthread_create_info_t {
  void* (*func)(void*);
  void* arg;
};

static _Atomic int nb_threads=0;

extern int current_process_ref;

static void _ezt_finalize_thread(OTF2_EvtWriter* e_writer,
				 enum ezt_trace_status *th_status,
				 uint64_t thread_id,
				 int is_locked) {

  while(*th_status != ezt_trace_status_running) {
    /* the thread is doing something, let's wait until it is safe to stop it */
    if(*th_status == ezt_trace_status_finalized)
      return;
    sched_yield();
  }

  if(eztrace_should_trace && e_writer) {
    EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadEnd(e_writer, NULL, ezt_get_timestamp(),
					    OTF2_UNDEFINED_COMM, thread_id));

    *th_status = ezt_trace_status_finalized;
  }
  *th_status = ezt_trace_status_finalized;
  if(is_locked)
    ezt_at_finalize_cancel_locked(thread_id);
  else
    ezt_at_finalize_cancel(thread_id);
}

void ezt_finalize_thread_locked(OTF2_EvtWriter* e_writer,
				enum ezt_trace_status *status,
				uint64_t thread_id) {
  _ezt_finalize_thread(e_writer, status, thread_id, 1);
}

void ezt_finalize_thread() {
  _ezt_finalize_thread(evt_writer, &thread_status, otf2_thread_id, 0);
}

struct ezt_thread_info registered_threads[1024];

struct ezt_thread_info* get_thread_info_by_pid(pid_t tid) {
  for(int i=0; i<nb_threads; i++)
    if(registered_threads[i].tid == tid)
      return &registered_threads[i];

  return NULL;
}

struct ezt_thread_info* get_thread_info(int thread_id) {
  assert(thread_id < nb_threads);
  return &registered_threads[thread_id];
}


static pid_t _gettid() {
#if defined(__APPLE__) && MAC_OS_X_VERSION_MAX_ALLOWED >= MAC_OS_X_VERSION_10_12
  uint64_t tid64;
  pthread_threadid_np(NULL, &tid64);
  pid_t tid = (pid_t)tid64;
#else
  pid_t tid = syscall(__NR_gettid);
#endif
  return tid;
}

static void _ezt_register_thread() {
  struct ezt_thread_info *info = &registered_threads[thread_rank];
  info->thread_rank = thread_rank;
  info->tid = _gettid();
  snprintf(info->thread_name, 50, "P#%dT#%d", ezt_mpi_rank, info->thread_rank);
  
  int thread_id = ezt_otf2_initialize_thread(thread_rank);
  assert(thread_id !=  -1);
  otf2_thread_id = thread_id;
  info->otf2_thread_id = otf2_thread_id;

  evt_writer = OTF2_Archive_GetEvtWriter( _ezt_trace.archive,
					  otf2_thread_id);
}

void ezt_thread_cleanup() {
  static _Thread_local int finalized = 0;
  if(finalized)
    return;
  finalized = 1;
  ezt_finalize_thread();
  ezt_at_finalize_cancel(otf2_thread_id);
  return;
}

void ezt_init_thread() {  
  if((thread_status >= ezt_trace_status_running))
    return;
  if(todo_get_status("eztrace_init") != init_complete)
    return;

  /* atomically increment nb_threads */
  thread_rank = nb_threads++;

  _ezt_register_thread();

  thread_status = ezt_trace_status_running;
  ezt_at_finalize(ezt_finalize_thread_locked, evt_writer, &thread_status, otf2_thread_id);

  RECORD_HW_COUNTERS();
  eztrace_set_alarm();
  
  todo_set_status("ezt_init_thread", init_complete);

  ezt_pthread_first_event();
}

void ezt_pthread_first_event() {
  if(EZTRACE_SAFE) {
    EZT_OTF2_CHECK(OTF2_EvtWriter_ThreadBegin(evt_writer, NULL, ezt_get_timestamp(),
					      OTF2_UNDEFINED_COMM, thread_rank));
  }
}

int ezt_get_nb_threads() {
  return nb_threads;
}

/* Invoked by pthread_create on the new thread */
static void*
_pthread_new_thread(void* arg) {
  struct _pthread_create_info_t* p_arg = (struct _pthread_create_info_t*)arg;
  void* (*f)(void*) = p_arg->func;
  void* _arg = p_arg->arg;
  free(p_arg);

  ezt_init_thread();

  void* res = (*f)(_arg);

  ezt_at_finalize_run(otf2_thread_id);
  return res;
}

int pthread_create(pthread_t* __restrict thread,
                   const pthread_attr_t* __restrict attr,
                   void* (*start_routine)(void*),
                   void* __restrict arg) {
  FUNCTION_ENTRY;
  struct _pthread_create_info_t* _args =
      (struct _pthread_create_info_t*)malloc(
          sizeof(struct _pthread_create_info_t));
  _args->func = start_routine;
  _args->arg = arg;

  if (!libpthread_create)
    INTERCEPT("pthread_create", libpthread_create);

  /* We do not call directly start_routine since we could not get the actual creation timestamp of
   * the thread. Let's invoke _pthread_new_thread that will PROF_EVENT the thread and call
   * start_routine.
   */
  int retval = libpthread_create(thread, attr, _pthread_new_thread, _args);

  FUNCTION_EXIT;
  return retval;
}

int pthread_join(pthread_t th, void** thread_return) {
  FUNCTION_ENTRY;

  if (!libpthread_join)
    INTERCEPT("pthread_join", libpthread_join);

  int ret = libpthread_join(th, thread_return);
  FUNCTION_EXIT;
  return ret;
}

void pthread_exit(void* thread_return) {
  ezt_thread_cleanup();
  libpthread_exit(thread_return);
  __builtin_unreachable();
}

PPTRACE_START_INTERCEPT_FUNCTIONS(eztrace_core)
  INTERCEPT3("pthread_create", libpthread_create)
  INTERCEPT3("pthread_join", libpthread_join)
  INTERCEPT3("pthread_exit", libpthread_exit)
PPTRACE_END_INTERCEPT_FUNCTIONS(eztrace_core)


static void init_pthread_core() {
  
  INSTRUMENT_FUNCTIONS(eztrace_core);

  if(pthread_create_id<0)
    pthread_create_id=ezt_otf2_register_function("pthread_create");
  if(pthread_join_id<0)
    pthread_join_id=ezt_otf2_register_function("pthread_join");
  if(pthread_exit_id<0)
    pthread_exit_id=ezt_otf2_register_function("pthread_exit");
  

  if (eztrace_autostart_enabled())
    eztrace_start();

  _pthread_core_initialized = 1;
}

static void finalize_pthread_core() {
  if(_pthread_core_initialized) {
    _pthread_core_initialized = 0;
    eztrace_stop();
  }
}

static void _pthread_core_init(void) __attribute__((constructor));
static void _pthread_core_init(void) {
  EZT_REGISTER_MODULE(eztrace_core, "EZTrace core",
		      init_pthread_core, finalize_pthread_core);
}

static void _pthread_core_conclude(void) __attribute__((destructor));
static void _pthread_core_conclude(void) {
  if(_pthread_core_initialized) {
    _pthread_core_initialized = 0;
    eztrace_stop();
  }
}
