#include "symbol.h"

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

#include <algorithm>
#include <vector>

#include "elf.h"

symbol_parser g_symbol_parser;

vma *symbol_parser::find_vma(pid_t pid, size_t pc) {
  std::map<int, proc_vma>::iterator it;

  it = machine_vma.find(pid);
  if (it == machine_vma.end()) {
    return NULL;
  }

  proc_vma::iterator vma_iter = it->second.upper_bound(pc);
  if (vma_iter == it->second.end() || vma_iter->second.end < pc) {
    return NULL;
  }

  if (vma_iter != it->second.begin()) {
    --vma_iter;
  }

  return &vma_iter->second;
}

bool symbol_parser::find_vma(pid_t pid, vma &vm) {
  std::map<int, proc_vma>::iterator proc_vma_map;

  proc_vma_map = machine_vma.find(pid);
  if (proc_vma_map == machine_vma.end()) {
    return false;
  }

  proc_vma::const_iterator vma_iter = proc_vma_map->second.upper_bound(vm.pc);
  if (vma_iter == proc_vma_map->second.end()) {
    return false;
  }
  if (vma_iter->second.end < vm.pc) {
    return false;
  }

  if (vma_iter != proc_vma_map->second.begin()) {
    --vma_iter;
  }

  vm.start = vma_iter->second.start;
  vm.end = vma_iter->second.end;
  vm.name = vma_iter->second.name;
  vm.offset = vma_iter->second.offset;

  return true;
}

bool symbol_parser::load_pid_maps(int pid) {
  std::map<int, proc_vma>::iterator it;
  it = machine_vma.find(pid);
  if (it != machine_vma.end()) {
    return true;
  }

  proc_vma proc;
  char fn[256];
  sprintf(fn, "/proc/%d/maps", pid);
  FILE *fp = fopen(fn, "r");
  if (!fp) {
    return false;
  }

  char buf[4096];
  char exename[4096];
  size_t start, end, offset;
  while (fgets(buf, sizeof(buf), fp) != NULL) {
    start = end = offset = 0;
    exename[0] = '\0';
    sscanf(buf, "%lx-%lx %*s %lx %*x:%*x %*u %s %*s\n", &start, &end, &offset,
           exename);
    if (exename[0] == '\0') {
      strcpy(exename, "[anon]");
    }
    vma vm(start, end, offset, exename);
    proc.insert(std::make_pair(vm.start, std::move(vm)));
  }

  fclose(fp);

  machine_vma.insert(std::make_pair(pid, std::move(proc)));
  it = machine_vma.find(pid);
  if (it == machine_vma.end()) {
    return false;
  }

  return true;
}

bool symbol_parser::find_symbol_in_cache(int tgid, unsigned long addr,
                                         std::string &symbol) {
  std::map<int, std::map<unsigned long, std::string>>::const_iterator it_pid =
      symbols_cache.find(tgid);

  if (it_pid != symbols_cache.end()) {
    std::map<unsigned long, std::string> map = symbols_cache[tgid];
    std::map<unsigned long, std::string>::const_iterator it_symbol =
        map.find(addr);

    if (it_symbol != map.end()) {
      symbol = map[addr];

      return true;
    }
  }

  return false;
}

bool symbol_parser::get_symbol_info(int pid, symbol &sym, elf_file &file) {
  std::map<int, proc_vma>::iterator proc_vma_info;

  if (java_only) {
    file.type = UNKNOWN;
    return true;
  }

  proc_vma_info = machine_vma.find(pid);
  if (proc_vma_info == machine_vma.end()) {
    if (!load_pid_maps(pid)) {
      return false;
    }
  }

  vma area(sym.ip);
  if (!find_vma(pid, area)) {
    return false;
  }
  if (area.name == "[anon]") {
    file.type = JIT_TYPE;
  }

  file.reset(area.name);
  if (file.type != JIT_TYPE) {
    sym.reset(area.map(sym.ip));
    sym.base = area.start; // add base address
  }

  return true;
}

bool symbol_parser::putin_symbol_cache(int tgid, unsigned long addr,
                                       std::string &symbol) {
  std::map<int, std::map<unsigned long, std::string>>::const_iterator it_pid =
      symbols_cache.find(tgid);

  if (it_pid == symbols_cache.end()) {
    std::map<unsigned long, std::string> map;
    symbols_cache.insert(std::make_pair(tgid, map));
  }

  std::map<unsigned long, std::string> &map = symbols_cache[tgid];
  std::map<unsigned long, std::string>::const_iterator it_symbol =
      map.find(addr);

  if (it_symbol == map.end()) {
    map[addr] = symbol;
    return true;
  }

  return false;
}

bool search_symbol(const std::set<symbol> &ss, symbol &sym) {
  // printf("search symbol2 0x%lx\n", sym.ip);
  std::set<symbol>::const_iterator it = ss.find(sym);

  if (it != ss.end()) {
    sym.end = it->end;
    sym.start = it->start;
    sym.name = it->name;

    return true;
  }

  return false;
}

bool symbol_parser::load_elf(pid_t pid, const elf_file &file,
                             size_t file_base_addr) {
  std::map<elf_file, std::set<symbol>>::iterator it;
  it = file_symbols.find(file);
  std::set<symbol> tmp;
  std::set<symbol> &syms = tmp;
  if (it != file_symbols.end()) {
    return true;
  }
  if (get_symbol_from_elf(syms, file.filename.c_str(), file_base_addr)) {
    // printf("load elf %s\n", file.filename.c_str());
    file_symbols.insert(make_pair(file, std::move(syms)));
    return true;
  }
  return false;
}

bool symbol_parser::find_elf_symbol(symbol &sym, const elf_file &file, int pid,
                                    int pid_ns) {
  std::map<elf_file, std::set<symbol>>::iterator it;
  it = file_symbols.find(file);
  std::set<symbol> ss;
  // printf("search symbol1 0x%lx, base 0x%lx\n", sym.ip, sym.base);
  if (it == file_symbols.end()) {
    if (!load_elf(pid, file, sym.base)) {
      // printf("load elf %s failed\n", file.filename.c_str());
      return false;
    }
    // printf("load elf %s success\n", file.filename.c_str());
    it = file_symbols.find(file);
    return search_symbol(it->second, sym);
  } else {
    return search_symbol(it->second, sym);
  }
  return true;
}

static bool load_kernel_symbol_list(std::vector<std::string> &sym_list) {
  FILE *fp = fopen("/proc/kallsyms", "r");
  if (!fp) {
    return -1;
  }

  char buf[256];
  char type;
  int len;
  while (fgets(buf, sizeof(buf), fp) != NULL) {
    sscanf(buf, "%*p %c %*s\n", &type);
    if ((type | 0x20) != 't') {
      continue;
    }
    len = strlen(buf);
    if (buf[len - 1] == '\n') {
      buf[len - 1] = ' ';
    }
    sym_list.push_back(buf);
  }
  fclose(fp);

  std::sort(sym_list.begin(), sym_list.end());
  return true;
}

bool is_space(int ch) { return std::isspace(ch); }

static inline void rtrim(std::string &s) {
  s.erase(std::find_if(s.rbegin(), s.rend(), is_space).base(), s.end());
}

static bool get_next_kernel_symbol(std::set<symbol> &syms,
                                   std::vector<std::string> &sym_list,
                                   std::vector<std::string>::iterator cursor) {
  if (cursor == sym_list.end()) {
    return false;
  }
  symbol sym;
  size_t start, end;
  sscanf(cursor->c_str(), "%p %*c %*s\n", (void **)&start);
  sym.name = cursor->c_str() + 19;
  rtrim(sym.name);
#if 0
    if (sym.name[sym.name.size()-1] == '\n') {
        sym.name[sym.name.size()-1] = '\0';
    }
#endif
  cursor++;
  if (cursor != sym_list.end()) {
    sscanf(cursor->c_str(), "%p %*c %*s\n", (void **)&end);
  } else {
    end = INVALID_ADDR;
  }
  sym.start = start;
  sym.end = end;
  sym.ip = start;

  syms.insert(sym);
  return true;
}

bool symbol_parser::load_kernel() {
  if (kernel_symbols.size() != 0) {
    return true;
  }

  std::vector<std::string> sym_list;
  if (!load_kernel_symbol_list(sym_list)) {
    exit(0);
    return false;
  }

  std::vector<std::string>::iterator cursor = sym_list.begin();
  while (get_next_kernel_symbol(kernel_symbols, sym_list, cursor)) {
    cursor++;
  }
  return true;
}

bool symbol_parser::find_kernel_symbol(symbol &sym) {
  load_kernel();
  sym.end = sym.start = 0;
  std::set<symbol>::iterator it = kernel_symbols.find(sym);
  if (it != kernel_symbols.end()) {
    sym.end = it->end;
    sym.start = it->start;
    sym.name = it->name;
    return true;
  }
  return false;
}