#include "moduler.h"

#include "common.h"
#include "files.h"
#include "hielf.h"
#include "histring.h"
#include "logger.h"
#include "memmap.h"
#include "symbols.h"
#include "types.h"
#include "vector.h"

#include <dlfcn.h>
#include <errno.h>
#include <fcntl.h>
#include <link.h>
#include <stdalign.h>
#include <stdbool.h>
#include <sys/mman.h>
#include <time.h>
#include <unistd.h>

typedef struct {
  const char *filename;
  void *dlhandle;
  MemorySpan memreg;
} PatchData;

typedef struct {
  /// The original module name
  const char *modname;
  /// The associated symbols for that module
  VectorSymbol symbols;
} SymbolData;

vector_def(SymbolData, SymbolData);

// TODO: Stop being lazy and put this with the rest of the data
static VectorSymbolData symbol_data;

static VectorSymbol *symdat_get(const char *modname) {
  for (size_t i = 0; i < vector_size(&symbol_data); ++i) {
    SymbolData *symdat = &vector_at(&symbol_data, i);
    if (strcmp(modname, symdat->modname) == 0) {
      return &symdat->symbols;
    }
  }
  return NULL;
}

static void *adjust_if_relative(void *ptr, void *module_base) {
  uptr p = (uptr)ptr;
  if (p && (p < (uptr)module_base)) {
    p = (uptr)module_base + (uptr)ptr;
  }
  return (void *)p;
}

static HiResult moduler_collect_symbols(VectorSymbol *symbols,
                                        const char *module_name,
                                        void *module_base) {
  symbols_clear(symbols);

  HiResult ret = HI_FAIL;

  int patchfd = open(module_name, O_RDONLY);
  if (patchfd == -1) {
    log_error("Failed to open %s: %s\n", module_name, strerror(errno));
    return HI_FAIL;
  }

  if (elf_version(EV_CURRENT) == EV_NONE) {
    log_error("Failed to initialize libelf\n");
    close(patchfd);
    return HI_FAIL;
  }

  Elf *elf = elf_begin(patchfd, ELF_C_READ, NULL);
  if (!elf) {
    log_error("Failed to open elf file: %s\n", elf_errmsg(elf_errno()));
    goto cleanup;
  }

  if (elf_kind(elf) != ELF_K_ELF) {
    log_error("Not a valid ELF file\n");
    goto cleanup;
  }

  // Get ELF header
  GElf_Ehdr ehdr;
  if (gelf_getehdr(elf, &ehdr) == NULL) {
    log_error("Failed to get ELF header: %s\n", elf_errmsg(elf_errno()));
    goto cleanup;
  }

  // Iterate through sections to find symbol tables
  Elf_Scn *scn = NULL;
  while ((scn = elf_nextscn(elf, scn)) != NULL) {
    GElf_Shdr shdr;
    if (gelf_getshdr(scn, &shdr) != &shdr) {
      log_error("Failed to get section header: %s\n", elf_errmsg(elf_errno()));
      continue;
    }

    if (shdr.sh_type == SHT_SYMTAB) {
      // Get the string table for symbol names
      Elf_Scn *str_scn = elf_getscn(elf, shdr.sh_link);
      if (!str_scn) {
        log_error("Failed to get string table section: %s\n",
                  elf_errmsg(elf_errno()));
        continue;
      }

      GElf_Shdr str_shdr;
      if (gelf_getshdr(str_scn, &str_shdr) != &str_shdr) {
        log_error("Failed to get string table header: %s\n",
                  elf_errmsg(elf_errno()));
        continue;
      }

      // Get string table data
      Elf_Data *str_data = elf_getdata(str_scn, NULL);
      if (!str_data) {
        log_error("Failed to get string table data: %s\n",
                  elf_errmsg(elf_errno()));
        continue;
      }

      // Get symbol table data
      Elf_Data *sym_data = elf_getdata(scn, NULL);
      if (!sym_data) {
        log_error("Failed to get symbol table data: %s\n",
                  elf_errmsg(elf_errno()));
        continue;
      }

      // Calculate number of symbols
      size_t sym_count = shdr.sh_size / shdr.sh_entsize;

      for (size_t i = 0; i < sym_count; ++i) {
        GElf_Sym sym;
        if (gelf_getsym(sym_data, i, &sym) != &sym) {
          log_error("Failed to get symbol: %s\n", elf_errmsg(elf_errno()));
          continue;
        }

        // Assumption: Only symbols with size are defined here, and those
        // without are safely ignored. Linker will handle them.
        // Assumption: completed.0 seems to be size 1 and appears locally, lets
        // assume its useless
        if (sym.st_size <= 1) {
          continue;
        }

        const char *name = (const char *)(str_data->d_buf) + sym.st_name;
        if (sym.st_name == 0 || strlen(name) == 0) {
          continue; // Skip unnamed symbols
        }

        void *sym_addr = (void *)((uintptr_t)module_base + sym.st_value);
        SymbolBind binding =
            symbol_bind_from_efibind(GELF_ST_BIND(sym.st_info));
        SymbolType type = symbol_type_from_efitype(GELF_ST_TYPE(sym.st_info));
        size_t size = sym.st_size;

        // Gather global symbols and local object symbols. Local functions are
        // handled automatically, but data from the running process needs to
        // be copied over
        if (binding == HI_SYMBOL_BIND_GLOBAL ||
            (binding == HI_SYMBOL_BIND_LOCAL &&
             type == HI_SYMBOL_TYPE_OBJECT)) {
          Symbol hisym = {.name = strdup(name),
                          .size = size,
                          .binding = binding,
                          .type = type,
                          .address = sym_addr};

          vector_add(symbols, hisym);
        }
      }
    }
  }

  ret = HI_OK;

cleanup:
  elf_end(elf);
  close(patchfd);

  return ret;
}

static HiResult moduler_patch_functions(VectorSymbol *psymbols,
                                        MemorySpan module_memory) {

  void *module_base = (void *)module_memory.start;
  size_t module_size = module_memory.end - module_memory.start;

  void *plt = NULL;
  ElfW(Rela) *jmprel = NULL;
  ElfW(Rela) *rela = NULL;
  ElfW(Sym) *symtab = NULL;
  char *strtab = NULL;
  size_t relasz = 0;

  Elf *elf = NULL;
  ElfW(Dyn) *dyn = hi_elf_find_dynseg(module_base, module_size, &elf);

  for (ElfW(Dyn) *d = dyn; d->d_tag != DT_NULL; d++) {

    log_debugv("%s\n", hi_elf_dyntostr(d->d_tag));
    switch (d->d_tag) {
    case DT_RELASZ:
      relasz = d->d_un.d_val;
      break;
    case DT_RELA:
      rela = (ElfW(Rela) *)d->d_un.d_ptr;
      break;
    case DT_PLTGOT:
      plt = (ElfW(Rela) *)d->d_un.d_ptr;
      break;
    case DT_JMPREL:
      jmprel = (ElfW(Rela) *)d->d_un.d_ptr;
      break;
    case DT_SYMTAB:
      symtab = (ElfW(Sym) *)d->d_un.d_ptr;
      break;
    case DT_STRTAB:
      strtab = (char *)d->d_un.d_ptr;
      break;
    }
  }

  // Adjust relative pointers by adding the module base address
  plt = adjust_if_relative(plt, module_base);
  rela = adjust_if_relative(rela, module_base);
  symtab = adjust_if_relative(symtab, module_base);
  strtab = adjust_if_relative(strtab, module_base);

  if ((!rela && !jmprel) || !symtab || !strtab || relasz == 0) {
    log_error("Missing required dynamic information\n");
    elf_end(elf);
    return HI_FAIL;
  }

  int found_count = 0;
  size_t rela_count = relasz / sizeof(ElfW(Rela));

  printf("Processing %zu relocations\n", rela_count);

  size_t num_symbols = vector_size(psymbols);
  ElfW(Rela) *r = NULL;
  for (size_t i = 0; i < rela_count; i++) {

    r = jmprel ? &jmprel[i] : &rela[i];
    int sym_idx = GELF_R_SYM(r->r_info);
    int rel_type = GELF_R_TYPE(r->r_info);

    // We're only interested in JMP_SLOT relocations (GOT entries for functions)
    if (rel_type != R_X86_64_JUMP_SLOT && rel_type != R_386_JMP_SLOT) {
      continue;
    }

    ElfW(Sym) *sym = &symtab[sym_idx];
    char *name = strtab + sym->st_name;

    // Calculate the GOT entry address
    void **got_entry = (void **)((char *)module_base + r->r_offset);

    // Check if this is a symbol we want to patch
    for (size_t j = 0; j < num_symbols; j++) {
      Symbol *s = &vector_at(psymbols, j);
      if (strcmp(s->name, name) == 0) {
        s->got_entry = got_entry;
        s->orig_address = *got_entry; // Save the original function

        *got_entry = s->address;
        log_debug("Found GOT entry for '%s' at %p (points to %p)\n", name,
                  got_entry, *got_entry);

        found_count++;
      }
    }
  }
  elf_end(elf);
  return HI_OK;
}

static PatchData moduler_create_patch(ModuleData *module) {

  time_t now = time(NULL);
  struct tm *t = localtime(&now);
  if (t == NULL) {
    return (PatchData){0};
  }

  char file_append[32];
  const char *patch_format = ".%Y%m%d%H%M%S.patch";
  if (strftime(file_append, sizeof(file_append), patch_format, t) == 0) {
    log_error("Failed to create patch filename.\n");
    return (PatchData){0};
  }

  char filename[512];
  size_t written = string_concat_buf(sizeof(filename), filename,
                                     module->original_name, file_append);

  if (written == 0) {
    log_error("Failed to concat %s and %s\n", module->original_name, ".patch");
    return (PatchData){0};
  }

  file_copy(module->original_name, filename);

  PatchData data = {.filename = strdup(filename)};
  return data;
}

HiResult moduler_reload(VectorModuleData *modules, size_t modindx) {

  ModuleData *module = &vector_at(modules, modindx);
  // Return OK because this isn't an error case. We just don't do it yet
  // because we only handle reloading a complete solib. Can't do that with
  // executables.
  if (modinfo_has(module->info, HI_MODULE_STATE_EXEC)) {
    module->info = modinfo_clear(module->info, HI_MODULE_STATE_DIRTY);
    return HI_OK;
  }

  PatchData patch = moduler_create_patch(module);
  if (!patch.filename) {
    log_error("Couldn't create patch for %s\n", module->name);
    return HI_FAIL;
  }

  // Load patch
  dlerror(); // clear any previous errors
  log_debug("Opening patch: %s\n", patch.filename);
  void *new_handle = dlopen(patch.filename, RTLD_LAZY);
  if (!new_handle) {
    log_errorv("Couldn't load: %s\n", dlerror());
    return HI_FAIL;
  }
  patch.dlhandle = new_handle;

  VectorMemoryMap memmaps;
  vector_init(&memmaps);
  if (!HIOK(memmaps_from_process(&memmaps))) {
    log_error("Failed to collect memory information for process\n");
    return HI_FAIL;
  }

  patch.memreg = memmaps_find_by_name(patch.filename, &memmaps);
  void *patch_base = (void *)patch.memreg.start;

  VectorSymbol patch_symbols;
  symbols_init(&patch_symbols);

  HiResult ret =
      moduler_collect_symbols(&patch_symbols, patch.filename, patch_base);
  if (!HIOK(ret)) {
    log_error("Failed to gather symbols for %s\n", patch.filename);
    return HI_FAIL;
  }

  for (size_t i = 0; i < vector_size(modules); ++i) {
    ModuleData mod = vector_at(modules, i);

    if (!modinfo_has(mod.info, HI_MODULE_STATE_PATCHABLE))
      continue;

    log_debug("Patching: %s\n", mod.name);
    MemorySpan module_memory = memmaps_find_by_name(mod.name, &memmaps);
    moduler_patch_functions(&patch_symbols, module_memory);

    // This relies on the patch filename being different only by the append
    if (strncmp(mod.original_name, patch.filename, strlen(mod.original_name)) ==
        0) {

      // If patch is for the same module, also collect local object symbols for
      // coping those over.
      VectorSymbol *module_symbols = symdat_get(mod.original_name);

      // Copy old data to new data location. Breaks if layout changes,
      // e.g. struct fields are moved around
      // Also update our symbol cache for the module. We need to add new
      // symbols, but don't care much about deleting old ones
      for (size_t j = 0; j < vector_size(&patch_symbols); ++j) {
        Symbol *psym = &vector_at(&patch_symbols, j);
        if (psym->type == HI_SYMBOL_TYPE_OBJECT) {
          Symbol *msym = symbol_find(module_symbols, psym);
          if (msym) {
            size_t copy_size = MIN(psym->size, msym->size);
            memcpy(psym->address, msym->address, copy_size);
            log_debug("Copied data for symbol: %s\n", psym->name);

            // Maintain our current symbol references with the new address
            msym->address = psym->address;
          } else {
            // A new symbol has been added in the patch
            symbols_add(module_symbols, symbol_copy(psym));
            log_debug("Found new symbol from patch: %s\n", psym->name);
          }
        }
      }
    }
  }
  symbols_term(&patch_symbols);

  module->info = modinfo_clear(module->info, HI_MODULE_STATE_DIRTY);
  dlclose(module->dlhandle);

  // Update module reference
  if (module->name != module->original_name)
    free((char *)module->name);
  module->name = patch.filename;
  module->address = patch.memreg.start;
  module->dlhandle = patch.dlhandle;

  return HI_OK;
}

void moduler_init(const VectorModuleData *modules) {
  vector_init(&symbol_data);

  ModuleData mod = {0};
  vector_foreach(modules, mod) {
    SymbolData symdata = {0};
    symdata.modname = mod.original_name;
    vector_add(&symbol_data, symdata);

    if (modinfo_has(mod.info, HI_MODULE_STATE_PATCHABLE)) {
      moduler_collect_symbols(&vector_last(&symbol_data).symbols,
                              mod.original_name, (void *)mod.address);
    }

    vector_init(&symdata.symbols);
  }
}

void moduler_deinit() {
  SymbolData symdata = {0};
  vector_foreach(&symbol_data, symdata) { symbols_term(&symdata.symbols); }
  vector_term(&symbol_data);
}