#include "elf.h"

#include "minicc.h"

#include <libreichwein/file.h>

#include <iostream>

// TODO
//#include <boost/endian/conversion.hpp>

#include <cstring>

#define PAD

// Helper Functions
namespace {
 // On amd64, we need at least 4096 bytes aligment. Otherwise we get segfaults at startup
 const size_t PAGE_SIZE = 0x1000;

 // Helper function: returns size with full pages: 0, 4096, ...
 size_t paged_size(const std::vector<uint8_t>& data)
 {
  return (data.size() + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);
 }

 void AddFileHeader(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
  Elf::FileHeader fh;

#ifdef PAD
  fh.e_shoff = PAGE_SIZE + paged_size(code) + data.size() + 23;
#else
  fh.e_shoff = sizeof(fh) + sizeof(Elf::ProgramHeader) * 2 + code.size() + data.size() + 23;
#endif

  size_t old_size {elf.size()};
  elf.resize(old_size + sizeof(fh));
  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&fh), sizeof(fh));
 }

 void AddProgramHeaderCode(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
  Elf::ProgramHeader ph;

#ifdef PAD
  ph.p_offset = PAGE_SIZE;
  ph.p_align = PAGE_SIZE;
#else
  ph.p_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2;
#endif
  ph.p_filesz = code.size();
  ph.p_memsz = code.size();

  size_t old_size {elf.size()};
  elf.resize(old_size + sizeof(ph));
  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&ph), sizeof(ph));
 }

 void AddProgramHeaderData(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
  Elf::ProgramHeader ph;

#ifdef PAD
  ph.p_offset = PAGE_SIZE + paged_size(code);
  ph.p_align = PAGE_SIZE;
#else
  ph.p_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2 + code.size();
#endif
  ph.p_filesz = data.size();
  ph.p_memsz = data.size();

  ph.p_vaddr = 0x401000 + paged_size(code);
  ph.p_paddr = 0x401000 + paged_size(code);
  ph.p_flags = 6; // RW

  size_t old_size {elf.size()};
  elf.resize(old_size + sizeof(ph));
  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&ph), sizeof(ph));
 }

 void AddSectionHeaderNull(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
  Elf::SectionHeader sh;

  sh.sh_name = 0; // offset in section
  sh.sh_size = 0;
  sh.sh_type = 0; // <null>
  sh.sh_flags = 0;
  sh.sh_addr = 0;
  sh.sh_offset = 0;

  size_t old_size {elf.size()};
  elf.resize(old_size + sizeof(sh));
  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
 }
 
 void AddSectionHeaderText(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
  Elf::SectionHeader sh;

  sh.sh_name = 1; // offset in section
  sh.sh_size = code.size();
  sh.sh_type = 1; // program
#ifdef PAD
  sh.sh_offset = PAGE_SIZE;
#else
  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2;
#endif
  sh.sh_addralign = 1;

  size_t old_size {elf.size()};
  elf.resize(old_size + sizeof(sh));
  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
 }
 
 void AddSectionHeaderData(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
  Elf::SectionHeader sh;

  sh.sh_name = 7; // offset in section
  sh.sh_size = data.size();
  sh.sh_type = 1; // program (also data)
#ifdef PAD
  sh.sh_offset = PAGE_SIZE + paged_size(code);
#else
  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2 + code.size();
#endif
  sh.sh_addralign = 1;

  size_t old_size {elf.size()};
  elf.resize(old_size + sizeof(sh));
  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
 }
 
 void AddSectionHeaderSectionNames(std::vector<uint8_t>& elf, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
 {
  Elf::SectionHeader sh;

  sh.sh_name = 13; // offset in section
  sh.sh_size = 23;
  sh.sh_type = 3; // section names
  sh.sh_flags = 0;
  sh.sh_addr = 0;
#ifdef PAD
  sh.sh_offset = PAGE_SIZE + paged_size(code) + data.size();
#else
  sh.sh_offset = sizeof(Elf::FileHeader) + sizeof(Elf::ProgramHeader) * 2 + code.size() + data.size();
#endif
  sh.sh_addralign = 1;

  size_t old_size {elf.size()};
  elf.resize(old_size + sizeof(sh));
  std::memcpy(elf.data() + old_size, reinterpret_cast<void*>(&sh), sizeof(sh));
 }

 void PadUpTo(std::vector<uint8_t>& elf, size_t size)
 {
  if (elf.size() < size) {
   elf.resize(size);
  } else if (elf.size() > size)
   throw std::runtime_error("Padding not possible. Too many bytes already. ("s + std::to_string(elf.size()) + " > "s + std::to_string(size) + ")"s);
 }
}

void Elf::Write(const std::filesystem::path& path, const std::vector<uint8_t>& code, const std::vector<uint8_t>& data)
{
 std::vector<uint8_t> elf;

 AddFileHeader(elf, code, data);
 AddProgramHeaderCode(elf, code, data);
 AddProgramHeaderData(elf, code, data);

#ifdef PAD
 PadUpTo(elf, PAGE_SIZE);
#endif

 elf.insert(elf.end(), code.begin(), code.end());

#ifdef PAD
 PadUpTo(elf, PAGE_SIZE + paged_size(code));
#endif
 
 elf.insert(elf.end(), data.begin(), data.end());

 std::string section_names("\0.text\0.data\0.shstrtab\0", 23);
 if (section_names.size() != 23)
  throw std::runtime_error("Bad size");

 size_t old_size {elf.size()};
 elf.resize(old_size + section_names.size());
 std::memcpy(elf.data() + old_size, section_names.data(), section_names.size());

 AddSectionHeaderNull(elf, code, data);
 AddSectionHeaderText(elf, code, data);
 AddSectionHeaderData(elf, code, data);
 AddSectionHeaderSectionNames(elf, code, data);

 Reichwein::File::setFile(path, elf);

 fs::permissions(path, fs::perms::owner_exec | fs::perms::group_exec | fs::perms::others_exec, fs::perm_options::add);
}

std::vector<uint8_t> Elf::Read(const std::filesystem::path& path)
{
 std::vector<uint8_t> result;

 // TODO

 return result;
}