#include "elf.h"

#include <elf.h>
#include <fcntl.h> // for open
#include <gelf.h>  // for GElf_Ehdr
#include <libelf.h>
#include <string.h>
#include <sys/utsname.h>
#include <unistd.h>

// #include <dwarf.h>
// #include <elfutils/libdw.h>
// #include <dwarf.h>
// #include <elfutils/libdw.h>

#define NS_NAME_LEN 128

struct sym_section_ctx {
  Elf_Data *syms;
  Elf_Data *symstrs;
  Elf_Data *rel_data;
  int is_reloc;
  int is_plt;
  int sym_count;
  int plt_rel_type;
  unsigned long plt_offset;
  unsigned long plt_entsize;
};

struct symbol_sections_ctx {
  sym_section_ctx symtab;
  sym_section_ctx symtab_in_dynsym;
  sym_section_ctx dynsymtab;
};

struct section_info {
  Elf_Scn *sec;
  GElf_Shdr *hdr;
};

struct plt_ctx {
  section_info dynsym;
  section_info plt_rel;
  section_info plt;
};

static int global_mnt_ns_fd = -1;
static char global_mnt_ns_name[NS_NAME_LEN];
static char g_mnt_ns_name[128];

int init_global_env(void) {
  struct utsname utsname;
  int ret = uname(&utsname);
  if (ret < 0) {
    return -1;
  }

  if (readlink("/proc/1/ns/mnt", g_mnt_ns_name, sizeof(g_mnt_ns_name)) < 0) {
    return -1;
  }

  ret = readlink("/proc/1/ns/mnt", global_mnt_ns_name, NS_NAME_LEN);
  if (ret <= 0) {
    return -1;
  } else {
    global_mnt_ns_fd = open("/proc/1/ns/mnt", 0);
    if (global_mnt_ns_fd < 0) {
      return -1;
    }
  }
  return 0;
}

static int get_symbols_in_section(sym_section_ctx *sym, Elf *elf, Elf_Scn *sec,
                                  GElf_Shdr *shdr, int is_reloc) {
  sym->syms = elf_getdata(sec, NULL);
  if (!sym->syms) {
    return -1;
  }

  Elf_Scn *symstrs_sec = elf_getscn(elf, shdr->sh_link);
  if (!sec) {
    return -1;
  }

  sym->symstrs = elf_getdata(symstrs_sec, NULL); // 获取section data
  if (!sym->symstrs) {
    return -1;
  }

  sym->sym_count =
      shdr->sh_size / shdr->sh_entsize; // 获取section中的symbol数量
  sym->is_plt = 0;
  sym->is_reloc = is_reloc; // 是否是可重定位文件

  return 0;
}

static int get_plt_symbols_in_section(sym_section_ctx *sym, Elf *elf,
                                      plt_ctx *plt) {
  sym->syms = elf_getdata(plt->dynsym.sec, NULL);
  if (!sym->syms) {
    return -1;
  }

  sym->rel_data = elf_getdata(plt->plt_rel.sec, NULL);
  if (!sym->rel_data) {
    return -1;
  }

  Elf_Scn *symstrs_sec = elf_getscn(elf, plt->dynsym.hdr->sh_link);
  if (!symstrs_sec) {
    return -1;
  }

  sym->symstrs = elf_getdata(symstrs_sec, NULL);
  if (!sym->symstrs) {
    return -1;
  }

  sym->is_plt = 1;
  sym->plt_entsize = plt->plt.hdr->sh_type;
  sym->plt_offset = plt->plt.hdr->sh_offset;
  sym->sym_count = plt->plt_rel.hdr->sh_size /
                   plt->plt_rel.hdr->sh_entsize; // 获取section中的symbol数量
  sym->plt_rel_type = plt->plt_rel.hdr->sh_type;

  return 0;
}

static int is_function(const GElf_Sym *sym) {
  return GELF_ST_TYPE(sym->st_info) == STT_FUNC && sym->st_name != 0 &&
         sym->st_shndx != SHN_UNDEF;
}

static void __get_symbol_without_plt(std::set<symbol> &ss, sym_section_ctx *tab,
                                     Elf *elf) {
  GElf_Sym sym;
  int index = 0;
  const char *sym_name;
  symbol s;
  s.end = 0;
  s.start = 0;

  while (index < tab->sym_count) {
    gelf_getsym(tab->syms, index, &sym);
    index++;
    if (sym.st_shndx == SHN_ABS) {
      continue;
    }
    if (!is_function(&sym)) {
      continue;
    }
    sym_name = (const char *)tab->symstrs->d_buf + sym.st_name;
    if (tab->is_reloc) {
      Elf_Scn *sec = elf_getscn(elf, sym.st_shndx);
      if (!sec) {
        continue;
      }
      GElf_Shdr shdr;
      gelf_getshdr(sec, &shdr);
      sym.st_value -= shdr.sh_addr - shdr.sh_offset;
    }
    s.start = sym.st_value & 0xffffffff;
    s.end = s.start + sym.st_size;
    s.ip = s.start;
    s.name = sym_name;
    // printf("name: %s, start: 0x%lx, end: 0x%lx\n", s.name.c_str(), s.start,
    // s.end);
    ss.insert(s);
  }
}

static void __get_symbol(std::set<symbol> &ss, symbol_sections_ctx *si,
                         Elf *elf) {
  symbol s;
  s.end = 0;
  s.start = 0;

  if (!si->symtab.syms && !si->dynsymtab.syms) {
    return;
  }

  sym_section_ctx *tab = &si->symtab;
  __get_symbol_without_plt(ss, tab, elf);
  tab = &si->symtab_in_dynsym;
  __get_symbol_without_plt(ss, tab, elf);
}

static void __get_plt_symbol(std::set<symbol> &ss, symbol_sections_ctx *si,
                             Elf *elf) {
  symbol s;
  GElf_Sym sym;
  int symidx;
  int index = 0;
  const char *sym_name = NULL;

  s.end = 0;
  s.start = 0;

  if (!si->dynsymtab.syms) {
    return;
  }

  while (index < si->dynsymtab.sym_count) {
    if (si->dynsymtab.plt_rel_type == SHT_RELA) {
      GElf_Rela pos_mem, *pos;
      pos = gelf_getrela(si->dynsymtab.rel_data, index, &pos_mem);
      symidx = GELF_R_SYM(pos->r_info);
    } else if (si->dynsymtab.plt_rel_type == SHT_REL) {
      GElf_Rel pos_mem, *pos;
      pos = gelf_getrel(si->dynsymtab.rel_data, index, &pos_mem);
      symidx = GELF_R_SYM(pos->r_info);
    } else {
      return;
    }
    index++;
    si->dynsymtab.plt_offset += si->dynsymtab.plt_entsize;
    gelf_getsym(si->dynsymtab.syms, symidx, &sym);

    sym_name = (const char *)si->dynsymtab.symstrs->d_buf + sym.st_name;
    s.start = si->dynsymtab.plt_offset;
    s.end = s.start + si->dynsymtab.plt_entsize;
    s.ip = s.start;
    s.name = sym_name;
    // printf("plt name: %s, start: 0x%lx, end: 0x%lx\n", s.name.c_str(),
    // s.start, s.end);
    ss.insert(s);
  }
}

static void get_all_symbols(std::set<symbol> &ss, symbol_sections_ctx *si,
                            Elf *elf) {
  __get_symbol(ss, si, elf);
  __get_plt_symbol(ss, si, elf);
}

// static void get_debug_symbols(std::set<symbol>& ss, Dwarf* dw) {
//   Dwarf_Off die_offset = 0, next_offset;
//   size_t header_size;
//   printf("1\n");
//   while (dwarf_nextcu(dw, die_offset, &next_offset, &header_size, NULL, NULL,
//   NULL) == 0) {
//     printf("2\n");
//     Dwarf_Die cudie;
//     printf("3\n");
//     if (dwarf_offdie(dw, die_offset + header_size, &cudie) == NULL) {
//       printf("4\n");
//       continue;
//     }
//     printf("4\n");
//     Dwarf_Die die;
//     if (dwarf_child(&cudie, &die) != 0) {
//       printf("5\n");
//       continue;
//     }

//     do {
//       const char* die_name = dwarf_diename(&die);
//       printf("6\n");
//       if (!die_name) {
//         continue;
//       }

//       Dwarf_Attribute attr_mem;
//       Dwarf_Attribute* attr = dwarf_attr(&die, DW_AT_low_pc, &attr_mem);
//       if (attr) {
//         Dwarf_Addr low_pc;
//         if (dwarf_formaddr(attr, &low_pc) == 0) {
//           printf("name: %s, low_pc: %lx\n", die_name, low_pc);
//           // symbol s;
//           // s.name = die_name;
//           // s.start = low_pc;
//           // s.end = s.start; // You need to find the high_pc to set the end.

//           // ss.insert(s);
//         }
//       }
//     } while (dwarf_siblingof(&die, &die) == 0);

//     die_offset = next_offset;
//   }
// }
//     die_offset = next_offset;
//   }
// }

bool is_stripped(const char *path) {
  char command[256];
  // snprintf(command, sizeof(command), "file %s", path);

  FILE *fp = popen(command, "r");
  if (fp == NULL) {
    // handle error
    return false;
  }

  char output[256];
  fgets(output, sizeof(output), fp);
  pclose(fp);

  return strstr(output, "stripped") != NULL;
}

#define MAX_LINE_LENGTH 1024
void get_symbol_from_elf_gdb(std::set<symbol> &ss, const char *path,
                             size_t file_base_addr = 0) {
  FILE *fp;
  char cmd[MAX_LINE_LENGTH];
  char line[MAX_LINE_LENGTH];
  // 构建 GDB 命令
  snprintf(cmd, sizeof(cmd),
           "gdb -batch -ex \"file %s\" -ex \"info functions\"", path);

  // 执行 GDB 命令并获取输出
  fp = popen(cmd, "r");
  if (fp == NULL) {
    perror("popen");
    return;
  }

  bool non_debugging_symbols = false;
  std::map<unsigned long long, std::string> symbol_map;

  // 读取并解析 GDB 的输出
  while (fgets(line, sizeof(line), fp) != NULL) {
    unsigned long long address;
    char name[MAX_LINE_LENGTH];
    // printf("line: %s", line);
    if (!non_debugging_symbols) {
      if (strstr(line, "Non-debugging symbols:") != NULL) {
        non_debugging_symbols = true;
      }
      continue;
    }

    char *token = strtok(line, " ");
    if (token != NULL) {
      sscanf(token, "%llx", &address);
      token = strtok(NULL, "\n");
      if (token != NULL) {
        strcpy(name, token);
      }
      // if name == ":", 跳过
      if (name[0] == ':') {
        continue;
      }
      symbol_map[address] = std::string(name);
    }
    // 解析函数名和地址
    // if (sscanf(line, "%llx %s", &address, name) == 2)
    // {
    //     // if name == ":", 跳过
    //     if (name[0] == ':')
    //     {
    //         continue;
    //     }
    //     // printf("Name %s, Address %llx\n", name, address);
    //     // name -> std::string
    //     symbol_map[address] = std::string(name);
    // }
  }

  symbol s;

  // 插入范围到新的map
  auto it = symbol_map.begin();
  auto next_it = it;
  ++next_it;

  for (; next_it != symbol_map.end(); ++it, ++next_it) {
    s.start = it->first - file_base_addr;
    s.end = next_it->first - file_base_addr;
    s.ip = s.start;
    s.name = it->second;
    // printf("gdb name: %s, start: 0x%lx, end: 0x%lx\n", s.name.c_str(),
    // s.start, s.end);
    ss.insert(s);
  }
  // 关闭 GDB 进程
  pclose(fp);
  //
}

bool get_symbol_from_elf(std::set<symbol> &ss, const char *path,
                         size_t file_base_addr) {
  static int first_init = 0;

  if (!first_init) {
    first_init = true;
    init_global_env();
  }

  int is_reloc = 0;
  elf_version(EV_CURRENT);

  int fd = open(path, O_RDONLY);

  Elf *elf = elf_begin(fd, ELF_C_READ, NULL); // 指向elf文件的指针
  if (elf == NULL) {
    close(fd);
    return false;
  }

  Elf_Kind ek = elf_kind(elf); // 获取elf文件类型
  if (ek != ELF_K_ELF) {
    elf_end(elf);
    close(fd);
    return false;
  }
  GElf_Ehdr hdr; // elf文件头
  if (gelf_getehdr(elf, &hdr) == NULL) {
    elf_end(elf);
    close(fd);
    return false;
  }

  if (hdr.e_type == ET_EXEC) { // 可执行文件
    is_reloc = 1;
  }

  if (!elf_rawdata(elf_getscn(elf, hdr.e_shstrndx), NULL)) { //
    elf_end(elf);
    close(fd);
    return false;
  }

  GElf_Shdr shdr;
  GElf_Shdr symtab_shdr;
  GElf_Shdr dynsym_shdr;
  GElf_Shdr plt_shdr;
  GElf_Shdr plt_rel_shdr;
  memset(&shdr, 0, sizeof(shdr));
  memset(&symtab_shdr, 0, sizeof(symtab_shdr));
  memset(&dynsym_shdr, 0, sizeof(dynsym_shdr));
  memset(&plt_shdr, 0, sizeof(plt_shdr));
  memset(&plt_rel_shdr, 0, sizeof(plt_rel_shdr));

  Elf_Scn *sec = NULL;
  Elf_Scn *dynsym_sec = NULL;
  Elf_Scn *symtab_sec = NULL;
  Elf_Scn *plt_sec = NULL;
  Elf_Scn *plt_rel_sec = NULL;
  Elf_Scn *debug_sec = NULL;
  // 遍历elf文件的section header table
  while ((sec = elf_nextscn(elf, sec)) != NULL) {
    char *str;
    gelf_getshdr(sec, &shdr); // 获取section header
    str = elf_strptr(elf, hdr.e_shstrndx, shdr.sh_name);

    if (str && strcmp(".symtab", str) == 0) { // .symtab section
      symtab_sec = sec;
      memcpy(&symtab_shdr, &shdr, sizeof(dynsym_shdr));
    }
    if (str && strcmp(".dynsym", str) == 0) { // .dynsym section
      dynsym_sec = sec;
      memcpy(&dynsym_shdr, &shdr, sizeof(dynsym_shdr));
    }
    if (str && strcmp(".rela.plt", str) == 0) { // .rela.plt section
      plt_rel_sec = sec;
      memcpy(&plt_rel_shdr, &shdr, sizeof(plt_rel_shdr));
    }
    if (str && strcmp(".plt", str) == 0) { // .plt section
      plt_sec = sec;
      memcpy(&plt_shdr, &shdr, sizeof(plt_shdr));
    }
    if (str && strcmp(".gnu.prelink_undo", str) == 0) {
      is_reloc = 1;
    }
  }

  plt_ctx plt;
  plt.dynsym.hdr = &dynsym_shdr;
  plt.dynsym.sec = dynsym_sec;
  plt.plt.hdr = &plt_shdr;
  plt.plt.sec = plt_sec;
  plt.plt_rel.hdr = &plt_rel_shdr;
  plt.plt_rel.sec = plt_rel_sec;

  symbol_sections_ctx si;
  memset(&si, 0, sizeof(si));
  if (symtab_sec) {
    get_symbols_in_section(&si.symtab, elf, symtab_sec, &symtab_shdr, is_reloc);
  } else {
    if (hdr.e_type != ET_DYN) { // no shared object
      // printf("stripped, path: %s\n", path);
      get_symbol_from_elf_gdb(ss, path, file_base_addr);
      return true;
    }
  }
  if (dynsym_sec) {
    get_symbols_in_section(&si.symtab_in_dynsym, elf, dynsym_sec, &dynsym_shdr,
                           is_reloc);
  }
  if (dynsym_sec && plt_sec) {
    get_plt_symbols_in_section(&si.dynsymtab, elf, &plt);
  }

  get_all_symbols(ss, &si, elf);
  elf_end(elf);

  close(fd);
  return true;
}