// Backtrace handling part of i77 perms.

#define _POSIX_C_SOURCE 200809L
#define _GNU_SOURCE  /* Required to enable dladdr */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <setjmp.h>
#include <unistd.h>
#include <string.h>
#include <backtrace.h>
#include <dlfcn.h>   /* Required for dladdr and Dl_info */

#include "pbacktrace.h"

void *__bt_state = NULL;
static int bt_started = 0;

int __attribute__((no_instrument_function))  bt_callback(void *v, uintptr_t u, const char *filename, int lineno, const char *function) {
  const char *func_name = function;
  int status;

  if ((filename || lineno || func_name) && (strcmp(func_name, "print_back_trace"))) {
    if (!bt_started) {
      fprintf(stderr, "\nMONITOR ENTERED FROM IMP\n\n"); bt_started = 1;
    }
    fprintf(stderr, "Entered from %s at %s:%d\n\n", func_name, filename, lineno);
  }
  return 0;
}

void __attribute__((no_instrument_function))  bt_error_callback(void *v, const char *msg, int errnum) {
  if (strcmp(msg, "no debug info in ELF executable")) {
    fprintf(stderr, "Error %d occurred when getting the stacktrace: %s", errnum, msg);
  }
}

void __attribute__((no_instrument_function))  bt_error_callback_create(void *v, const char *msg, int errnum) {
  fprintf(stderr, "Error %d occurred when initializing the stacktrace: %s", errnum, msg);
}

void __attribute__((no_instrument_function))  init_back_trace(const char *filename) {
  // too late to init when a backtrace is needed - has to be done at startup.
  // Fortunately I believe there actually is a way with gcc to have some init
  // code run when main() is started.  TO DO.  This will also work for runtime tracing.
  __bt_state = backtrace_create_state(filename, 0, bt_error_callback_create, NULL);
}

void __attribute__((no_instrument_function))  print_back_trace(void) {
  if (!__bt_state) { /// make sure init_back_trace() is called
    fprintf(stderr, "Make sure init_back_trace() is called before calling print_stack_trace()\n");
    abort();
  }
  bt_started = 0;
  backtrace_full((struct backtrace_state *) __bt_state, 0, bt_callback, bt_error_callback, NULL);
  bt_started = 0;
}

void __cyg_profile_func_enter(void *this_fn, void *call_site) __attribute__((no_instrument_function));
void __cyg_profile_func_exit(void *this_fn, void *call_site) __attribute__((no_instrument_function));

static int tracing = 1;
static int trace_depth = 0;

// Global backtrace state - trace code is duplicating what was done for backtrace - need to merge...
static struct backtrace_state *bt_state = NULL;
const char *called_fn_name = "unknown";
const char *caller_name = "unknown";


// Callback called when libbacktrace successfully resolves the PC address
int __attribute__((no_instrument_function))  get_called_fn_name(void *v, uintptr_t u, const char *filename, int lineno, const char *function) {
  if (function) {
    //printf("Entered function: %s [at %s:%d]\n", function, filename, lineno);
    if (strncmp(function, "_imp_", 5) == 0) called_fn_name = function+5; else called_fn_name = function;
  } else {
    //printf("Entered unknown function at address %p\n", (void*)pc);
    called_fn_name = "<unknown function>";
  }
}

int __attribute__((no_instrument_function))  get_caller_name(void *v, uintptr_t u, const char *filename, int lineno, const char *function) {
  if (function) {
    //printf("Entered function: %s [at %s:%d]\n", function, filename, lineno);
    if (strncmp(function, "_imp_", 5) == 0) caller_name = function+5; else caller_name = function;
  } else {
    //printf("Entered unknown function at address %p\n", (void*)pc);
    caller_name  = "<unknown caller>";
  }
}

// Error callback required by libbacktrace
static void __attribute__((no_instrument_function))  error_callback(void *data, const char *msg, int errnum) {
  // Avoid printf here if called in a sensitive context
  called_fn_name = "<unknown function>";
  caller_name  = "<unknown caller>";
}

void __attribute__((no_instrument_function))  __cyg_profile_func_enter(void *this_fn, void *call_site) {
  if (tracing) {
#ifdef USE_DWARF
    Dl_info current_fn_info;
    Dl_info caller_fn_info;
    
    /* Look up the current function name */
    if (dladdr(this_fn, &current_fn_info) && current_fn_info.dli_sname) {
        called_fn_name = current_fn_info.dli_sname;
    }

    /* Look up the calling function name */
    if (dladdr(call_site, &caller_fn_info) && caller_fn_info.dli_sname) {
        caller_name = caller_fn_info.dli_sname;
    }
#else // ELF
    // Initialize the backtrace state once
    if (bt_state == NULL) {
      bt_state = backtrace_create_state(NULL, 0, error_callback, NULL);
      fprintf(stderr, "? To redirect tracing to a file, run with 2> tracefile.log\n");
    }
    
    if (bt_state) {
      // Query the exact program counter of the function block
      backtrace_pcinfo(bt_state, (uintptr_t)this_fn, get_called_fn_name, error_callback, NULL);
    }
#endif
    for (int i = trace_depth; i > 0; i -= 1) fprintf(stderr, " ");
    fprintf(stderr, ">> %s", called_fn_name);
    //fprintf(stderr, " (Called from: %s)", caller_name);
    fprintf(stderr, "\n");
    trace_depth += 2;
  }
}

void __attribute__((no_instrument_function))  __cyg_profile_func_exit(void *this_fn, void *call_site) {
  if (tracing) {
    Dl_info current_fn_info;

#ifdef USE_DWARF
    if (dladdr(this_fn, &current_fn_info) && current_fn_info.dli_sname) {
        called_fn_name = current_fn_info.dli_sname;
    }
#else // ELF
    // Initialize the backtrace state once
    if (bt_state == NULL) {
      bt_state = backtrace_create_state(NULL, 0, error_callback, NULL);
      fprintf(stderr, "? To redirect tracing to a file, run with 2> tracefile.log\n");
    }
    
    if (bt_state) {
      // Query the exact program counter of the function block
      // called:
      backtrace_pcinfo(bt_state, (uintptr_t)this_fn, get_called_fn_name, error_callback, NULL);
      // caller: (appears to be unreliable)
      //backtrace_pcinfo(bt_state, (uintptr_t)call_site, get_caller_name, error_callback, NULL);
    }
#endif
    trace_depth -= 2;
    for (int i = trace_depth; i > 0; i -= 1) fprintf(stderr, " ");
    fprintf(stderr, "<< %s\n", called_fn_name);
  }
}
