#include "hiload/hiload.h"

#include "logger.h"
#include "memory.h"
#include "symbols.h"
#include "types.h"

#include <assert.h>
#include <dlfcn.h>
#include <link.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
  struct sc_array_str names;    // Array of library names
  struct sc_array_ptr handles;  // Array of library handles
  struct sc_array_syms symbols; // Symbol info for modules
  size_t count;                 // Number of modules
} ModuleInfos;

typedef struct {
  struct sc_array_memreg memory_regions;
} HiloadContext;

static HiloadContext context = {0};
static ModuleInfos *module_infos = 0;

// if the pathname contains these, skip them from early gathering
// as they are unlikely to be changed, and would bloat our memory
const char *path_filter_list[] = {"libstdc++.", "libc++.",   "libc.", "libm.",
                                  "libgcc",     "ld-linux-", NULL};

static inline int if_load_symbols_for_module(struct dl_phdr_info *info) {
  const char *name = info->dlpi_name;

  for (int i = 0; path_filter_list[i] != NULL; i++) {
    if (strstr(name, path_filter_list[i]) != NULL) {
      return 0;
    }
  }
  return 1;
}

// Callback function for dl_iterate_phdr
static
int gather_module_infos_callback(struct dl_phdr_info *info, size_t size,
                                        void *data) {

  ModuleInfos *infos = (ModuleInfos *)data;

  // Store the module name
  sc_array_add(&infos->names, strdup(info->dlpi_name));

  sc_log_info("Processing: '%s'\n", info->dlpi_name);

  // Try to get the handle
  void *handle = dlopen(info->dlpi_name, RTLD_LAZY | RTLD_NOLOAD);
  assert(handle);
  infos->count++;

  sc_array_add(&infos->handles, handle);
  sc_log_debug("  size: %u\n", size);
  sc_log_debug("  handle: %p\n", sc_array_last(&infos->handles));
  sc_log_debug("  dlpi_addr: %p\n", info->dlpi_addr);
  sc_log_debug("  dlpi_tls_modid: %zu\n", info->dlpi_tls_modid);
  sc_log_debug("  dlpi_tls_data: %p\n", info->dlpi_tls_data);

  SymbolInfos symbol_info = {0};
  sc_array_add(&infos->symbols, symbol_info);

  return 0; // Continue iteration
}

static
void free_module_infos(ModuleInfos *modules) {
  if (!modules)
    return;

  for (size_t i = 0; i < modules->count; i++) {
    // Free char* before clearing the array
    const char *n = 0;
    sc_array_foreach(&modules->names, n) { free((void *)n); }
    sc_array_term(&modules->names);

    // Use a destructor for the symbolinfos
    hi_free_symbol_info(&(sc_array_at(&modules->symbols, i)));
  }

  sc_array_term(&modules->names);
  sc_array_term(&modules->handles);
  sc_array_term(&modules->symbols);
  free(modules);
}

static
ModuleInfos *gather_module_infos(void) {
  ModuleInfos *result = NULL;

  // Allocate the result structure
  result = calloc(1, sizeof(ModuleInfos));
  if (!result) {
    return NULL;
  }

  // Iterate over all loaded shared objects
  if (dl_iterate_phdr(gather_module_infos_callback, result) != 0) {
    // Error occurred in callback
    free_module_infos(result);
    return NULL;
  }

  return result;
}

static
ReloadResult reload_module(ModuleInfos *modules, const char *filename,
                                  void **updated_handle) {
  if (!modules || !filename) {
    return HI_RELOAD_NOT_FOUND;
  }

  // Find the module by filename
  int found = 0;
  size_t index = 0;

  for (size_t i = 0; i < modules->count; i++) {
    // Check if this is the module we're looking for
    const char *pathname = sc_array_at(&modules->names, i);

    if (pathname && strcmp(pathname, filename) == 0) {
      found = 1;
      index = i;
      break;
    }

    // Also check if the filename is at the end of the path
    if (pathname) {
      const char *basename = strrchr(pathname, '/');
      if (basename && strcmp(basename + 1, filename) == 0) {
        found = 1;
        index = i;
        break;
      }
    }
  }

  if (!found) {
    return HI_RELOAD_NOT_FOUND;
  }

  // Save the full path
  const char *fullpath = sc_array_at(&modules->names, index);
  if (!fullpath) {
    return HI_RELOAD_OPEN_ERROR;
  }

  // Close the old handle
  void *handle = sc_array_at(&modules->handles, index);
  if (handle) {
    if (dlclose(handle) != 0) {
      return HI_RELOAD_CLOSE_ERROR;
    }
  }

  // Open the module again with RTLD_NOW
  void *new_handle = dlopen(fullpath, RTLD_NOW);
  if (!new_handle) {
    sc_log_error("Error reloading module: %s\n", dlerror());
    return HI_RELOAD_OPEN_ERROR;
  }

  // Update the handle in our structure
  modules->handles.elems[index] = new_handle;

  // If the caller wants the new handle, provide it
  if (updated_handle) {
    *updated_handle = new_handle;
  }

  return HI_RELOAD_SUCCESS;
}

ReloadResult hi_reload_module(const char *module_name) {
  assert(module_infos);

  void *new_handle = NULL;
  ReloadResult result = reload_module(module_infos, module_name, &new_handle);

  return result;
}

void hi_print_module_infos() {
  assert(module_infos);

  const ModuleInfos *modules = module_infos;
  if (!modules) {
    sc_log_error("No module information available.\n");
    return;
  }

  sc_log_info("Found %zu loaded modules:\n\n", modules->count);

  for (size_t i = 0; i < modules->count; i++) {
    // Get handle information where possible
    Dl_info info = {0};
    int has_info = 0;

    sc_log_debug("'%s': %p\n", sc_array_at(&modules->names, i),
                 sc_array_at(&modules->handles, i));

    const SymbolInfos *symbols = &sc_array_at(&modules->symbols, i);
    for (int j = 0; j < sc_array_size(&symbols->names); j++) {
      const void *addr = sc_array_at(&symbols->addresses, j);
      const char *name = sc_array_at(&symbols->names, j);
      sc_log_debug("  %p: %s\n", addr, name);
    }

    sc_log_debug("\n");
  }
}

int hi_init() {
  assert(!module_infos);

  if (sc_log_init() != 0) {
    fprintf(stderr, "Failed to init logger.\n");
    return 1;
  }
  sc_log_set_level("DEBUG");

  if (read_memory_maps_self(&context.memory_regions) != HILOAD_OK) {
    sc_log_error("Could not populate program memory maps.\n");
    return HILOAD_FAIL;
  }

  ModuleInfos *infos = gather_module_infos();
  if (!infos) {
    sc_log_error("Failed to gather module infos.\n");
    return 1;
  }

  // Override existing
  if (module_infos) {
    free_module_infos(module_infos);
  }

  module_infos = infos;
  hi_print_module_infos();

  return 0;
}

void hi_deinit() {
  free_module_infos(module_infos);
  sc_log_term();
}