#define _GNU_SOURCE
#include "heload/symbols.h"

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

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

  if (!symbols)
    return CREATE_FAILED;

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

    printf("Dynamic Header:\n");
    printf("p_type: %u\n", phdr->p_type);
    printf("p_flags: %u\n", phdr->p_flags);
    printf("p_offset: %#06x\n", phdr->p_offset);
    printf("p_vaddr: %#06x\n", phdr->p_vaddr);
    printf("p_paddr: %#06x\n", phdr->p_paddr);
    printf("p_filesz: %zu\n", phdr->p_filesz);
    printf("p_memsz: %zu\n", phdr->p_memsz);
    printf("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;

        // Parse the dynamic table
    for (; dyn->d_tag != DT_NULL; dyn++) {
      if (dyn->d_tag == DT_STRTAB) {
        strtab = (const char *)(dyn->d_un.d_ptr);
      } else if (dyn->d_tag == DT_STRSZ) {
        strtab_size = dyn->d_un.d_val;
      } else if (dyn->d_tag == DT_SYMTAB) {
        symtab = (const ElfW(Sym) *)(dyn->d_un.d_ptr);
      } 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) {
      fprintf(stderr, "Failed to find symbol or string table in %s\n",
              info->dlpi_name);
      return CREATE_FAILED;
    }

    symbols->capacity = symtab_size / sizeof(ElfW(Sym));

    symbols->names = calloc(symbols->capacity, sizeof(char *));
    if (!symbols->names) {
      fprintf(stderr, "Failed to allocate memory for symbol names.\n");
      return CREATE_FAILED;
    }

    symbols->addresses = calloc(symbols->capacity, sizeof(void *));
    if (!symbols->addresses) {
      fprintf(stderr, "Failed to allocate memory for symbol addresses.\n");
      return CREATE_FAILED;
    }

    // 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 struct of arrays
        if (symbols->count < symbols->capacity) {
          symbols->names[symbols->count] = (char *)name;
          symbols->addresses[symbols->count] = address;
          symbols->count++;
        } else {
          fprintf(stderr, "Symbol table capacity exceeded!\n");
          return CREATE_FAILED;
        }
      }
    }
  }

  return CREATE_SUCCESS;
}

void he_free_symbol_info(SymbolInfos *symbols) {
  for (size_t i = 0; i < symbols->count; i++) {
    free(symbols->names[i]);
  }
  free(symbols->names);
  free(symbols->addresses);
  symbols->count = 0;
  symbols->capacity = 0;
}