#include <decomp/routine.hpp>

namespace theo::decomp {
routine_t::routine_t(coff::symbol_t* sym,
                     coff::image_t* img,
                     coff::section_header_t* scn,
                     std::vector<std::uint8_t>& fn,
                     decomp_type_t dcmp_type)
    : m_img(img), m_scn(scn), m_data(fn), m_dcmp_type(dcmp_type), m_sym(sym) {}

std::vector<decomp::symbol_t> routine_t::decompose() {
  std::vector<decomp::symbol_t> result;

  switch (m_dcmp_type) {
    case none: {
      std::vector<comp::reloc_t> relocs;
      auto scn_relocs = reinterpret_cast<coff::reloc_t*>(
          m_scn->ptr_relocs + reinterpret_cast<std::uint8_t*>(m_img));

      for (auto idx = 0u; idx < m_scn->num_relocs; ++idx) {
        auto scn_reloc = &scn_relocs[idx];
        auto sym_reloc = m_img->get_symbol(scn_relocs[idx].symbol_index);
        auto sym_name = sym_reloc->name.to_string(m_img->get_strings());
        auto sym_hash = decomp::symbol_t::hash(sym_name.data());

        spdlog::info("{} reloc to: {} at offset: {}",
                     m_sym->name.to_string(m_img->get_strings()), sym_name,
                     scn_reloc->virtual_address);

        relocs.push_back(comp::reloc_t(scn_reloc->virtual_address, sym_hash));
      }

      result.push_back(decomp::symbol_t(
          m_sym->name.to_string(m_img->get_strings()).data(), m_sym->value,
          m_data, m_scn, m_sym, relocs, decomp_type_t::none));
      break;
    }
    case instr_split: {
      std::uint32_t offset = 0u;
      xed_error_enum_t err;

      xed_decoded_inst_t instr;
      xed_state_t istate{XED_MACHINE_MODE_LONG_64, XED_ADDRESS_WIDTH_64b};
      xed_decoded_inst_zero_set_mode(&instr, &istate);

      // keep looping over the section, lower the number of bytes each time...
      //
      while ((err = xed_decode(&instr, m_data.data() + offset,
                               m_data.size() - offset)) == XED_ERROR_NONE) {
        // symbol name is of the format: symbol@instroffset, I.E: main@11...
        //
        auto new_sym_name =
            std::string(m_sym->name.to_string(m_img->get_strings()));

        // first instruction doesnt need the @offset...
        //
        if (offset)
          new_sym_name.append("@").append(std::to_string(offset));

        std::vector<comp::reloc_t> relocs;
        auto scn_relocs = reinterpret_cast<coff::reloc_t*>(
            m_scn->ptr_relocs + reinterpret_cast<std::uint8_t*>(m_img));

        // find if this instruction has a relocation or not...
        // if so, return the reloc_t...
        //
        auto reloc = std::find_if(
            scn_relocs, scn_relocs + m_scn->num_relocs,
            [&](coff::reloc_t reloc) {
              return reloc.virtual_address >= offset &&
                     reloc.virtual_address <
                         offset + xed_decoded_inst_get_length(&instr);
            });

        // if there is indeed a reloc for this instruction...
        //
        if (reloc != scn_relocs + m_scn->num_relocs) {
          auto sym_reloc = m_img->get_symbol(reloc->symbol_index);
          auto sym_name = sym_reloc->name.to_string(m_img->get_strings());
          auto sym_hash = decomp::symbol_t::hash(sym_name.data());
          auto reloc_offset = reloc->virtual_address - offset;

          spdlog::info("{} reloc to: {} at offset: {}", new_sym_name, sym_name,
                       reloc_offset);

          relocs.push_back(comp::reloc_t(reloc_offset, sym_hash));
        }

        // add a reloc to the next instruction...
        // note that the offset is ZERO... comp_t will understand that
        // relocs with offset ZERO means the next instructions...
        //
        auto next_inst_sym =
            std::string(m_sym->name.to_string(m_img->get_strings()))
                .append("@")
                .append(std::to_string(offset +
                                       xed_decoded_inst_get_length(&instr)));

        relocs.push_back(
            comp::reloc_t(0, decomp::symbol_t::hash(next_inst_sym)));

        // get the instructions bytes
        //
        std::vector<std::uint8_t> inst_bytes(
            m_data.data() + offset,
            m_data.data() + offset + xed_decoded_inst_get_length(&instr));

        result.push_back(decomp::symbol_t(new_sym_name, offset, inst_bytes,
                                          m_scn, m_sym, relocs,
                                          decomp_type_t::instr_split));

        // after creating the symbol and dealing with relocs then print the
        // information we have concluded...
        //
        char buff[255];
        offset += xed_decoded_inst_get_length(&instr);
        xed_format_context(XED_SYNTAX_INTEL, &instr, buff, sizeof buff, NULL,
                           NULL, NULL);

        spdlog::info("{}: {}", new_sym_name, buff);
        // need to set this so that instr can be used to decode again...
        xed_decoded_inst_zero_set_mode(&instr, &istate);
      }

      // remove the relocation to the next symbol from the last instruction
      //
      auto last_inst = result.back();
      auto last_inst_relocs = last_inst.relocs();
      last_inst_relocs.erase(last_inst_relocs.end() - 1);
      break;
    }
    default:
      break;
  }

  return result;
}

coff::section_header_t* routine_t::scn() {
  return m_scn;
}

std::vector<std::uint8_t> routine_t::data() {
  return m_data;
}
}  // namespace theo::decomp