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

#include <unistd.h>
#include <sys/types.h>
#include <string.h>
#include <stdlib.h>
#include <stdarg.h>
#include <errno.h>
#include <assert.h>
#include <pthread.h>
#include <signal.h>
#include <execinfo.h>
#include <sys/syscall.h>

#include "ev_codes.h"
#include "eztrace.h"
#include "eztrace_config.h"

#if HAVE_LIBBACKTRACE
#include <libbacktrace/backtrace.h>
#include <libbacktrace/backtrace-supported.h>
#endif

static void __eztrace_run_atexit_functions();

static pthread_key_t protect_on;

struct __ezt_write_trace __ezt_trace;

 
static void init_recursion_shield() {
  static int init_done = 0;
  if (!init_done) {
    pthread_key_create(&protect_on, NULL);
    init_done = 1;
  }
}

#ifndef EZTRACE_AUTOSTART

void eztrace_register_init_routine(eztrace_function_t init_func) {
  __ezt_trace.init_routines[__ezt_trace.nb_module++] = init_func;
}

#endif	/* EZTRACE_AUTOSTART */

static void __eztrace_set_buffer_size() {
  char* res = getenv("EZTRACE_BUFFER_SIZE");
  if (res) {
    __ezt_trace.buffer_size = atoll(res);
    if(__ezt_trace.buffer_size > UINT32_MAX) {
      __ezt_trace.buffer_size = UINT32_MAX - 1024;
      fprintf(stderr, "[EZTrace] Buffer size is too large. Reducing it to the maximum: %lu bytes\n", __ezt_trace.buffer_size);
    }
    fprintf(stderr, "[EZTrace] Buffer size: %lu bytes\n", __ezt_trace.buffer_size);
  }
}

static char* __eztrace_get_filedir() {
  char* res = getenv("EZTRACE_TRACE_DIR");
  if (!res)
    asprintf(&res, ".");
  return res;
}

void eztrace_set_filename(char* name) {
  int total_size = strlen(__eztrace_get_filedir());
  total_size += strlen(getenv("USER"));
  total_size += strlen(name);
  total_size += 3;		/* add 3 chars ('/', '_', '\0') */

  if(__ezt_trace.filename)
    free(__ezt_trace.filename);
  __ezt_trace.filename = malloc(total_size);

  sprintf(__ezt_trace.filename, "%s/%s_%s", __eztrace_get_filedir(), getenv("USER"),
	  name);
  litl_write_set_filename(__ezt_trace.litl_trace, __ezt_trace.filename);
}

/* return the library name from a symbol string */
static char* get_lib_name(char*symbol) {
  char* ret;
  int begin = 0;
  int end = -1;
  int i = 0;

  /* the format of symbol is : '/path/to/libxxx.so (function+0xoffset) [0xaddress]*/
  /* so we need to locate the last / and the first ( */
  while (symbol[i] != 0) {
    if (symbol[i] == '/')
      begin = i + 1;
    if (symbol[i] == '(') {
      end = i;
      break;
    }
    i++;
  }

  ret = &symbol[begin];
  /* replace ( with a \0 */
  if (end >= 0)
    symbol[end] = 0;
  return ret;

}

/* return the function name (or library name if the function name is unknown)
 * from a symbol string.
 * This function may alter the symbol string.
 */
static char* get_function_name(char* symbol) {
  char* ret = symbol;
  int len = strlen(symbol);
  int begin=0;
  int end=0;
  int i;

  /* the format of symbol is : 'libxxx.so (function+0xoffset) [0xaddress]*/
  /* the goal is to retrieve the function+0xoffset string */
  for (i = 0; i < len; i++) {
    if (symbol[i] == '(') {
      begin = i;
      if (symbol[i + 1] == '+' || symbol[i + 1] == ')') {
	return get_lib_name(symbol);
      }
    }
    if (symbol[i] == ')') {
      end = i;
      break;
    }
  }
  if(begin == 0) {
    return symbol;
  }
  ret = &symbol[begin + 1];
  symbol[end] = 0;
  return ret;
}


void eztrace_error_handler(int signo) {
  static volatile int shield = 0;
  /* in case several thread receive signals simultaneously, only the first one
   * handle it.
   */
  while (shield)
    ;
  shield = 1;

  set_recursion_shield_on();

  EZT_PRINTF(0, "[EZTrace] signal %d catched. my pid=%d\n", signo, getpid());

  void* buffer[50];
  /* get pointers to functions */
  int nb_calls = backtrace(buffer, 50);
  char **functions;

  functions = backtrace_symbols(buffer, nb_calls);
  int i;

  EZTRACE_EVENT_PACKED_2(EZTRACE_SIGNAL_RECEIVED, signo, nb_calls);
  eztrace_record_backtrace(nb_calls);
  set_recursion_shield_off();
}

void eztrace_signal_handler(int signo) {
  static volatile int shield = 0;
  /* in case several thread receive signals simultaneously, only the first one
   * handle it.
   */
  while (shield)
    ;
  shield = 1;

  EZT_PRINTF(0, "EZTrace received signal %d...\n", signo);
  if (signo == SIGSEGV)
    eztrace_error_handler(signo);

  eztrace_stop();
  EZT_PRINTF(0, "Signal handling done\n");
  exit(EXIT_FAILURE);
  //signal(signo, SIG_DFL);
}

/* when an alarm signal is received, check for sampling information */
void __eztrace_alarm_sighandler(int signo __attribute__((unused))){
  ezt_sampling_check_callbacks();
}

#ifdef __linux__
long __ezt_alarm_interval = 0;
int alarm_enabled = 0;
int alarm_set = 0;

void eztrace_set_alarm() {
  if(__ezt_alarm_interval>=0 && alarm_enabled && (! alarm_set)) {
    alarm_set = 1;

    struct sigevent sevp;
    sevp.sigev_notify=SIGEV_THREAD_ID | SIGEV_SIGNAL;
    sevp.sigev_signo=SIGALRM;
    sevp.sigev_value.sival_int=0;
    sevp.sigev_notify_function = NULL;
    sevp.sigev_notify_attributes=NULL;
    sevp._sigev_un._tid = gettid();

    timer_t *t = malloc(sizeof(timer_t));
    int ret = timer_create(CLOCK_REALTIME, &sevp, t);
    if(ret != 0){
      perror("timer create failed");
      abort();
    }

    struct itimerspec new_value, old_value;
    new_value.it_interval.tv_sec=0;
    new_value.it_interval.tv_nsec=__ezt_alarm_interval;

    new_value.it_value.tv_sec=0;
    new_value.it_value.tv_nsec=__ezt_alarm_interval;

    ret = timer_settime(*t,0, &new_value, &old_value);
    if(ret != 0){
      perror("timer settime failed");
      abort();
    }
  }
}
#else
void eztrace_set_alarm() {
}
#endif

static void __eztrace_set_sighandler() {
  char* res = getenv("EZTRACE_NO_SIGNAL_HANDLER");

  if (!res || !strncmp(res, "0", 2)) {
    signal(SIGSEGV, eztrace_signal_handler);
    signal(SIGINT, eztrace_signal_handler);
    signal(SIGTERM, eztrace_signal_handler);
    signal(SIGABRT, eztrace_signal_handler);
    signal(SIGHUP, eztrace_signal_handler);
  }

#ifdef __linux__
  res = getenv("EZTRACE_SIGALARM");
  if(res && (strncmp(res, "0", 2) != 0)) {
    /* convert from ms to ns */
    alarm_enabled = 1;
    __ezt_alarm_interval = atoi(res)*1000000;
    printf("[EZTrace] Setting an alarm every %d ms\n", atoi(res));
    signal(SIGALRM, __eztrace_alarm_sighandler);
    eztrace_set_alarm();
  }
#endif
}

void eztrace_start_() {
    eztrace_start();
}

void eztrace_start() {
  /* avoid executing this function several times */
  if(__ezt_trace.status >= ezt_trace_status_running)
    return;

  __ezt_trace.status = ezt_trace_status_uninitialized;
  __ezt_trace.debug_level = 0;
  __ezt_trace.buffer_size = (16 * 1024 * 1024); // 16MB

  char* debug_mode = getenv("EZTRACE_DEBUG");
  if (debug_mode) {
      __ezt_trace.debug_level = atoi(debug_mode);
    EZT_PRINTF(0, "EZTrace Debug mode enabled (trace level: %d)\n",
	       __ezt_trace.debug_level);
  }

  EZT_PRINTF(0, "Starting EZTrace... ");
   __eztrace_set_buffer_size();

  /* make sure eztrace_stop is called when the program stops */
  atexit(eztrace_stop);
  __eztrace_set_sighandler();


  char* allow_flush = getenv("EZTRACE_FLUSH");
  // start LiTL
  __ezt_trace.litl_trace = litl_write_init_trace(__ezt_trace.buffer_size);
  // the recording should be paused, because some further functions, e.g. *_set_filename() can be intercepted by eztrace
  litl_write_pause_recording(__ezt_trace.litl_trace);
  litl_write_tid_recording_on(__ezt_trace.litl_trace);

  if (allow_flush && strncmp(allow_flush, "0", 2)) {
    litl_write_buffer_flush_on(__ezt_trace.litl_trace);
    EZT_PRINTF(0, "EZTrace Flush enabled\n");
  }

  eztrace_set_filename("eztrace_log_rank_1");

  litl_write_resume_recording(__ezt_trace.litl_trace);

  init_recursion_shield();

  __ezt_trace.status = ezt_trace_status_running;

  /* the current thread needs to be registered since eztrace won't
   * intercept any pthread_create for this one
   */
  EZTRACE_EVENT0(EZTRACE_NEW_THREAD);

#ifndef EZTRACE_AUTOSTART
  /* call the initialisation routines that were registered */
  int i;
  for (i = 0; i < __ezt_trace.nb_module; i++)
    __ezt_trace.init_routines[i]();
#endif	/* EZTRACE_AUTOSTART */
  EZT_PRINTF(0, "done\n");
}

void eztrace_stop_() {
    eztrace_stop();
}

void eztrace_stop() {
  if(__ezt_trace.status >= ezt_trace_status_being_finalized ||
     __ezt_trace.status < ezt_trace_status_running)
    return;
  __ezt_trace.status = ezt_trace_status_being_finalized;

  __eztrace_run_atexit_functions();

  EZTRACE_EVENT0(EZTRACE_END_THREAD);
  __ezt_trace.status = ezt_trace_status_finalized;

  EZTRACE_FIN_TRACE();

  EZT_PRINTF(0, "Stopping EZTrace... saving trace  %s\n", __ezt_trace.filename);
}

void eztrace_silent_start() {
  eztrace_start();
  __ezt_trace.status = ezt_trace_status_paused;
}

void eztrace_pause() {
  __ezt_trace.status = ezt_trace_status_paused;
}

void eztrace_resume() {
  __ezt_trace.status = ezt_trace_status_running;
}

void eztrace_enter_event(char *name, enum ezt_color color) {
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
  switch (color) {
   case EZTRACE_RED:
     litl_write_probe_raw(__ezt_trace.litl_trace,
			 EZTRACE_ENTER_EVENT_RED,
			 strlen(name) + 1,
			 (litl_data_t*) name);
    break;
  case EZTRACE_YELLOW:
    litl_write_probe_raw(__ezt_trace.litl_trace,
			 EZTRACE_ENTER_EVENT_YELLOW,
			 strlen(name) + 1,
			 (litl_data_t*) name);
    break;
      case EZTRACE_BLUE:
    litl_write_probe_raw(__ezt_trace.litl_trace,
			 EZTRACE_ENTER_EVENT_BLUE,
			 strlen(name) + 1,
			 (litl_data_t*) name);
    break;
      case EZTRACE_GREEN:
    litl_write_probe_raw(__ezt_trace.litl_trace,
			 EZTRACE_ENTER_EVENT_GREEN,
			 strlen(name) + 1,
			 (litl_data_t*) name);
    break;
      case EZTRACE_PINK:
    litl_write_probe_raw(__ezt_trace.litl_trace,
			 EZTRACE_ENTER_EVENT_PINK,
			 strlen(name) + 1,
			 (litl_data_t*) name);
    break;
  }
  EZTRACE_PROTECT_OFF();
  }
}

void eztrace_leave_event(char *name) {
  EZTRACE_EVENT0(EZTRACE_LEAVE_EVENT);
}


void eztrace_generic(uint32_t code, int nbargs, ...) {
  int i;
  va_list args;
  uint64_t arg_array[9];

  va_start(args, nbargs);
  for (i = 0; i < nbargs; i++)
    arg_array[i] = va_arg(args, uint64_t);

  switch (nbargs) {
  case 0:
    EZTRACE_EVENT0(code)
    ;
    break;
  case 1:
    EZTRACE_EVENT1(code, arg_array[0])
    ;
    break;
  case 2:
    EZTRACE_EVENT2(code, arg_array[0], arg_array[1])
    ;
    break;
  case 3:
    EZTRACE_EVENT3(code, arg_array[0], arg_array[1], arg_array[2])
    ;
    break;
  case 4:
    EZTRACE_EVENT4(code, arg_array[0], arg_array[1], arg_array[2], arg_array[3])
    ;
    break;
    case 5:
    EZTRACE_EVENT5(code, arg_array[0], arg_array[1], arg_array[2], arg_array[3], arg_array[4]);
    break;
    case 6:
    EZTRACE_EVENT6(code, arg_array[0], arg_array[1], arg_array[2], arg_array[3], arg_array[4], arg_array[5]);
    break;
    case 7:
    EZTRACE_EVENT7(code, arg_array[0], arg_array[1], arg_array[2], arg_array[3], arg_array[4], arg_array[5], arg_array[6]);
    break;
    case 8:
    EZTRACE_EVENT8(code, arg_array[0], arg_array[1], arg_array[2], arg_array[3], arg_array[4], arg_array[5], arg_array[6], arg_array[7]);
    break;
    case 9:
    EZTRACE_EVENT9(code, arg_array[0], arg_array[1], arg_array[2], arg_array[3], arg_array[4], arg_array[5], arg_array[6], arg_array[7], arg_array[8]);
    break;
  }
}

void eztrace_code0(uint32_t code) {
  EZTRACE_EVENT0(code);
}

void eztrace_code1(uint32_t code, uint64_t arg1) {
  EZTRACE_EVENT1(code, arg1);
}

void eztrace_code2(uint32_t code, uint64_t arg1, uint64_t arg2) {
  EZTRACE_EVENT2(code, arg1, arg2);
}

void eztrace_code3(uint32_t code, uint64_t arg1, uint64_t arg2, uint64_t arg3) {
  EZTRACE_EVENT3(code, arg1, arg2, arg3);
}

void eztrace_code4(uint32_t code, uint64_t arg1, uint64_t arg2, uint64_t arg3,
		   uint64_t arg4) {
  EZTRACE_EVENT4(code, arg1, arg2, arg3, arg4);
}

void eztrace_code5(uint32_t code, uint64_t arg1, uint64_t arg2, uint64_t arg3,
		   uint64_t arg4, uint64_t arg5) {
  EZTRACE_EVENT5(code, arg1, arg2, arg3, arg4, arg5);
}

void eztrace_code6(uint32_t code, uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4, uint64_t arg5, uint64_t arg6) {
  EZTRACE_EVENT6(code, arg1, arg2, arg3, arg4, arg5, arg6);
}

void eztrace_code7(uint32_t code, uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4, uint64_t arg5, uint64_t arg6, uint64_t arg7) {
  EZTRACE_EVENT7(code, arg1, arg2, arg3, arg4, arg5, arg6, arg7);
}

void eztrace_code8(uint32_t code, uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4, uint64_t arg5, uint64_t arg6, uint64_t arg7, uint64_t arg8) {
  EZTRACE_EVENT8(code, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8);
}

void eztrace_code9(uint32_t code, uint64_t arg1, uint64_t arg2, uint64_t arg3, uint64_t arg4, uint64_t arg5, uint64_t arg6, uint64_t arg7, uint64_t arg8, uint64_t arg9) {
  EZTRACE_EVENT9(code, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9);
}

int recursion_shield_on() {
  init_recursion_shield();
  void* ret = NULL;
  ret = pthread_getspecific(protect_on);
  return (ret != NULL);
}

void set_recursion_shield_on() {
  init_recursion_shield();
   pthread_setspecific(protect_on, (void*) 1);
}
void set_recursion_shield_off() {
  pthread_setspecific(protect_on, (void*) NULL);
}

static void  __eztrace_get_function_name_from_pointer_fallback(void* pointer,
							       char *output_str,
							       int buffer_len) {
  char **functions;
  functions = backtrace_symbols(pointer, 1);
  char* ptr = get_function_name(functions[0]);
  strncpy(output_str, ptr, strlen(ptr)+1);
  free(functions);
}

#if HAVE_LIBBACKTRACE

__thread struct backtrace_state *backtrace_state = NULL;
__thread char current_frame[1024];

static void error_callback(void *data, const char *msg, int errnum) {
  fprintf(stderr, "ERROR: %s (%d)", msg, errnum);
}

static int backtrace_callback (void *data, uintptr_t pc,
			       const char *filename, int lineno,
			       const char *function) {
  if(filename == NULL) {
    /* cannot find debugging info */
    __eztrace_get_function_name_from_pointer_fallback((void*)pc, current_frame, 1024);
  } else {
    snprintf(current_frame, 1024, "at %s:%d %s", filename, lineno, function);
  }
  return 0;
}
#endif /* HAVE_LIBBACKTRACE */


/* retrieve the name of the function corresponding to pointer
 * the collected information is copied into output_str
 * buffer_len is the maximum number of bytes written to output_str
 */
static void __eztrace_get_function_name_from_pointer(void* pointer,
						     char *output_str,
						     int buffer_len) {
#if HAVE_LIBBACKTRACE
  if(!backtrace_state) {
    backtrace_state = backtrace_create_state (NULL, 0, error_callback, NULL);
  }
  backtrace_pcinfo (backtrace_state, (uintptr_t)pointer,
		    backtrace_callback,
		    error_callback,
		    NULL);

  strncpy(output_str, current_frame, strlen(current_frame)+1);
#else
  __eztrace_get_function_name_from_pointer_fallback(pointer, output_str, buffer_len);
#endif

}

/* 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) {
  void* buffer[frameid+1];

  /* get pointers to functions */
  int nb_calls = backtrace(buffer, frameid+1);
  __eztrace_get_function_name_from_pointer(buffer[frameid],
					   output_str,
					   buffer_len);
}

#define RECORD_BACKTRACES 1
#ifdef RECORD_BACKTRACES

/* record events (code=EZTRACE_CALLING_FUNCTION) containing the backtrace */
void eztrace_record_backtrace(int backtrace_depth) {
  /* the current backtrace looks like this:
   * 0 - record_backtrace()
   * 1 - eztrace_callback()
   * 2 - calling_function
   *
   * So, we need to get the name of the function in frame 2.
   */
  void* bt_buffer[backtrace_depth];

  /* get pointers to functions */
  int nb_calls = backtrace(bt_buffer, backtrace_depth);

  EZTRACE_EVENT1_PACKED_UNPROTECTED(EZTRACE_BACKTRACE_DEPTH, nb_calls);
  int i;
  for (i = 0; i < nb_calls; i++) {
    char buffer[1024];
    __eztrace_get_function_name_from_pointer(bt_buffer[i], buffer, 1024);
    litl_write_probe_raw(__ezt_trace.litl_trace, EZTRACE_CALLING_FUNCTION, strlen(buffer), buffer);
  }
}

void record_backtrace() {
  EZTRACE_PROTECT {
    EZTRACE_PROTECT_ON();
    eztrace_record_backtrace(15);
    EZTRACE_PROTECT_OFF();
  }
}

#else
void eztrace_record_backtrace(int backtrace_depth) {
}
void record_backtrace() {
}
#endif


struct eztrace_atexit_token_t {
  eztrace_atexit_function_t func;
  void* param;
};

struct eztrace_atexit_list_t {
  struct eztrace_atexit_token_t* list;
  int nb_allocated;
  int nb_functions;
};

struct eztrace_atexit_list_t* atexit_list = NULL;

/* register a function to be called before eztrace_stop. */
void eztrace_atexit(eztrace_atexit_function_t f, void* param) {
  if(!atexit_list) {
    /* first time this function is called. Allocate/initialize the atexit_list structure */
    atexit_list = malloc(sizeof(struct eztrace_atexit_list_t));
    atexit_list->nb_allocated = 10;
    atexit_list->list = malloc(sizeof(struct eztrace_atexit_token_t) * atexit_list->nb_allocated);
    atexit_list->nb_functions = 0;
  }
  int num = atexit_list->nb_functions++;

  if(num >= atexit_list->nb_allocated) {
    /* too many atexit callbacks. We need to expand the array */
    atexit_list->nb_allocated *= 2;
    void* ptr = realloc(atexit_list->list, sizeof(struct eztrace_atexit_token_t)*atexit_list->nb_allocated);
    assert(ptr);
    atexit_list->list = ptr;
  }

  atexit_list->list[num].func = f;
  atexit_list->list[num].param = param;
}

/* run all the atexit callbacks that were registered */
static void __eztrace_run_atexit_functions() {
  if(atexit_list) {
    int i;
    for(i=0; i< atexit_list->nb_functions; i++) {
      atexit_list->list[i].func(atexit_list->list[i].param);
    }
  }
}
