#include "symbols.h"

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

#include <dlfcn.h>
#include <elf.h>
#include <link.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static inline uptr
add_ptr_offset_if_invalid(uptr p, uptr offset,
                          struct sc_array_memreg *const regions) {
  if (memory_find_pointer(p, regions, NULL) != HILOAD_OK)
    return p + offset;
  return p;
}

/**
 * Gathers and populates symbols, given a dynamic module info
 *
 * Will clear and free the given SymbolInfo struct. Allocates enough memory to
 * hold found symbols.
 */
HiloadResult hi_create_symbol_info(SymbolInfos *symbols,
                                   struct sc_array_memreg *const regions,
                                   struct dl_phdr_info *info) {

  if (!symbols)
    return HILOAD_FAIL;

  hi_free_symbol_info(symbols);

  for (int i = 0; i < info->dlpi_phnum; i++) {
    const ElfW(Phdr) *phdr = &info->dlpi_phdr[i];

    // Look for the dynamic segment
    if (phdr->p_type != PT_DYNAMIC)
      continue;

    sc_log_debug("Dynamic Header:\n");
    sc_log_debug("p_type: %u\n", phdr->p_type);
    sc_log_debug("p_flags: %u\n", phdr->p_flags);
    sc_log_debug("p_offset: %#06lx\n", phdr->p_offset);
    sc_log_debug("p_vaddr: %#06lx\n", phdr->p_vaddr);
    sc_log_debug("p_paddr: %#06lx\n", phdr->p_paddr);
    sc_log_debug("p_filesz: %zu\n", phdr->p_filesz);
    sc_log_debug("p_memsz: %zu\n", phdr->p_memsz);
    sc_log_debug("p_align: %zu\n", phdr->p_align);

    const ElfW(Dyn) *dyn = (const ElfW(Dyn) *)(info->dlpi_addr + phdr->p_vaddr);
    const char *strtab = NULL;
    const ElfW(Sym) *symtab = NULL;
    size_t symtab_size = 0;
    size_t strtab_size = 0;

    uptr off = info->dlpi_addr;
    // Parse the dynamic table. Add offset if address is not in executable memory.
    // NOTE: Haven't found a better way to differentiate with items that have
    // relative address, and items that don't.
    for (; dyn->d_tag != DT_NULL; dyn++) {

      if (dyn->d_tag == DT_STRTAB) {
        uptr p = dyn->d_un.d_ptr;
        p = add_ptr_offset_if_invalid(p, off, regions);
        strtab = (const char *)p;

      } else if (dyn->d_tag == DT_STRSZ) {
        strtab_size = dyn->d_un.d_val;

      } else if (dyn->d_tag == DT_SYMTAB) {
        uptr p = dyn->d_un.d_ptr;
        p = add_ptr_offset_if_invalid(p, off, regions);
        symtab = (const ElfW(Sym) *)(p);

      } else if (dyn->d_tag == DT_SYMENT) {
        symtab_size = dyn->d_un.d_val;
      }
    }

    // Ensure we found the symbol and string tables
    if (!strtab || !symtab || strtab_size == 0 || symtab_size == 0) {
      sc_log_error("Failed to find symbol or string table in %s\n",
              info->dlpi_name);
      return HILOAD_FAIL;
    }

    // Iterate over the symbol table
    for (const ElfW(Sym) *sym = symtab;
         (const char *)sym < (const char *)symtab + symtab_size; sym++) {
      if (ELF64_ST_TYPE(sym->st_info) == STT_FUNC ||
          ELF64_ST_TYPE(sym->st_info) == STT_OBJECT) {
        const char *name = strdup(&strtab[sym->st_name]);
        void *address = (void *)(info->dlpi_addr + sym->st_value);

        // Store the symbol information in the  arrays
        sc_array_add(&symbols->names, name);
        sc_array_add(&symbols->addresses, address);
      }
    }
  }

  return HILOAD_OK;
}

void hi_free_symbol_info(SymbolInfos *symbols) {
  for (size_t i = 0; i < symbols->names.size; i++) {
    free((void *)symbols->names.elems[i]);
  }

  sc_array_term(&symbols->names);
  sc_array_term(&symbols->addresses);
}