// clang++ -std=c++17 -o jit jit.cpp -lcapstone

#include <capstone/capstone.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>

class x86 {
 protected:
  uint64_t size = 0;
  uint64_t cap = 32;
  uint8_t *buf;

 public:
  inline x86() {
    cap = 32;
    buf = new uint8_t[cap];
  }

  template <typename T>
  inline T read(uint64_t i) const {
    return *(T *)(void *)(buf + i);
  }

  template <typename T>
  inline uint64_t write_to(uint64_t i, T val) const {
    *(T *)(void *)(buf + i) = val;
    return size;
  }

  template <typename T>
  inline uint64_t write(T val) {
    if (size + sizeof(val) >= cap - 1) {
      uint8_t *new_code = new uint8_t[cap * 2];
      memcpy(new_code, buf, cap);
      cap *= 2;
      delete buf;
      buf = new_code;
    }
    uint64_t addr = write_to(size, val);
    size += sizeof(T);
    return addr;
  }

  inline void bytes(void) {}
  inline void bytes(unsigned char c) { write<unsigned char>(c); }

  template <typename... Args>
  inline void bytes(unsigned char c, Args const &... args) {
    bytes(c);
    bytes(args...);
  }

  inline uint64_t get_size() { return size; }
  inline uint64_t get_cap() { return cap; }

  inline void copy(void *dest) { memcpy(dest, buf, size); }

  inline int dump(void) {
    csh handle;
    cs_insn *insn;
    size_t count;

    if (cs_open(CS_ARCH_X86, CS_MODE_64, &handle) != CS_ERR_OK) return -1;
    count = cs_disasm(handle, buf, size, 0x1000, 0, &insn);
    if (count > 0) {
      size_t j;
      for (j = 0; j < count; j++) {
        printf("0x%" PRIx64 ":\t%s\t%s\n", insn[j].address, insn[j].mnemonic,
               insn[j].op_str);
      }

      cs_free(insn, count);
    } else
      printf("ERROR: Failed to disassemble given code!\n");

    cs_close(&handle);
    return 1;
  }
};

typedef unsigned char(bytefunc)(void);

int main() {
  x86 code;

  code.bytes(0xB8);
  code.write<int>(42);  // mov eax, 42

  code.bytes(0xC3);     // ret

  code.dump();


  void *func_ptr = mmap(NULL, code.get_size(),
                        PROT_READ | PROT_EXEC | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);

  code.copy(func_ptr);

  auto func = (bytefunc *)func_ptr;

  unsigned char c = func();
  printf("%d\n", c);
}