#include "PlatformUtil/ProcessRuntimeUtility.h"

#include <elf.h>
#include <dlfcn.h>
#include <link.h>
#include <sys/mman.h>

#include <string>
#include <string.h>

#include <vector>

#define LINE_MAX 2048

// ================================================================
// GetProcessMemoryLayout

static bool memory_region_comparator(MemoryRegion a, MemoryRegion b) {
    return ((addr_t) a.address < (addr_t) b.address);
}

std::vector<MemoryRegion> ProcessRuntimeUtility::GetProcessMemoryLayout() {
    std::vector<MemoryRegion> ProcessMemoryLayout;

    FILE *fp = fopen("/proc/self/maps", "r");
    if (fp == nullptr)
        return ProcessMemoryLayout;

    while (!feof(fp)) {
        char line_buffer[LINE_MAX + 1];
        fgets(line_buffer, LINE_MAX, fp);

        // ignore the rest of characters
        if (strlen(line_buffer) == LINE_MAX && line_buffer[LINE_MAX] != '\n') {
            // Entry not describing executable data. Skip to end of line to set up
            // reading the next entry.
            int c;
            do {
                c = getc(fp);
            } while ((c != EOF) && (c != '\n'));
            if (c == EOF)
                break;
        }

        addr_t region_start, region_end;
        addr_t region_offset;
        char permissions[5] = {'\0'}; // Ensure NUL-terminated string.
        uint8_t dev_major = 0;
        uint8_t dev_minor = 0;
        long inode = 0;
        int path_index = 0;

        // Sample format from man 5 proc:
        //
        // address           perms offset  dev   inode   pathname
        // 08048000-08056000 r-xp 00000000 03:0c 64593   /usr/sbin/gpm
        //
        // The final %n term captures the offset in the input string, which is used
        // to determine the path name. It *does not* increment the return value.
        // Refer to man 3 sscanf for details.
        if (sscanf(line_buffer,
                   "%" PRIxPTR "-%" PRIxPTR " %4c "
                   "%" PRIxPTR " %hhx:%hhx %ld %n",
                   &region_start, &region_end, permissions, &region_offset, &dev_major, &dev_minor,
                   &inode,
                   &path_index) < 7) {
            FATAL("/proc/self/maps parse failed!");
            fclose(fp);
            return ProcessMemoryLayout;
        }

        MemoryPermission permission;
        if (permissions[0] == 'r' && permissions[1] == 'w') {
            permission = MemoryPermission::kReadWrite;
        } else if (permissions[0] == 'r' && permissions[2] == 'x') {
            permission = MemoryPermission::kReadExecute;
        } else if (permissions[0] == 'r' && permissions[1] == 'w' && permissions[2] == 'x') {
            permission = MemoryPermission::kReadWriteExecute;
        } else {
            permission = MemoryPermission::kNoAccess;
        }

#if 0
        DLOG(0, "%p --- %p", region_start, region_end);
#endif

        ProcessMemoryLayout.push_back(
                MemoryRegion{(void *) region_start, region_end - region_start, permission});
    }
    std::sort(ProcessMemoryLayout.begin(), ProcessMemoryLayout.end(), memory_region_comparator);

    fclose(fp);
    return ProcessMemoryLayout;
}

// ================================================================
// GetProcessModuleMap

static std::vector<RuntimeModule> get_process_map_with_proc_maps() {
    std::vector<RuntimeModule> ProcessModuleMap;

    FILE *fp = fopen("/proc/self/maps", "r");
    if (fp == nullptr)
        return ProcessModuleMap;

    while (!feof(fp)) {
        char line_buffer[LINE_MAX + 1];
        fgets(line_buffer, LINE_MAX, fp);

        // ignore the rest of characters
        if (strlen(line_buffer) == LINE_MAX && line_buffer[LINE_MAX] != '\n') {
            // Entry not describing executable data. Skip to end of line to set up
            // reading the next entry.
            int c;
            do {
                c = getc(fp);
            } while ((c != EOF) && (c != '\n'));
            if (c == EOF)
                break;
        }

        addr_t region_start, region_end;
        addr_t region_offset;
        char permissions[5] = {'\0'}; // Ensure NUL-terminated string.
        uint8_t dev_major = 0;
        uint8_t dev_minor = 0;
        long inode = 0;
        int path_index = 0;

        // Sample format from man 5 proc:
        //
        // address           perms offset  dev   inode   pathname
        // 08048000-08056000 r-xp 00000000 03:0c 64593   /usr/sbin/gpm
        //
        // The final %n term captures the offset in the input string, which is used
        // to determine the path name. It *does not* increment the return value.
        // Refer to man 3 sscanf for details.
        if (sscanf(line_buffer,
                   "%" PRIxPTR "-%" PRIxPTR " %4c "
                   "%" PRIxPTR " %hhx:%hhx %ld %n",
                   &region_start, &region_end, permissions, &region_offset, &dev_major, &dev_minor,
                   &inode,
                   &path_index) < 7) {
            FATAL("/proc/self/maps parse failed!");
            fclose(fp);
            return ProcessModuleMap;
        }

        // check header section permission
        if (strcmp(permissions, "r--p") != 0 && strcmp(permissions, "r-xp") != 0)
            continue;

        // check elf magic number
        ElfW(Ehdr) *header = (ElfW(Ehdr) *) region_start;
        if (memcmp(header->e_ident, ELFMAG, SELFMAG) != 0) {
            continue;
        }

        char *path_buffer = line_buffer + path_index;
        if (*path_buffer == 0 || *path_buffer == '\n' || *path_buffer == '[')
            continue;
        RuntimeModule module;

        // strip
        if (path_buffer[strlen(path_buffer) - 1] == '\n') {
            path_buffer[strlen(path_buffer) - 1] = 0;
        }
        strncpy(module.path, path_buffer, sizeof(module.path));
        module.load_address = (void *) region_start;
        ProcessModuleMap.push_back(module);

#if 0
        DLOG(0, "module: %s", module.path);
#endif
    }

    fclose(fp);
    return ProcessModuleMap;
}

#if defined(__LP64__)
static std::vector<RuntimeModule> get_process_map_with_linker_iterator() {
  std::vector<RuntimeModule> ProcessModuleMap;

  static int (*dl_iterate_phdr_ptr)(int (*)(struct dl_phdr_info *, size_t, void *), void *);
  dl_iterate_phdr_ptr = (__typeof(dl_iterate_phdr_ptr))dlsym(RTLD_DEFAULT, "dl_iterate_phdr");
  if (dl_iterate_phdr_ptr == NULL) {
    return ProcessModuleMap;
  }

  dl_iterate_phdr_ptr(
      [](dl_phdr_info *info, size_t size, void *data) {
        RuntimeModule module = {0};
        if (info->dlpi_name && info->dlpi_name[0] == '/')
          strcpy(module.path, info->dlpi_name);

        module.load_address = (void *)info->dlpi_addr;
        for (size_t i = 0; i < info->dlpi_phnum; ++i) {
          if (info->dlpi_phdr[i].p_type == PT_LOAD) {
            uintptr_t load_bias = (info->dlpi_phdr[i].p_vaddr - info->dlpi_phdr[i].p_offset);
            module.load_address = (void *)((addr_t)module.load_address + load_bias);
            break;
          }
        }

        // push to vector
        auto ProcessModuleMap = reinterpret_cast<std::vector<RuntimeModule> *>(data);
        ProcessModuleMap->push_back(module);
        return 0;
      },
      (void *)&ProcessModuleMap);

  return ProcessModuleMap;
}
#endif

std::vector<RuntimeModule> ProcessRuntimeUtility::GetProcessModuleMap() {
#if defined(__LP64__) && 0
    // TODO: won't resolve main binary
    return get_process_map_with_linker_iterator();
#else
    return get_process_map_with_proc_maps();
#endif
}

RuntimeModule ProcessRuntimeUtility::GetProcessModule(const char *name) {
    std::vector<RuntimeModule> ProcessModuleMap = GetProcessModuleMap();
    for (auto module : ProcessModuleMap) {
        if (strstr(module.path, name) != 0) {
            return module;
        }
    }
    return RuntimeModule{0};
}