/* -*- c-file-style: "GNU" -*- */
/*
 */

#ifndef _REENTRANT
#define _REENTRANT
#endif

#include <Python.h>
#include <errno.h>
#include <eztrace_python.h>
#include <eztrace-core/eztrace_config.h>
#include <eztrace-instrumentation/pptrace.h>
#include <eztrace-lib/eztrace.h>
#include <eztrace-lib/eztrace_module.h>
#include <frameobject.h>
#include <otf2/OTF2_AttributeList.h>
#include <regex.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <time.h>
#include <unistd.h>

static volatile int _python_initialized = 0;

#define CURRENT_MODULE python
DECLARE_CURRENT_MODULE;

#define NAME_SIZE 100
#define LINE_SIZE 1000
#define HASH_TABLE_SIZE 10000

struct filters {
  int function_count;
  char** functions_filter;
  int module_count;
  char** modules_filter;
  int source_filter;
  int caller_filter;
  char* source_module_name;
  char* source_filename;
};

struct timers {
  double search_function_timer;
  double register_filters_timer;
  double free_list_timer;
  double eztrace_python_enter_timer;
  double eztrace_python_leave_timer;
  double tracer_timer;
  double filter_timer;
};

struct time_calc {
  struct timespec start;
  struct timespec end;
};

struct Function {
  int id;
  char name[NAME_SIZE];
  struct Function* next;
};

// Variables
struct Function* function_hash_table[HASH_TABLE_SIZE] = {NULL};
struct filters my_filters = {0};
struct timers my_timers = {0, 0, 0, 0, 0};
struct time_calc timer_calc_0;
struct time_calc timer_calc_1;
int timer_enabled = 0;
int nb_functions = 0;

// Start timer
void start_timer(struct time_calc* timer) {
  if (timer_enabled) {
    clock_gettime(CLOCK_MONOTONIC, &(timer->start));
  }
}

// End timer
void end_timer(struct time_calc* timer, double* acc_timer) {
  if (timer_enabled) {
    clock_gettime(CLOCK_MONOTONIC, &(timer->end));
    long seconds = timer->end.tv_sec - timer->start.tv_sec;
    long nanoseconds = timer->end.tv_nsec - timer->start.tv_nsec;
    *(acc_timer) += seconds + nanoseconds * 1e-9;
  }
}

// Hash function : djb2
unsigned long hash_function(char* name, size_t table_size) {
  unsigned long hash = 5381;
  int c;
  while ((c = *name++)) {
    hash = ((hash << 5) + hash) + c; /* hash * 33 + c */
  }
  return hash % table_size;
}

// Adds a function to the hash table
void add_function(char* name) {
  struct Function* new_function = (struct Function*)malloc(sizeof(struct Function));
  if (new_function == NULL) {
    perror("Failed to allocate memory for a new function");
    exit(EXIT_FAILURE);
  }
  strncpy(new_function->name, name, NAME_SIZE - 1);
  new_function->name[NAME_SIZE - 1] = '\0';
  new_function->id = ezt_otf2_register_function(name);
  // Insert into the hash table
  unsigned int index = hash_function(name, HASH_TABLE_SIZE);
  new_function->next = function_hash_table[index];
  function_hash_table[index] = new_function;
  nb_functions++;
}

// Search for a function by name in the hash table
int search_function(char* name) {
  start_timer(&timer_calc_0);
  // Calculate the hash value
  unsigned int index = hash_function(name, HASH_TABLE_SIZE);
  struct Function* current = function_hash_table[index];
  while (current != NULL) {
    if (strcmp(current->name, name) == 0) {
      // Found a match
      end_timer(&timer_calc_0, &(my_timers.search_function_timer));
      return current->id;
    }
    current = current->next;
  }
  // Name not found in the hash table, register
  add_function(name);
  end_timer(&timer_calc_0, &(my_timers.search_function_timer));
  return function_hash_table[index]->id;
}

void free_hash_table() {
  start_timer(&timer_calc_0);
  for (int i = 0; i < HASH_TABLE_SIZE; ++i) {
    struct Function* current = function_hash_table[i];
    while (current != NULL) {
      struct Function* next = current->next;
      free(current);
      current = next;
    }
    function_hash_table[i] = NULL;
  }
  nb_functions = 0;
  end_timer(&timer_calc_0, &(my_timers.free_list_timer));
}

// Python cleanup wrapper for free_hash_table
static PyObject*
cleanup(PyObject* self, PyObject* args) {
  free_hash_table();
  Py_RETURN_NONE;
}

char** read_file_content(char* filename, int* word_count) {
  FILE* file = fopen(filename, "r");
  if (file == NULL) {
    eztrace_warn("Cannot open filter file %s: %s\n", filename, strerror(errno));
    return NULL;
  }

  int array_size = 10;
  char** words = (char**)malloc(array_size * sizeof(char*));
  if (words == NULL) {
    perror("Failed to allocate memory for words");
    exit(EXIT_FAILURE);
  }
  *word_count = 0;
  char line[LINE_SIZE];

  while (fgets(line, sizeof(line), file) != NULL) {
    char* token = strtok(line, " \t\n");

    while (token != NULL) {
      if (*word_count >= array_size) {
        array_size *= 2;
        char** temp = (char**)realloc(words, array_size * sizeof(char*));
        if (temp == NULL) {
          perror("Failed to reallocate memory for words");
          for (int i = 0; i < *word_count; i++) {
            free(words[i]);
          }
          free(words);
          exit(EXIT_FAILURE);
        }
        words = temp;
      }
      words[*word_count] = strndup(token, NAME_SIZE);
      (*word_count)++;
      token = strtok(NULL, " \t\n");
    }
  }
  fclose(file);
  return words;
}

char** get_filter_from_env(char* env_name, int* word_count) {
  char* env_value = getenv(env_name);
  if (env_value) {
    // If the env variable is a filename
    const char* pattern = "\\b\\S*\\.txt\\b";
    regex_t regex;
    if (regcomp(&regex, pattern, REG_EXTENDED) != 0) {
      perror("Failed to compile regex pattern\n");
      exit(EXIT_FAILURE);
    }
    regmatch_t match;
    if (regexec(&regex, env_value, 1, &match, 0) == 0) {
      char** content = read_file_content(env_value, word_count);
      return content;
    } else {
      // If the env variable is a string
      int array_size = 10;
      char** content = (char**)malloc(array_size * sizeof(char*));
      if (content == NULL) {
        perror("Failed to allocate mempry for words");
        exit(EXIT_FAILURE);
      }
      *word_count = 0;
      char* token = strtok(env_value, " ");
      while (token != NULL) {
        if (*word_count >= array_size) {
          array_size *= 2;
          char** temp = realloc(content, array_size * sizeof(char*));
          if (temp == NULL) {
            perror("Failed to reallocate memory for words");
            for (int i = 0; i < *word_count; i++) {
              free(content[i]);
            }
            free(content);
            exit(EXIT_FAILURE);
          }
          content = temp;
        }
        content[*word_count] = strndup(token, NAME_SIZE);
        token = strtok(NULL, " ");
        word_count++;
      }
      return content;
    }
  } else {
    return NULL;
  }
}

void get_filters(char* source_module_name, char* source_filename) {
  start_timer(&timer_calc_0);

  int function_list_size = -1;
  char** function_list = get_filter_from_env("EZTPY_FUNCTION_FILTER", &function_list_size);
  int module_list_size = -1;
  char** module_list = get_filter_from_env("EZTPY_MODULE_FILTER", &module_list_size);

  char* source_filter_str = getenv("EZTPY_SOURCE_FILTER");
  if (source_filter_str) {
    my_filters.source_filter = atoi(source_filter_str);
  }

  char* caller_filter_str = getenv("EZTPY_CALLER_FILTER");
  if (caller_filter_str) {
    my_filters.caller_filter = atoi(caller_filter_str);
  }

  my_filters.function_count = function_list_size;
  my_filters.functions_filter = function_list;
  my_filters.module_count = module_list_size;
  my_filters.modules_filter = module_list;
  my_filters.source_module_name = source_module_name;
  my_filters.source_filename = source_filename;
  end_timer(&timer_calc_0, &(my_timers.register_filters_timer));
}

// Python Wrapper for get_filters
static PyObject* register_filters(PyObject* self, PyObject* args) {
  char* source_module_name;
  char* source_filename;
  if (!PyArg_ParseTuple(args, "ss", &source_module_name, &source_filename)) {
    return NULL;
  }
  get_filters(source_module_name, source_filename);
  Py_RETURN_NONE;
}

void unregister_filters() {
  for (int i = 0; i < my_filters.function_count; i++) {
    free(my_filters.functions_filter[i]);
  }
  for (int i = 0; i < my_filters.module_count; i++) {
    free(my_filters.modules_filter[i]);
  }
  free(my_filters.functions_filter);
  free(my_filters.modules_filter);
}

void print_filters_c() {
  printf("%s\n", "Function Filter :");
  if (my_filters.function_count == -1) {
    printf("%s", "No function filter set up.");
  } else {
    for (int i = 0; i < my_filters.function_count; i++) {
      puts(my_filters.functions_filter[i]);
    }
  }
  printf("%s\n", "Module Filter :");
  if (my_filters.module_count == -1) {
    printf("%s\n", "No module filter set up.");
  } else {
    for (int i = 0; i < my_filters.module_count; i++) {
      puts(my_filters.modules_filter[i]);
    }
  }
  if (my_filters.source_filter == -1) {
    printf("%s\n", "No source filter set up.");
  } else {
    printf("Source filter = %d", my_filters.source_filter);
  }
  if (my_filters.caller_filter == -1) {
    printf("%s\n", "No caller filter set up.");
  } else {
    printf("Caller filter = %d", my_filters.caller_filter);
  }
  printf("Tested python script module name : %s\n", my_filters.source_module_name);
}

// Python wrapper for printing filters
static PyObject* print_filters(PyObject* self, PyObject* args) {
  print_filters_c();
  Py_RETURN_NONE;
}

void print_timers_c() {
  printf("tracer: %f seconds\n", my_timers.tracer_timer);
  printf("filtering: %f seconds\n", my_timers.filter_timer);
  printf("eztrace_python_enter: %f seconds\n", my_timers.eztrace_python_enter_timer);
  printf("eztrace_python_leave: %f seconds\n", my_timers.eztrace_python_leave_timer);
  printf("search_function: %f seconds\n", my_timers.search_function_timer);
  printf("register_filters: %f seconds\n", my_timers.register_filters_timer);
  printf("free_list: %f seconds\n", my_timers.free_list_timer);
}

static PyObject* print_timers(PyObject* self, PyObject* args) {
  print_timers_c();
  Py_RETURN_NONE;
}

int filter_function(char* func_name) {
  if (my_filters.function_count == -1) {
    return 0;
  }
  for (int i = 0; i < my_filters.function_count; i++) {
    if (strcmp(func_name, my_filters.functions_filter[i]) == 0) {
      return 0;
    }
  }
  return 1;
}

int filter_module(const char* module_name) {
  if (module_name == NULL) {
    return 0;
  }
  if (my_filters.module_count == -1) {
    return 0;
  }
  for (int i = 0; i < my_filters.module_count; i++) {
    if (strcmp(module_name, my_filters.modules_filter[i]) == 0) {
      return 0;
    }
  }
  return 1;
}

void eztrace_python_enter(char* function_name) {
  if (EZTRACE_SAFE) {
    int function_id = search_function(function_name);
    start_timer(&timer_calc_0);
    EZT_OTF2_EvtWriter_Enter(evt_writer,
			     NULL,
			     ezt_get_timestamp(),
			     function_id);
    end_timer(&timer_calc_0, &(my_timers.eztrace_python_enter_timer));
  }
}

void eztrace_python_leave(char* function_name) {
  if (EZTRACE_SAFE) {
    int function_id = search_function(function_name);
    start_timer(&timer_calc_0);
    EZT_OTF2_EvtWriter_Leave(evt_writer,
			     NULL,
			     ezt_get_timestamp(),
			     function_id);
    end_timer(&timer_calc_0, &(my_timers.eztrace_python_leave_timer));
  }
}

/* test if a function should be traced
 * return 1 if the function should be trace, or 0 otherwise
 */
int filtering_in_c(PyFrameObject* frame, char** function_name) {
  start_timer(&timer_calc_0);
  const char* module_name = NULL;

// Extract module name from frame
#if PYTHON_VERSION_INT < 311
  PyObject* module_name_obj = PyDict_GetItemString(frame->f_globals, "__name__");
#elif PYTHON_VERSION_INT >= 311
  PyObject* frame_globals = PyFrame_GetGlobals(frame);
  PyObject* module_name_obj = PyDict_GetItemString(frame_globals, "__name__");
#endif
  if (module_name_obj != NULL && PyUnicode_Check(module_name_obj)) {
    module_name = PyUnicode_AsUTF8(module_name_obj);
  }

  if(!frame) return 0;

// Extract function name from frame
#if PYTHON_VERSION_INT < 311
  if(! frame->f_code) return 0;
  const char* func_name = PyUnicode_AsUTF8(frame->f_code->co_name);
#elif PYTHON_VERSION_INT >= 311
  PyCodeObject* code_object = PyFrame_GetCode(frame);
  if(!code_object) return 0;
  const char* func_name = PyUnicode_AsUTF8(code_object->co_name);
#endif
  if (func_name == NULL) {
    PyErr_SetString(PyExc_RuntimeError, "Failed to get valid function name from frame->f_code->co_name");
    PyErr_Print();
    return -1;
  }
  *function_name = strdup(func_name);
  if (*function_name == NULL) {
    PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for function_name");
    PyErr_Print();
    return -1;
  }

  // importlib filter
  if (module_name != NULL && strncmp(module_name, "importlib", strlen("importlib")) == 0) {
    end_timer(&timer_calc_0, &(my_timers.filter_timer));
    return 1;
  }

  // source filter
  if (my_filters.source_filter == 1 && module_name != NULL && strncmp(module_name, my_filters.source_module_name, strlen(my_filters.source_module_name)) != 0) {
    end_timer(&timer_calc_0, &(my_timers.filter_timer));
    return 1;
  }

  // caller filter
  if (my_filters.caller_filter == 1) {
#if PYTHON_VERSION_INT < 311
    PyFrameObject* back_frame = frame->f_back;
#elif PYTHON_VERSION_INT >= 311
    PyFrameObject* back_frame = PyFrame_GetBack(frame);
#endif
    if (back_frame != NULL) {
#if PYTHON_VERSION_INT < 311
      const char* caller_filename_const = PyUnicode_AsUTF8(back_frame->f_code->co_filename);
#elif PYTHON_VERSION_INT >= 311
      PyCodeObject* back_frame_code = PyFrame_GetCode(back_frame);
      const char* caller_filename_const = PyUnicode_AsUTF8(back_frame_code->co_filename);
#endif
      if (caller_filename_const != NULL) {
        char* caller_filename = strdup(caller_filename_const);
        if (caller_filename == NULL) {
          PyErr_SetString(PyExc_MemoryError, "Failed to allocate memory for function_name");
          PyErr_Print();
          return -1;
        }
        char* last_slash = strrchr(caller_filename, '/');
        if (last_slash != NULL) {
          caller_filename = last_slash + 1;
          if (strncmp(caller_filename, my_filters.source_filename, sizeof(my_filters.source_filename)) != 0) {
            end_timer(&timer_calc_0, &(my_timers.filter_timer));
            return 1;
          }
        }
      }
    }
  }

  // function filter
  if (filter_function(*function_name) == 1) {
    end_timer(&timer_calc_0, &(my_timers.filter_timer));
    return 1;
  }

  // module filter
  if (filter_module(module_name) == 1) {
    end_timer(&timer_calc_0, &(my_timers.filter_timer));
    return 1;
  }

  // Not filtered
  end_timer(&timer_calc_0, &(my_timers.filter_timer));
  return 0;
}

// Tracer function
int tracer(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg) {
  start_timer(&timer_calc_1);
  const char* event_str;
  char* func_name = NULL;
  if (what == PyTrace_CALL) {
    event_str = "CALL";
    if (filtering_in_c(frame, &func_name) == 0) {
      eztrace_python_enter(func_name);
    }
    end_timer(&timer_calc_1, &(my_timers.tracer_timer));
    return 0; // Continue tracing
  } else if (what == PyTrace_RETURN) {
    event_str = "RETURN";
    if (filtering_in_c(frame, &func_name) == 0) {
      eztrace_python_leave(func_name);
    }
    end_timer(&timer_calc_1, &(my_timers.tracer_timer));
    return 0; // Continue tracing
  } else {
    end_timer(&timer_calc_1, &(my_timers.tracer_timer));
    return 0; // Continues tracing
  }
}

// Python wrapper function for setting up the tracer
static PyObject* set_tracer(PyObject* self, PyObject* args) {
  // Set the tracer function
  PyEval_SetTrace(tracer, NULL);
  Py_RETURN_NONE; // Return None to indicate success
}

// Python wrapper function for unsetting the tracer
static PyObject* unset_tracer(PyObject* self, PyObject* args) {
  // Set the tracer function
  PyEval_SetTrace(NULL, NULL);
  Py_RETURN_NONE; // Return None to indicate success
}

static PyObject* print_python_version(PyObject* self, PyObject* args) {
#if defined(PYTHON_VERSION_INT)
  printf("Python version = %d \n", PYTHON_VERSION_INT);
#endif
  Py_RETURN_NONE;
}

PPTRACE_START_INTERCEPT_FUNCTIONS(python)
PPTRACE_END_INTERCEPT_FUNCTIONS(python)

static void init_python() {
  INSTRUMENT_FUNCTIONS(python);

  if (eztrace_autostart_enabled())
    eztrace_start();

  _python_initialized = 1;
}

static void finalize_python() {
  _python_initialized = 0;

  eztrace_stop();
}

static void _python_init(void) __attribute__((constructor));
static void _python_init(void) {
  char* timer_str = getenv("EZTRACEPYTHON_TESTING");
  if (timer_str) {
    if (strcmp(timer_str, "FILTERING_PYTHON_TIMERS") == 0 || strcmp(timer_str, "FILTERING_C_TIMERS") == 0 || strcmp(timer_str, "TRACING_FULL_C_TIMERS") == 0) {
      timer_enabled = 1;
    }
  }
  eztrace_log(dbg_lvl_debug, "eztrace_python constructor starts\n");
  EZT_REGISTER_MODULE(python, "Module for python functions",
                      init_python, finalize_python);
  eztrace_log(dbg_lvl_debug, "eztrace_pythonconstructor ends\n");
}

// Method definition for the module
static PyMethodDef module_methods[] = {
    {"set_tracer", set_tracer, METH_NOARGS, "Set up the tracer function."},
    {"unset_tracer", unset_tracer, METH_NOARGS, "Unsets the tracer function (sets it to None)"},
    {"cleanup", cleanup, METH_NOARGS, "Cleanup function : frees up internal data"},
    {"register_filters", register_filters, METH_VARARGS, "Populate filter data from environment variables set by user"},
    {"print_filters", print_filters, METH_NOARGS, "Print filters registered by the module"},
    {"print_timers", print_timers, METH_NOARGS, "Print C function timers"},
    {"print_python_version", print_python_version, METH_NOARGS, "Prints detected python version"},
    {NULL, NULL, 0, NULL} // Sentinel
};

// Module definition structure
static struct PyModuleDef eztracepython = {PyModuleDef_HEAD_INIT,
                                           "eztracepython", // Module name
                                           NULL,            // Documentation string
                                           -1,              // Size of per-interpreter state or -1 if the module keeps state in global variables.
                                           module_methods};

// Module initialization function
PyMODINIT_FUNC
PyInit_eztracepython(void) {
#ifndef PYTHON_VERSION_INT
#error "PYTHON_VERSION_INT constant is not defined! Please define it."
#endif
  return PyModule_Create(&eztracepython);
};
