#include "ELFLoader.h"
#include "LogManager.h"
#include <cstring>
#include <elf.h>
#include <fstream>
#include <stdint.h>
#include <vector>

namespace ELFLoader {

ELFContainer::ELFContainer(std::string const &Filename) {
  std::fstream ELFFile;
  size_t FileSize{0};
  ELFFile.open(Filename, std::fstream::in | std::fstream::binary);
  LogMan::Throw::A(ELFFile.is_open(), "Failed to open file");

  ELFFile.seekg(0, ELFFile.end);
  FileSize = ELFFile.tellg();
  ELFFile.seekg(0, ELFFile.beg);

  RawFile.resize(FileSize);

  ELFFile.read(&RawFile.at(0), FileSize);

  ELFFile.close();

  memcpy(&Header, reinterpret_cast<Elf64_Ehdr *>(&RawFile.at(0)),
         sizeof(Elf64_Ehdr));
  LogMan::Throw::A(Header.e_phentsize == 56, "PH Entry size wasn't 56");
  LogMan::Throw::A(Header.e_shentsize == 64, "PH Entry size wasn't 64");

  SectionHeaders.resize(Header.e_shnum);
  ProgramHeaders.resize(Header.e_phnum);

  Elf64_Shdr *RawShdrs =
      reinterpret_cast<Elf64_Shdr *>(&RawFile.at(Header.e_shoff));
  Elf64_Phdr *RawPhdrs =
      reinterpret_cast<Elf64_Phdr *>(&RawFile.at(Header.e_phoff));

  for (uint32_t i = 0; i < Header.e_shnum; ++i) {
    SectionHeaders[i] = &RawShdrs[i];
  }

  for (uint32_t i = 0; i < Header.e_phnum; ++i) {
    ProgramHeaders[i] = &RawPhdrs[i];
  }

  CalculateMemoryLayouts();
  CalculateSymbols();

  // Print Information
  PrintHeader();
  //PrintSectionHeaders();
  //PrintProgramHeaders();
  //PrintSymbolTable();

  LogMan::Throw::A(IsStatic, "Can only handle static programs");
}

ELFContainer::~ELFContainer() {
  SymbolMapByAddress.clear();
  SymbolMap.clear();
  Symbols.clear();
  ProgramHeaders.clear();
  SectionHeaders.clear();
  RawFile.clear();
}

void ELFContainer::WriteLoadableSections(MemoryWriter Writer) {
  for (uint32_t i = 0; i < ProgramHeaders.size(); ++i) {
    Elf64_Phdr const *hdr = ProgramHeaders.at(i);
    if (hdr->p_type != PT_LOAD)
      continue;
    Writer(&RawFile.at(hdr->p_offset), hdr->p_paddr, hdr->p_filesz);
  }
}

ELFSymbol const *ELFContainer::GetSymbol(char const *Name) {
  auto Sym = SymbolMap.find(Name);
  if (Sym == SymbolMap.end())
    return nullptr;
  return Sym->second;
}
ELFSymbol const *ELFContainer::GetSymbol(uint64_t Address) {
  auto Sym = SymbolMapByAddress.find(Address);
  if (Sym == SymbolMapByAddress.end())
    return nullptr;
  return Sym->second;
}
ELFSymbol const *ELFContainer::GetSymbolInRange(RangeType Address) {
  auto Sym = SymbolMapByAddress.upper_bound(Address.first);
  if (Sym != SymbolMapByAddress.begin())
    --Sym;
  if (Sym == SymbolMapByAddress.end())
    return nullptr;

  if ((Sym->second->Address + Sym->second->Size) < Address.first)
    return nullptr;

  return Sym->second;
}

void ELFContainer::CalculateMemoryLayouts() {
  uint64_t MinPhysAddr = ~0ULL;
  uint64_t MaxPhysAddr = 0;
  uint64_t PhysMemSize = 0;
  for (uint32_t i = 0; i < ProgramHeaders.size(); ++i) {
    Elf64_Phdr const *hdr = ProgramHeaders.at(i);
    MinPhysAddr = std::min(MinPhysAddr, hdr->p_paddr);
    MaxPhysAddr = std::max(MaxPhysAddr, hdr->p_paddr + hdr->p_memsz);

    if (hdr->p_type == PT_INTERP) {
      IsStatic = false;
    }
  }

  PhysMemSize = MaxPhysAddr - MinPhysAddr;

  MinPhysicalMemoryLocation = MinPhysAddr;
  MaxPhysicalMemoryLocation = MaxPhysAddr;
  PhysicalMemorySize = PhysMemSize;
  LogMan::Msg::D("Min PAddr: 0x%lx", MinPhysAddr);
  LogMan::Msg::D("Max PAddr: 0x%lx", MaxPhysAddr);
  LogMan::Msg::D("Physical Size: 0x%lx", PhysMemSize);
}

void ELFContainer::CalculateSymbols() {
  // Find the symbol table
  Elf64_Shdr const *SymTabHeader{nullptr};
  Elf64_Shdr const *StringTableHeader{nullptr};
  char const *StrTab{nullptr};
  for (uint32_t i = 0; i < SectionHeaders.size(); ++i) {
    Elf64_Shdr const *hdr = SectionHeaders.at(i);
    if (hdr->sh_type == SHT_SYMTAB) {
      SymTabHeader = hdr;
      break;
    }
  }
  if (!SymTabHeader) {
    LogMan::Msg::I("No Symbol table");
    return;
  }

  LogMan::Throw::A(SymTabHeader->sh_link < SectionHeaders.size(),
                   "Symbol table string table section is wrong");
  LogMan::Throw::A(SymTabHeader->sh_entsize == sizeof(Elf64_Sym),
                   "Entry size doesn't match symbol entry");

  StringTableHeader = SectionHeaders.at(SymTabHeader->sh_link);
  StrTab = &RawFile.at(StringTableHeader->sh_offset);

  uint64_t NumSymbols = SymTabHeader->sh_size / SymTabHeader->sh_entsize;
  Symbols.resize(NumSymbols);
  for (uint64_t i = 0; i < NumSymbols; ++i) {
    uint64_t offset = SymTabHeader->sh_offset + i * SymTabHeader->sh_entsize;
    Elf64_Sym const *Symbol =
        reinterpret_cast<Elf64_Sym const *>(&RawFile.at(offset));
    if (ELF64_ST_VISIBILITY(Symbol->st_other) != STV_HIDDEN) {
      ELFSymbol *DefinedSymbol = &Symbols.at(i);
      DefinedSymbol->FileOffset = offset;
      DefinedSymbol->Address = Symbol->st_value;
      DefinedSymbol->Size = Symbol->st_size;
      DefinedSymbol->Type = ELF64_ST_TYPE(Symbol->st_info);
      DefinedSymbol->Bind = ELF64_ST_BIND(Symbol->st_info);
      DefinedSymbol->Name = &StrTab[Symbol->st_name];

      SymbolMap[DefinedSymbol->Name] = DefinedSymbol;
      SymbolMapByAddress[DefinedSymbol->Address] = DefinedSymbol;
    }
  }
}

void ELFContainer::PrintHeader() const {
  LogMan::Msg::I("Type: %d", Header.e_type);
  LogMan::Msg::I("Machine: %d", Header.e_machine);
  LogMan::Msg::I("Version: %d", Header.e_version);
  LogMan::Msg::I("Entry point: 0x%lx", Header.e_entry);
  LogMan::Msg::I("PH Off: %d", Header.e_phoff);
  LogMan::Msg::I("SH Off: %d", Header.e_shoff);
  LogMan::Msg::I("Flags: %d", Header.e_flags);
  LogMan::Msg::I("EH Size: %d", Header.e_ehsize);
  LogMan::Msg::I("PH Num: %d", Header.e_phnum);
  LogMan::Msg::I("SH Num: %d", Header.e_shnum);
  LogMan::Msg::I("PH Entry Size: %d", Header.e_phentsize);
  LogMan::Msg::I("SH Entry Size: %d", Header.e_shentsize);
  LogMan::Msg::I("SH Str Index: %d", Header.e_shstrndx);
}

void ELFContainer::PrintSectionHeaders() const {
  LogMan::Throw::A(Header.e_shstrndx < SectionHeaders.size(),
                   "String index section is wrong index!");
  Elf64_Shdr const *StrHeader = SectionHeaders.at(Header.e_shstrndx);
  char const *SHStrings = &RawFile.at(StrHeader->sh_offset);
  for (uint32_t i = 0; i < SectionHeaders.size(); ++i) {
    Elf64_Shdr const *hdr = SectionHeaders.at(i);
    LogMan::Msg::I("Index: %d", i);
    LogMan::Msg::I("Name:       %s", &SHStrings[hdr->sh_name]);
    LogMan::Msg::I("Type:       %d", hdr->sh_type);
    LogMan::Msg::I("Flags:      %d", hdr->sh_flags);
    LogMan::Msg::I("Addr:       0x%lx", hdr->sh_addr);
    LogMan::Msg::I("Offset:     0x%lx", hdr->sh_offset);
    LogMan::Msg::I("Size:       %d", hdr->sh_size);
    LogMan::Msg::I("Link:       %d", hdr->sh_link);
    LogMan::Msg::I("Info:       %d", hdr->sh_info);
    LogMan::Msg::I("AddrAlign:  %d", hdr->sh_addralign);
    LogMan::Msg::I("Entry Size: %d", hdr->sh_entsize);
  }
}

void ELFContainer::PrintProgramHeaders() const {
  LogMan::Throw::A(Header.e_shstrndx < SectionHeaders.size(),
                   "String index section is wrong index!");
  Elf64_Shdr const *StrHeader = SectionHeaders.at(Header.e_shstrndx);
  char const *SHStrings = &RawFile.at(StrHeader->sh_offset);
  for (uint32_t i = 0; i < ProgramHeaders.size(); ++i) {
    Elf64_Phdr const *hdr = ProgramHeaders.at(i);
    LogMan::Msg::I("Type:    %d", hdr->p_type);
    LogMan::Msg::I("Flags:   %d", hdr->p_flags);
    LogMan::Msg::I("Offset:  %d", hdr->p_offset);
    LogMan::Msg::I("VAddr:   0x%lx", hdr->p_vaddr);
    LogMan::Msg::I("PAddr:   0x%lx", hdr->p_paddr);
    LogMan::Msg::I("FSize:   %d", hdr->p_filesz);
    LogMan::Msg::I("MemSize: %d", hdr->p_memsz);
    LogMan::Msg::I("Align:   %d", hdr->p_align);
  }
}

void ELFContainer::PrintSymbolTable() const {
  // Find the symbol table
  Elf64_Shdr const *SymTabHeader{nullptr};
  Elf64_Shdr const *StringTableHeader{nullptr};
  char const *StrTab{nullptr};
  for (uint32_t i = 0; i < SectionHeaders.size(); ++i) {
    Elf64_Shdr const *hdr = SectionHeaders.at(i);
    if (hdr->sh_type == SHT_SYMTAB) {
      SymTabHeader = hdr;
      break;
    }
  }
  if (!SymTabHeader) {
    LogMan::Msg::I("No Symbol table");
    return;
  }

  LogMan::Throw::A(SymTabHeader->sh_link < SectionHeaders.size(),
                   "Symbol table string table section is wrong");
  LogMan::Throw::A(SymTabHeader->sh_entsize == sizeof(Elf64_Sym),
                   "Entry size doesn't match symbol entry");

  StringTableHeader = SectionHeaders.at(SymTabHeader->sh_link);
  StrTab = &RawFile.at(StringTableHeader->sh_offset);

  uint64_t NumSymbols = SymTabHeader->sh_size / SymTabHeader->sh_entsize;
  for (uint64_t i = 0; i < NumSymbols; ++i) {
    uint64_t offset = SymTabHeader->sh_offset + i * SymTabHeader->sh_entsize;
    Elf64_Sym const *Symbol =
        reinterpret_cast<Elf64_Sym const *>(&RawFile.at(offset));
    std::ostringstream Str{};
    Str << i << " : " << std::hex << Symbol->st_value << std::dec << " "
        << Symbol->st_size << " " << uint32_t(Symbol->st_info) << " "
        << uint32_t(Symbol->st_other) << " " << Symbol->st_shndx << " "
        << &StrTab[Symbol->st_name];
    LogMan::Msg::I("%s", Str.str().c_str());
  }
}

} // namespace ELFLoader