// Required for dlinfo and other GNU specific linker code
#define _GNU_SOURCE

#include "hiload/hiload.h"
#include "hiload/symbols.h"

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

typedef struct {
  char **names;         // Array of library names
  void **handles;       // Array of library handles
  SymbolInfos *symbols; // Symbol info for modules
  size_t count;         // Number of libraries
  size_t capacity;      // Allocated capacity
} ModuleInfos;

static ModuleInfos *module_infos = 0;

// 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;

  // Resize arrays if needed
  if (infos->count >= infos->capacity) {
    infos->capacity *= 2;
    char **new_names = realloc(infos->names, infos->capacity * sizeof(char *));
    void **new_handles =
        realloc(infos->handles, infos->capacity * sizeof(void *));
    SymbolInfos *new_symbols =
        realloc(infos->symbols, infos->capacity * sizeof(SymbolInfos));

    if (!new_names || !new_handles || !new_symbols) {
      return 1; // Stop iteration on error
    }

    infos->names = new_names;
    infos->handles = new_handles;
    infos->symbols = new_symbols;
  }

  // Store the module name
  infos->names[infos->count] = strdup(info->dlpi_name);
  if (!infos->names[infos->count]) {
    return 1; // Stop iteration on error
  }

  // Try to get the handle
  infos->handles[infos->count] =
      dlopen(info->dlpi_name, RTLD_LAZY | RTLD_NOLOAD);

  if (he_create_symbol_info(&infos->symbols[infos->count], info) != CREATE_SUCCESS) {
    fprintf(stderr, "Failed to create symbol info for %s\n", infos->names[infos->count]);
  }
  infos->count++;
  return 0; // Continue iteration
}

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

  for (size_t i = 0; i < modules->count; i++) {
    if (modules->names[i])
      free(modules->names[i]);
    he_free_symbol_info(&modules->symbols[i]);
  }

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

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

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

  // Initial capacity
  result->capacity = 16;
  result->names = calloc(result->capacity, sizeof(char *));
  result->handles = calloc(result->capacity, sizeof(void *));
  result->symbols = calloc(result->capacity, sizeof(SymbolInfos));

  if (!result->names || !result->handles || !result->symbols) {
    if (result->names)
      free(result->names);
    if (result->handles)
      free(result->handles);
    if (result->symbols)
      free(result->symbols);
    free(result);
    return NULL;
  }

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

  return result;
}

int he_init() {
  assert(!module_infos);

  ModuleInfos *infos = gather_shared_libraries();
  if (!infos) {
    fprintf(stderr, "Failed to gather module infos.\n");
    return 1;
  }
  if (module_infos) {
    free_module_infos(module_infos);
  }
  module_infos = infos;
  return 0;
}

void he_deinit() { free_module_infos(module_infos); }

/**
 * Reloads a shared library module
 *
 * @param modules The ModuleInfos structure containing loaded modules
 * @param filename The name of the module to reload
 * @param updated_handle Pointer to store the new handle (can be NULL)
 * @return ReloadResult indicating success or failure
 */
static ReloadResult reload_module(ModuleInfos *modules, const char *filename,
                                  void **updated_handle) {
  if (!modules || !filename) {
    return 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
    if (modules->names[i] && strcmp(modules->names[i], filename) == 0) {
      found = 1;
      index = i;
      break;
    }

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

  if (!found) {
    return RELOAD_NOT_FOUND;
  }

  // Save the full path
  char *fullpath = strdup(modules->names[index]);
  if (!fullpath) {
    return RELOAD_OPEN_ERROR;
  }

  // Close the old handle
  if (modules->handles[index]) {
    if (dlclose(modules->handles[index]) != 0) {
      free(fullpath);
      return RELOAD_CLOSE_ERROR;
    }
  }

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

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

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

  free(fullpath);
  return RELOAD_SUCCESS;
}
/**
 * Helper function to print the result of a module reload
 */
static void print_reload_result(ReloadResult result, const char *filename) {
  switch (result) {
  case RELOAD_SUCCESS:
    printf("Successfully reloaded module: %s\n", filename);
    break;
  case RELOAD_NOT_FOUND:
    printf("Module not found: %s\n", filename);
    break;
  case RELOAD_CLOSE_ERROR:
    printf("Error closing module: %s\n", filename);
    break;
  case RELOAD_OPEN_ERROR:
    printf("Error reopening module: %s\n", filename);
    break;
  default:
    printf("Unknown error reloading module: %s\n", filename);
  }
}

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

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

  return result;
}

void he_print_module_infos() {
  assert(module_infos);

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

  printf("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;

    printf("%s: %p\n", modules->names[i], modules->handles[i]);
    if (modules->symbols) {
      for (int j = 0; j < modules->symbols->count; j++) {
        const void *addr = modules->symbols->addresses[j];
        const char *name = modules->symbols->names[j];
        printf("  %p: %s\n", addr, name);
      }
    }

    printf("\n");
  }
}