#pragma once
#include <spdlog/spdlog.h>
#include <functional>
#include <map>
#include <random>

#define XED_ENCODER
extern "C" {
#include <xed-decode.h>
#include <xed-interface.h>
}

namespace theo::obf::transform {
using transform_t = std::function<std::size_t(std::size_t, std::size_t)>;

class operation_t {
 public:
  explicit operation_t(transform_t op, xed_iclass_enum_t type)
      : m_transform(op), m_type(type) {}

  std::vector<std::uint8_t> native(const xed_decoded_inst_t* inst,
                                   std::size_t imm) {
    xed_error_enum_t err;
    xed_encoder_request_t req;
    xed_state_t istate{XED_MACHINE_MODE_LONG_64, XED_ADDRESS_WIDTH_64b};

    xed_encoder_request_zero_set_mode(&req, &istate);
    xed_encoder_request_set_effective_operand_width(&req, 64);
    xed_encoder_request_set_iclass(&req, m_type);
    xed_encoder_request_set_reg(&req, XED_OPERAND_REG0, XED_REG_RAX);
    xed_encoder_request_set_operand_order(&req, 0, XED_OPERAND_REG0);
    xed_encoder_request_set_uimm0(&req, imm, 4);
    xed_encoder_request_set_operand_order(&req, 1, XED_OPERAND_IMM0);

    std::uint8_t inst_buff[XED_MAX_INSTRUCTION_BYTES];
    std::uint32_t inst_len = {};

    if ((err = xed_encode(&req, inst_buff, sizeof(inst_buff), &inst_len)) !=
        XED_ERROR_NONE) {
      spdlog::error("failed to encode instruction... reason: {}",
                    xed_error_enum_t2str(err));

      assert(err == XED_ERROR_NONE);
    }

    return std::vector<std::uint8_t>(inst_buff, inst_buff + inst_len);
  }

  xed_iclass_enum_t inverse() { return m_inverse_op[m_type]; }
  transform_t* get_transform() { return &m_transform; }

  static std::size_t random(std::size_t lowest, std::size_t largest) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<std::size_t> distr(lowest, largest);
    return distr(gen);
  }

 private:
  transform_t m_transform;
  xed_iclass_enum_t m_type;

  std::map<xed_iclass_enum_t, xed_iclass_enum_t> m_inverse_op = {
      {XED_ICLASS_ADD, XED_ICLASS_SUB}, {XED_ICLASS_SUB, XED_ICLASS_ADD},
      {XED_ICLASS_AND, XED_ICLASS_OR},  {XED_ICLASS_OR, XED_ICLASS_AND},
      {XED_ICLASS_ROL, XED_ICLASS_ROR}, {XED_ICLASS_ROR, XED_ICLASS_ROL},
      {XED_ICLASS_XOR, XED_ICLASS_XOR}};
};

}  // namespace theo::obf::transform