#include "../executor.h"
#include "../../include/std/alloc.h"
#include "../../include/std/buffer.h"
#include "asm.h"
#include <assert.h>

/*
    TODO: Currently almost all operations are performed on memory. This should be optimized
    to take advantage of registers.

    TODO: Operations with memory registers could access outside the stack. Should this be checked?
*/

typedef struct {
    u32 asm_index;
    u16 func_index;
} CallDefer;

typedef struct {
    Asm asm;
    usize *function_indices;
    u16 num_functions;
    u16 func_counter;
    Buffer/*CallDefer*/ call_defer;
} amal_executor_impl;

#define IMPL \
    amal_executor_impl *impl; \
    impl = (amal_executor_impl*)self;

#define get_register_stack_offset(reg) -(i32)(reg * (int)sizeof(usize) + (int)sizeof(usize))

static i64 abs_i64(i64 value) {
    return value >= 0 ? value : -value;
}

int amal_executor_init(amal_executor **self) {
    amal_executor_impl **impl;
    impl = (amal_executor_impl**)self;
    *impl = NULL;
    return_if_error(am_malloc(sizeof(amal_executor_impl), (void**)impl));
    (*impl)->function_indices = NULL;
    (*impl)->num_functions = 0;
    (*impl)->func_counter = 0;
    ignore_result_int(buffer_init(&(*impl)->call_defer, NULL));
    return asm_init(&(*impl)->asm);
}

void amal_executor_deinit(amal_executor *self) {
    IMPL
    buffer_deinit(&impl->call_defer);
    am_free(impl->function_indices);
    asm_deinit(&impl->asm);
    am_free(impl);
}

int amal_executor_run(amal_executor *self) {
    IMPL
    return asm_execute(&impl->asm);
}

int amal_executor_instructions_start(amal_executor *self, u16 num_functions) {
    void *new_data;
    IMPL
    return_if_error(am_realloc(impl->function_indices, num_functions * sizeof(*impl->function_indices), &new_data));
    impl->function_indices = new_data;
    impl->num_functions = num_functions;
    impl->func_counter = 0;
    buffer_clear(&impl->call_defer);
    return 0;
}

int amal_executor_instructions_end(amal_executor *self) {
    CallDefer *call_defer, *call_defer_end;
    IMPL

    call_defer = buffer_begin(&impl->call_defer);
    call_defer_end = buffer_end(&impl->call_defer);
    for(; call_defer != call_defer_end; ++call_defer) {
        const isize func_offset = (isize)impl->function_indices[call_defer->func_index] - (isize)call_defer->asm_index;
        asm_override_call_rel32(&impl->asm, call_defer->asm_index, func_offset);
    }
    return 0;
}

int amal_exec_nop(amal_executor *self) {
    IMPL
    return asm_nop(&impl->asm);
}

int amal_exec_setz(amal_executor *self, u8 dst_reg) {
    AsmPtr dst;
    IMPL
    asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
    return asm_mov_mi(&impl->asm, &dst, 0);
}

int amal_exec_mov(amal_executor *self, u8 dst_reg, u8 src_reg) {
    AsmPtr ptr;
    IMPL

    asm_ptr_init_disp(&ptr, RBP, get_register_stack_offset(src_reg));
    return_if_error(asm_mov_rm(&impl->asm, RAX, &ptr));

    asm_ptr_init_disp(&ptr, RBP, get_register_stack_offset(dst_reg));
    return asm_mov_mr(&impl->asm, &ptr, RAX);
}

int amal_exec_movi(amal_executor *self, u8 dst_reg, i64 imm) {
    IMPL
    /* TODO: if @number is a float then use float instructions */
    if(abs_i64(imm) <= INT32_MAX) {
        AsmPtr dst;
        asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
        return_if_error(asm_mov_mi(&impl->asm, &dst, imm));
    } else {
        AsmPtr dst;
        asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
        return_if_error(asm_mov_ri(&impl->asm, RAX, imm));
        return_if_error(asm_mov_mr(&impl->asm, &dst, RAX));
    }
    return 0;
}

int amal_exec_movd(amal_executor *self, u8 dst_reg, BufferView data) {
    AsmPtr dst;
    IMPL

    asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
    return_if_error(asm_mov_ri(&impl->asm, RAX, (uintptr_t)data.data));
    return asm_mov_mr(&impl->asm, &dst, RAX);
}

int amal_exec_add(amal_executor *self, u8 dst_reg, u8 src_reg1, u8 src_reg2) {
    AsmPtr dst;
    AsmPtr reg1;
    AsmPtr reg2;
    IMPL

    asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
    asm_ptr_init_disp(&reg1, RBP, get_register_stack_offset(src_reg1));
    asm_ptr_init_disp(&reg2, RBP, get_register_stack_offset(src_reg2));

    return_if_error(asm_mov_rm(&impl->asm, RAX, &reg1));
    return_if_error(asm_mov_rm(&impl->asm, RCX, &reg2));
    return_if_error(asm_add_rr(&impl->asm, RAX, RCX));
    return asm_mov_mr(&impl->asm, &dst, RAX);
}

int amal_exec_sub(amal_executor *self, u8 dst_reg, u8 src_reg1, u8 src_reg2) {
    AsmPtr dst;
    AsmPtr reg1;
    AsmPtr reg2;
    IMPL

    asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
    asm_ptr_init_disp(&reg1, RBP, get_register_stack_offset(src_reg1));
    asm_ptr_init_disp(&reg2, RBP, get_register_stack_offset(src_reg2));

    return_if_error(asm_mov_rm(&impl->asm, RAX, &reg1));
    return_if_error(asm_mov_rm(&impl->asm, RCX, &reg2));
    return_if_error(asm_sub_rr(&impl->asm, RAX, RCX));
    return asm_mov_mr(&impl->asm, &dst, RAX);
}

int amal_exec_imul(amal_executor *self, u8 dst_reg, u8 src_reg1, u8 src_reg2) {
    AsmPtr dst;
    AsmPtr reg1;
    AsmPtr reg2;
    IMPL

    asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
    asm_ptr_init_disp(&reg1, RBP, get_register_stack_offset(src_reg1));
    asm_ptr_init_disp(&reg2, RBP, get_register_stack_offset(src_reg2));

    return_if_error(asm_mov_rm(&impl->asm, RAX, &reg1));
    return_if_error(asm_mov_rm(&impl->asm, RCX, &reg2));
    return_if_error(asm_imul_rr(&impl->asm, RAX, RCX));
    return asm_mov_mr(&impl->asm, &dst, RAX);
}

int amal_exec_mul(amal_executor *self, u8 dst_reg, u8 src_reg1, u8 src_reg2) {
    (void)self;
    (void)dst_reg;
    (void)src_reg1;
    (void)src_reg2;
    /* TODO: Implement! */
    #if 0
    AsmPtr dst;
    AsmPtr reg1;
    AsmPtr reg2;

    asm_ptr_init_disp(&dst, RBP, -(i32)get_register_at_offset(0));
    asm_ptr_init_disp(&reg1, RBP, -(i32)get_register_at_offset(1));
    asm_ptr_init_disp(&reg2, RBP, -(i32)get_register_at_offset(2));

    return_if_error(asm_mov_rm(&self->asm, RAX, &reg1));
    return_if_error(asm_mov_rm(&self->asm, RCX, &reg2));
    return_if_error(asm_mul_rr(&self->asm, RAX, RCX));
    return_if_error(asm_mov_mr(&self->asm, &dst, RAX));
    #endif
    assert(bool_false && "TODO: Implement!");
    return 0;
}

int amal_exec_idiv(amal_executor *self, u8 dst_reg, u8 src_reg1, u8 src_reg2) {
    AsmPtr dst;
    AsmPtr reg1;
    AsmPtr reg2;
    IMPL

    asm_ptr_init_disp(&dst, RBP, get_register_stack_offset(dst_reg));
    asm_ptr_init_disp(&reg1, RBP, get_register_stack_offset(src_reg1));
    asm_ptr_init_disp(&reg2, RBP, get_register_stack_offset(src_reg2));

    return_if_error(asm_mov_rm(&impl->asm, RAX, &reg1));
    return_if_error(asm_mov_rm(&impl->asm, RCX, &reg2));
    return_if_error(asm_cqo(&impl->asm));
    return_if_error(asm_idiv_rr(&impl->asm, RCX));
    return asm_mov_mr(&impl->asm, &dst, RAX);
}

int amal_exec_div(amal_executor *self, u8 dst_reg, u8 src_reg1, u8 src_reg2) {
    (void)self;
    (void)dst_reg;
    (void)src_reg1;
    (void)src_reg2;
    /* TODO: Implement! */
    assert(bool_false && "TODO: Implement!");
    return 0;
}

int amal_exec_push(amal_executor *self, u8 reg) {
    AsmPtr reg_ptr;
    IMPL

    asm_ptr_init_disp(&reg_ptr, RBP, get_register_stack_offset(reg));
    return_if_error(asm_mov_rm(&impl->asm, RAX, &reg_ptr));
    return asm_pushr(&impl->asm, RAX);
}

int amal_exec_pushi(amal_executor *self, i64 imm) {
    (void)self;
    (void)imm;
    /* TODO: Implement! */
    assert(bool_false && "TODO: Implement!");
    return 0;
}

int amal_exec_pushd(amal_executor *self, BufferView data) {
    (void)self;
    (void)data;
    /* TODO: Implement! */
    assert(bool_false && "TODO: Implement!");
    return 0;
}

int amal_exec_call(amal_executor *self, u16 func_index, u8 num_args) {
    isize asm_size;
    IMPL
    /* TODO: Preserve necessary registers before call? */
    /* TODO: This assumes all arguments are isize */
    asm_size = asm_get_size(&impl->asm);
    if(func_index < impl->func_counter) {
        return_if_error(asm_call_rel32(&impl->asm, (isize)impl->function_indices[func_index] - asm_size));
    } else {
        /*
            The location of the function has not been defined yet. Use call instruction with dummy data and change
            the location once the location to the function is known
        */
        CallDefer call_defer;
        call_defer.asm_index = asm_size;
        call_defer.func_index = func_index;
        return_if_error(buffer_append(&impl->call_defer, &call_defer, sizeof(call_defer)));
        return_if_error(asm_call_rel32(&impl->asm, 0));
    }

    if(num_args > 0)
        return asm_add_rm64_imm(&impl->asm, RSP, num_args * sizeof(isize));
    return 0;
}

/*
int amal_exec_callr(u8 dst_reg, BufferView data) {

}
*/

int amal_exec_cmp(amal_executor *self, u8 dst_reg, u8 src_reg1, u8 src_reg2) {
    (void)self;
    (void)dst_reg;
    (void)src_reg1;
    (void)src_reg2;
    /* TODO: Implement! */
    assert(bool_false && "TODO: Implement!");
    return 0;
}

int amal_exec_jz(amal_executor *self, u8 dst_reg, i16 offset) {
    (void)self;
    (void)dst_reg;
    (void)offset;
    /* TODO: Implement! */
    assert(bool_false && "TODO: Implement!");
    return 0;
}

int amal_exec_jmp(amal_executor *self, i16 offset) {
    (void)self;
    (void)offset;
    /* TODO: Implement! */
    assert(bool_false && "TODO: Implement!");
    return 0;
}

int amal_exec_ret(amal_executor *self, u8 reg) {
    AsmPtr ret_reg;
    IMPL

    asm_ptr_init_disp(&ret_reg, RBP, get_register_stack_offset(reg));
    /* Result is returned in RAX register. TODO: Make this work for larger data */
    return_if_error(asm_mov_rm(&impl->asm, RAX, &ret_reg));
    return amal_exec_func_end(self);
}

int amal_exec_func_start(amal_executor *self, u16 num_regs) {
    /*
        TODO: Validate stack size, or maybe remove all validation? do we really need validation?
        If we need security, we could fork the process instead.
    */

    /*
        Some registers need to be preserved before entering a function scope and these registers are different on different platforms.
        32-bit: EBX, ESI, EDI, EBP
        64-bit Windows: RBX, RSI, RDI, RBP, R12-R15, XMM6-XMM15
        64-bit Linux,BSD,Mac: RBX, RBP, R12-R15
    */
    IMPL
    impl->function_indices[impl->func_counter++] = asm_get_size(&impl->asm);
    return_if_error(asm_pushr(&impl->asm, RBX));
    return_if_error(asm_pushr(&impl->asm, RBP));
    return_if_error(asm_mov_rr(&impl->asm, RBP, RSP));
    return asm_sub_rm64_imm(&impl->asm, RSP, num_regs * sizeof(isize));
}

int amal_exec_func_end(amal_executor *self) {
    IMPL
    return_if_error(asm_mov_rr(&impl->asm, RSP, RBP));
    return_if_error(asm_popr(&impl->asm, RBP));
    return_if_error(asm_popr(&impl->asm, RBX));
    return asm_ret(&impl->asm, 0);
}