#include "eztrace-lib/eztrace_module.h"
#include "eztrace-lib/eztrace_internals.h"

#include "eztrace-core/eztrace_config.h"
#include "eztrace-core/eztrace_list.h"
#include "eztrace-lib/eztrace.h"

#include <string.h>
#include <stdio.h>
#include <dlfcn.h>
#include <dirent.h>


static struct ezt_list_t module_list;
char* cur_module_name = NULL;
#define SO_STRING DYNLIB_EXT

int module_verbose = 0;
static int module_initialized = 0;

void _core_init(void) __attribute__((constructor));
void _core_init(void) {
  if (!module_initialized) {
    module_initialized = 1;
    ezt_list_new(&module_list);
  }
}

/* This function calls the init handler for all the registered modules */
void _init_modules() {
  static int recursion_shield = 0;
  if(recursion_shield)
    return;
  recursion_shield = 1;
  struct ezt_list_token_t* token;

  ezt_list_foreach(&module_list, token) {

    struct eztrace_module* p_module =
        (struct eztrace_module*)token->data;
    p_module->init();
  }

  todo_set_status("ezt_init_modules", init_complete);

  recursion_shield = 0;
  return;
}

/* This function calls the init handler for all the registered modules */
void finalize_modules() {
  static int recursion_shield = 0;
  if(recursion_shield)
    return;
  recursion_shield = 1;
  struct ezt_list_token_t* token;

  ezt_list_foreach(&module_list, token) {

    struct eztrace_module* p_module =
        (struct eztrace_module*)token->data;

    p_module->finalize();
  }

  recursion_shield = 0;
  return;
}

int is_registered(const char* module_name) {
  struct ezt_list_token_t* token;
  ezt_list_foreach(&module_list, token) {
    struct eztrace_module* p_module = (struct eztrace_module*)token->data;
    if(strcmp(p_module->name, module_name)== 0)
      return 1;
  }
  return 0;
}

int initialize_modules() {
  static int already_initialized = 0;
  if(already_initialized)
    return 0;
  char*str = getenv("EZTRACE_TRACE");
  if(str) {
    char* buffer = malloc((strlen(str)+1)*sizeof(char));
    memcpy(buffer, str, (strlen(str)+1)*sizeof(char));

    if(! is_registered("eztrace_core")) {
      eztrace_warn("module %s is not registered yet !\n", "lib");
      return 0;
    }

    /* EZTRACE_TRACE is declared.
     * it should contain modules separated by " ", such
     * as "mpi memory pthread"
     * Let's iterate over these modules and check if they are registered
     */
    char* save_ptr = buffer;
    char* module = strtok_r(buffer, " ", &save_ptr);
    while (module) {
      if(! is_registered(module)) {
	return 0;
      }

      module = strtok_r(NULL, " ", &save_ptr);
    }

    _init_modules();
    already_initialized = 1;
    todo_wait("eztrace", init_complete);
    return 1;
  }
  return 0;
}

void eztrace_print_module_list() {
  struct ezt_list_token_t* token;
  /* Let's check wether another module with the same prefix is already registered */
  ezt_list_foreach(&module_list, token) {

    struct eztrace_module* p_mod = ezt_container_of(token, struct eztrace_module, token);
    if(strcmp(p_mod->name, "eztrace_core") != 0)
      printf("%s\t%s\n", p_mod->name, p_mod->description);
  }
}

void eztrace_register_module(struct eztrace_module* p_module) {
  struct ezt_list_token_t* token;

  if(module_verbose)
    eztrace_log(dbg_lvl_normal, "Register module %s\n", p_module->name);
  
  /* Let's check wether another module with the same prefix is already registered */
  ezt_list_foreach(&module_list, token) {
    struct eztrace_module* p_mod =
        (struct eztrace_module*)token->data;

    if (strcmp(p_module->name, p_mod->name) == 0) {
      eztrace_warn("Trying to register a module that is already registered. Module name is %s\n",
		   p_mod->name);
      return;
    }
  }

  p_module->token.data = p_module;
  ezt_list_add(&module_list, &p_module->token);
  if (module_verbose)
    eztrace_log(dbg_lvl_normal, "module %s registered\n", p_module->name);

  initialize_modules();
}

/* return 1 if the filename matches "libeztrace-*.so" */
static int filter(const struct dirent* entry) {
  const char* filename = entry->d_name;

  /* check wether the string starts with "libeztrace-" */
  if (strncmp(filename, "libeztrace-", strlen("libeztrace-")))
    /* the string doesn't start with "libeztrace-" */
    return 0;

  /* check wether the next chars correspond to the module name
   * that we are looking for
   */
  filename += strlen("libeztrace-");
  if (cur_module_name) {
    if (strncmp(filename, cur_module_name, strlen(cur_module_name)))
      /* the string doesn't start with "libeztrace-" */
      return 0;
    filename += strlen(cur_module_name);
  } else {
    filename = entry->d_name + strlen(entry->d_name) - strlen(SO_STRING);
  }
  if (strncmp(filename, SO_STRING, strlen(SO_STRING)))
    /* the string doesn't end with ".so" */
    return 0;

  /* check wether there are remaining chars after .so */
  filename += strlen(SO_STRING);
  if (filename[0])
    return 0;
  return 1;
}

struct module{
  char fullpath[4096];
  char module_name[128];  
};

struct module modules[1024];
int nb_modules=0;



static int _load_module(struct module* m) {
  if(!m)
    return 0;

  if(module_verbose)
    eztrace_log(dbg_lvl_normal, "Trying to load module %s\n", m->module_name);

  void *handle = dlopen(m->fullpath, RTLD_LAZY);
  if(!handle) {
    eztrace_error("failed to open %s: %s\n", m->fullpath, dlerror());
    return 0;
  }

  /* search for the list of functions to instrument in the library we just opened */
  char symbol_alias [256];
  PPTRACE_SYMBOL_ALIAS_STRING(symbol_alias, m->module_name, 256);
  void* hijack_address = dlsym(handle, symbol_alias);
  if(! hijack_address) {
    dlclose(handle);
    if(module_verbose)
        eztrace_log(dbg_lvl_normal, "\tfailed\n");
    return 0;
  }
  if(module_verbose)
    eztrace_log(dbg_lvl_normal, "\tsuccess\n");
  return 1;
}

/* search for files that could be modules in dirname */
static void _list_modules_in_dir(const char* dirname) {
  struct dirent** namelist;
  int n;
  /* Get the list of files that match the module name in the libdir directory */
  n = scandir(dirname, &namelist, filter, alphasort);
  if (n < 0)
    perror("scandir");
  else {
    while (n--) {
      /* Get the full name of the file (path/libname.so) */
      struct module* m=&modules[nb_modules++];
      snprintf(m->fullpath, 4096, "%s/%s", dirname, namelist[n]->d_name);

      // file name is   libeztrace-SOMETHING.so
      // let's extract SOMETHING
      int prefix_len=strlen("libeztrace-");
      int suffix_len=strlen(SO_STRING);
      int len = strlen(namelist[n]->d_name);
      strncpy(m->module_name, &namelist[n]->d_name[prefix_len], len-(suffix_len+prefix_len));
      m->module_name[len-(suffix_len+prefix_len+1)]='\0';
      free(namelist[n]);
    }
    free(namelist);
  }
}

/* search for files that could be modules in all the directories */
static void _list_modules() {
  /* search in the default directory*/
  _list_modules_in_dir(EZTRACE_LIB_DIR);

  char* lib_path = getenv("EZTRACE_LIBRARY_PATH");
  if (!lib_path) {
    /* No lib_path specified, we can't find any more plugin */
    goto out;
  }

  /* Iterate over the lib_path specified.
   * lib_path are separated by ':'
   */
  char* save_ptr = lib_path;
  char*cur_path = strtok_r(lib_path, ":", &save_ptr);
  while (cur_path) {
    _list_modules_in_dir(cur_path);
    cur_path = strtok_r(NULL, ":", &save_ptr);
  }

 out:
  /* we have a list of files that could contain modules, let check them */
  if(module_verbose) {
    eztrace_log(dbg_lvl_normal, "%d potential modules:\n", nb_modules);
    for(int i=0; i<nb_modules; i++) {
      eztrace_log(dbg_lvl_normal, "\t%s\n", modules[i].module_name);
    }
  }
}

int eztrace_load_module(const char* mod_name) {
  for(int i=0; i<nb_modules; i++) {
    if(strcmp(modules[i].module_name, mod_name)==0) {
      int ret = _load_module(&modules[i]);
      if(ret)
	return 1;
    }
  }
  return 0;

}

void eztrace_load_all_modules(int mod_verb) {
  module_verbose = mod_verb;

  _list_modules();
  int nb_loaded = 0;
  for(int i=0; i<nb_modules; i++) {
    nb_loaded += eztrace_load_module(modules[i].module_name);
  }

  if(module_verbose)
    eztrace_log(dbg_lvl_normal, "%d modules loaded\n", nb_loaded);
}

/* load all the modules specified by the EZTRACE_TRACE variable */
void eztrace_load_modules(int mod_verb) {
  todo_wait("eztrace_init", init_complete);

  char* module_list = getenv("EZTRACE_TRACE");
  module_verbose = mod_verb;
  int nb_loaded = 0;
  char* save_ptr = NULL;
  char* module = NULL;

  _list_modules();
  if (!module_list) {
    /* no env declares, so we only load the "core" module */
    nb_loaded = eztrace_load_module("lib");

    /* We have loaded the pthread_core module which is not a
     * 'real' module, so let's decrement nb_loaded
     */
    nb_loaded--;
    goto out;
  }

  /* EZTRACE_TRACE is declared.
   * it should contain modules separated by " ", such
   * as "mpi coreblas pthread"
   * Let's iterate over these modules and load them once at a time
   */
  save_ptr = module_list;
  module = strtok_r(module_list, " ", &save_ptr);
  while (module) {
    int loaded = eztrace_load_module(module);
    if (!loaded)
      eztrace_error("Cannot find module '%s'\n", module);
    nb_loaded += loaded;
    module = strtok_r(NULL, " ", &save_ptr);
  }

 out:
  if (mod_verb)
    eztrace_log(dbg_lvl_normal, "%d modules loaded\n", nb_loaded);
  return;
}
