#include "libafl/tcg.h"
#include "libafl/hooks/tcg/jmp.h"

static TCGHelperInfo libafl_exec_jmp_hook_info = {
    .func = NULL,
    .name = "libafl_exec_jmp_hook",
    .flags = dh_callflag(void),
    .typemask = dh_typemask(void, 0) | dh_typemask(tl, 1) | dh_typemask(tl, 2) | dh_typemask(i64, 3) | dh_typemask(i64, 4),
};

struct libafl_jmp_hook* libafl_jmp_hooks;
size_t libafl_jmp_hooks_num = 0;

size_t libafl_add_jmp_hook(uint64_t (*gen_cb)(uint64_t data, target_ulong src, target_ulong dst),
                          void (*exec_cb)(uint64_t data, target_ulong src, target_ulong dst, uint64_t id),
                          uint64_t data)
{
    struct libafl_jmp_hook* hook = calloc(sizeof(struct libafl_jmp_hook), 1);
    hook->gen_cb = gen_cb;
    hook->exec_cb = exec_cb;
    hook->num = libafl_jmp_hooks_num++;
    hook->data = data;
    hook->next = libafl_jmp_hooks;
    libafl_jmp_hooks = hook;

    memcpy(&hook->helper_info, &libafl_exec_jmp_hook_info, sizeof(TCGHelperInfo));
    hook->helper_info.func = exec_cb;
    return hook->num;
}

void libafl_gen_jmp(target_ulong src, target_ulong dst)
{

    struct libafl_jmp_hook* hook = libafl_jmp_hooks;
    while (hook) {
        uint64_t cur_id = 0;
        if (hook->gen_cb)
            cur_id = hook->gen_cb(hook->data, src, dst);
        if (cur_id != (uint64_t)-1 && hook->exec_cb) {
            TCGv_i64 tmp0 = tcg_constant_i64(hook->data);
            TCGv_i64 tmp1 = tcg_constant_i64(src);
            TCGv_i64 tmp2 = tcg_constant_i64(dst);
            TCGv_i64 tmp3 = tcg_constant_i64(cur_id);
            TCGTemp *tmp4[4] = { tcgv_i64_temp(tmp0), tcgv_i64_temp(tmp1), tcgv_i64_temp(tmp2), tcgv_i64_temp(tmp3) };
            tcg_gen_callN(hook->helper_info.func, &hook->helper_info, NULL, tmp4);
            tcg_temp_free_i64(tmp0);
            tcg_temp_free_i64(tmp1);
            tcg_temp_free_i64(tmp2);
            tcg_temp_free_i64(tmp3);
        }
        hook = hook->next;
    }
}

void libafl_gen_jmp_dynamic(target_ulong src, TCGv_i32 dst)
{
    struct libafl_jmp_hook* hook = libafl_jmp_hooks;
    while (hook) {
        uint64_t cur_id = 0;
        if (hook->gen_cb)
            cur_id = hook->gen_cb(hook->data, src, 0);   // target is not statically known, signal with 0
        if (cur_id != (uint64_t)-1 && hook->exec_cb) {
            TCGv_i64 tmp0 = tcg_constant_i64(hook->data);
            TCGv_i64 tmp1 = tcg_constant_i64(src);
           // TCGv_i32 tmp2 = dst;
            TCGv_i64 tmp3 = tcg_constant_i64(cur_id);
            TCGTemp *tmp4[4] = { tcgv_i64_temp(tmp0), tcgv_i64_temp(tmp1), tcgv_i32_temp(dst), tcgv_i64_temp(tmp3) };
            tcg_gen_callN(hook->helper_info.func, &hook->helper_info, NULL, tmp4);
            tcg_temp_free_i64(tmp0);
            tcg_temp_free_i64(tmp1);
            //tcg_temp_free_i64(tmp2);
            tcg_temp_free_i64(tmp3);
        }
        hook = hook->next;
    }
}

GEN_REMOVE_HOOK(jmp)