#include "patcher.h"
#include "debug.h"
#include "nyx/disassembler.h"
#include "nyx/memory_access.h"
#include "nyx/state/state.h"

uint8_t cmp_patch_data[] = {
    0x38, 0xC0, [2 ... MAX_INSTRUCTION_SIZE] = 0x90
}; // CMP AL,AL; NOP, NOP ...
const uint8_t *cmp_patch = &cmp_patch_data[0];

/*
 * Private Helper Functions Declarations
 */
static void _patcher_apply_patch(patcher_t *self, size_t index);

static void _patcher_restore_patch(patcher_t *self, size_t index);

static void _patcher_save_patch(patcher_t *self,
                                size_t     index,
                                uint8_t   *data,
                                size_t     instruction_size,
                                uint64_t   addr);

static size_t _patcher_disassemble_size(patcher_t *self,
                                        uint8_t   *data,
                                        uint64_t   addr,
                                        x86_insn   id);

static void _patcher_alloc_patch_infos(patcher_t *self, size_t num_patches);

static void _patcher_free_patch_infos(patcher_t *self);

static redqueen_t *_redq_ptr(patcher_t *self);


/*
 * Public Functions
 */

patcher_t *patcher_new(CPUState *cpu)
{
    patcher_t *res            = malloc(sizeof(patcher_t));
    res->cpu                  = cpu;
    res->num_patches          = 0;
    res->patches              = NULL;
    res->is_currently_applied = false;
    return res;
}

void patcher_free(patcher_t *self)
{
    assert(!self->is_currently_applied);
    _patcher_free_patch_infos(self);
    free(self);
}

void patcher_apply_all(patcher_t *self)
{
    assert(!self->is_currently_applied);
    assert(!_redq_ptr(self)->hooks_applied);
    // assert(patcher_validate_patches(self));
    for (size_t i = 0; i < self->num_patches; i++) {
        _patcher_apply_patch(self, i);
    }
    self->is_currently_applied = true;
}

void patcher_restore_all(patcher_t *self)
{
    assert(self->is_currently_applied);
    assert(!_redq_ptr(self)->hooks_applied);
    // assert(patcher_validate_patches(self));
    for (size_t i = 0; i < self->num_patches; i++) {
        _patcher_restore_patch(self, i);
    }
    self->is_currently_applied = false;
}

void patcher_set_addrs(patcher_t *self, uint64_t *addrs, size_t num_addrs)
{
    _patcher_free_patch_infos(self);
    _patcher_alloc_patch_infos(self, num_addrs);
    uint8_t curr_instruction_code[MAX_INSTRUCTION_SIZE];
    memset(&curr_instruction_code[0], 0, MAX_INSTRUCTION_SIZE);

    for (size_t i = 0; i < self->num_patches; i++) {
        // nyx_debug_p(REDQUEEN_PREFIX, "patching %lx", addrs[i]);
        if (read_virtual_memory(addrs[i], &curr_instruction_code[0],
                                MAX_INSTRUCTION_SIZE, self->cpu))
        {
            size_t size = _patcher_disassemble_size(self, &curr_instruction_code[0],
                                                    addrs[i], X86_INS_CMP);
            assert(size != 0); // csopen failed, shouldn't happen
            _patcher_save_patch(self, i, &curr_instruction_code[0], size, addrs[i]);
        }
    }
}

static void print_hexdump(const uint8_t *addr, size_t size)
{
    for (size_t i = 0; i < size; i++) {
        printf(" %02x", addr[i]);
    }
    printf("\n");
}

bool patcher_validate_patches(patcher_t *self)
{
    bool was_rq = _redq_ptr(self)->hooks_applied;
    if (was_rq)
        redqueen_remove_hooks(_redq_ptr(self));
    if (!self->patches) {
        return true;
    }
    for (size_t i = 0; i < self->num_patches; i++) {
        uint8_t buf[MAX_INSTRUCTION_SIZE];
        read_virtual_memory(self->patches[i].addr, &buf[0], MAX_INSTRUCTION_SIZE,
                            self->cpu);
        const uint8_t *should_value = NULL;
        if (self->is_currently_applied) {
            should_value = cmp_patch;
        } else {
            should_value = &self->patches[i].orig_bytes[0];
        }

        nyx_debug_p(REDQUEEN_PREFIX, "Validating, mem:");
        print_hexdump(&buf[0], self->patches[i].size);
        nyx_debug_p(REDQUEEN_PREFIX, "should_be:");
        print_hexdump(should_value, self->patches[i].size);
        if (0 != memcmp(&buf[0], should_value, self->patches[i].size)) {
            nyx_debug_p(REDQUEEN_PREFIX,
                        "validating patches failed self->is_currently_applied = %d",
                        self->is_currently_applied);
            return false;
        }
    }
    if (was_rq)
        redqueen_insert_hooks(_redq_ptr(self));
    return true;
}


/*
 * Private Helper Functions Definitions
 */

static void _patcher_apply_patch(patcher_t *self, size_t index)
{
    abort(); // deprecated function -> remove this code later
}

static void _patcher_restore_patch(patcher_t *self, size_t index)
{
    abort(); // deprecated function -> remove this code later
}

static void _patcher_save_patch(patcher_t *self,
                                size_t     index,
                                uint8_t   *data,
                                size_t     instruction_size,
                                uint64_t   addr)
{
    assert(instruction_size >= 2);
    assert(instruction_size < MAX_INSTRUCTION_SIZE);
    patch_info_t *info = &self->patches[index];
    memset(&info->orig_bytes[0], 0, MAX_INSTRUCTION_SIZE);
    memcpy(&info->orig_bytes[0], data, instruction_size);
    info->addr = addr;
    info->size = instruction_size;
}

static size_t _patcher_disassemble_size(patcher_t *self,
                                        uint8_t   *data,
                                        uint64_t   addr,
                                        x86_insn   type)
{
    csh handle;
    if (cs_open(CS_ARCH_X86,
                get_capstone_mode(GET_GLOBAL_STATE()->disassembler_word_width),
                &handle) == CS_ERR_OK)
    {
        cs_insn *insn       = cs_malloc(handle);
        uint8_t *cur_offset = data;
        uint64_t cs_address = addr;
        uint64_t code_size  = MAX_INSTRUCTION_SIZE;
        cs_disasm_iter(handle, (const uint8_t **)&cur_offset, &code_size,
                       &cs_address, insn);
        size_t size = insn->size;
        if (type != X86_INS_INVALID) {
            assert(insn->id == type);
        }
        cs_free(insn, 1);
        cs_close(&handle);
        return size;
    }
    return 0;
}

static void _patcher_alloc_patch_infos(patcher_t *self, size_t num_patches)
{
    assert(self->num_patches == 0);
    assert(self->patches == NULL);
    assert(num_patches < 10000);
    self->num_patches = num_patches;
    self->patches     = malloc(sizeof(patch_info_t) * num_patches);
}

static void _patcher_free_patch_infos(patcher_t *self)
{
    assert(!self->is_currently_applied);
    free(self->patches);
    self->patches     = NULL;
    self->num_patches = 0;
}

static redqueen_t *_redq_ptr(patcher_t *self)
{
    redqueen_t *res = GET_GLOBAL_STATE()->redqueen_state; // self->cpu->redqueen_state;
    return res;
}