#include "asm.h"

#include "../../include/std/mem.h"
#include "../../include/std/log.h"
#include <errno.h>
#include <assert.h>
#include <stdio.h>

#include <sys/mman.h>

#define REX_W 0x48

void asm_ptr_init(AsmPtr *self, Reg64 base) {
    self->base = base;
    self->index = -1;
    self->disp = 0;
    self->scale = 0;
}

void asm_ptr_init_index(AsmPtr *self, Reg64 base, Reg64 index) {
    self->base = base;
    self->index = index;
    self->disp = 0;
    self->scale = 0;
}

void asm_ptr_init_disp(AsmPtr *self, Reg64 base, i32 disp) {
    self->base = base;
    self->index = -1;
    self->disp = disp;
    self->scale = 0;
}

void asm_ptr_init_index_disp(AsmPtr *self, Reg64 base, Reg64 index, i32 disp) {
    self->base = base;
    self->index = index;
    self->disp = disp;
    self->scale = 0;
}

int asm_init(Asm *self) {
    self->size = am_pagesize();
    amal_log_debug("asm: page size: %u", self->size);
    self->code = mmap(NULL, self->size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if(self->code == MAP_FAILED)
        return -errno;
    self->code_it = self->code;
    return 0;
}

void asm_deinit(Asm *self) {
    if(self->code)
        munmap(self->code, self->size);
    self->code = NULL;
    self->code_it = NULL;
    self->size = 0;
}

static void asm_print_code_hex(Asm *self) {
    u8 *ptr;
    int off;
    ptr = self->code;
    off = 0;
    while(ptr != self->code_it) {
        printf("%02x", *ptr);
        ++ptr;
        ++off;
        if(off == 8) {
            putc('\n', stdout);
            off = 0;
        } else {
            putc(' ', stdout);
        }
    }
    if(off != 0)
        putc('\n', stdout);
}

int asm_execute(Asm *self) {
    void (*func)();
    if(mprotect(self->code, self->size, PROT_READ | PROT_EXEC) != 0)
        return -errno;

    asm_print_code_hex(self);

    /* TODO: Verify if this is valid on all platforms. According to ISO C standard it isn't? */
    *(void**)(&func) = self->code;
    func();
    return 0;
}

/* TODO: See how this can be optimized */
static CHECK_RESULT int asm_ensure_capacity(Asm *self, usize size) {
    usize current_offset;
    current_offset = (u8*)self->code_it - (u8*)self->code;
    if(current_offset + size > self->size) {
        void *new_mem;
        usize new_size;
        new_size = self->size + am_pagesize();
        new_mem = mmap(NULL, new_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if(self->code == MAP_FAILED)
            return -errno;

        am_memcpy(new_mem, self->code, self->size);
        self->code = new_mem;
        self->size = new_size;
        self->code_it = (u8*)self->code + current_offset;
    }
    return 0;
}

#ifdef DEBUG
static isize asm_get_capacity_left(Asm *self) {
    return (isize)self->size - (isize)((u8*)self->code_it - (u8*)self->code);
}
#endif

int asm_nop(Asm *self) {
    return_if_error(asm_ensure_capacity(self, 1));
    *self->code_it++ = 0x90;
    return 0;
}

static i32 abs_i32(i32 value) {
    return value >= 0 ? value : -value;
}

/*
TODO: Implement 1 and 2 byte displacement?
There has to be at least 6 bytes left in the asm buffer before calling this function.
*/
static void asm_rm(Asm *self, AsmPtr *mem, Reg64 reg) {
    u8 rm_byte;
    u8 disp_bytes;
    assert(asm_get_capacity_left(self) >= 6);
    if((int)mem->index != -1) {
        u8 sib_offset;
        if(mem->disp == 0) {
            rm_byte = 0x04;
            disp_bytes = 0;
        } else if(abs_i32(mem->disp) <= INT8_MAX) {
            rm_byte = 0x44;
            disp_bytes = 1;
        } else {
            rm_byte = 0x84;
            disp_bytes = 4;
        }

        #ifdef DEBUG
        if(mem->scale != 0 && mem->scale != 2 && mem->scale != 4 && mem->scale != 8) {
            amal_log_error("Invalid scale %d, expected 0, 2, 4, or 8", mem->scale);
            assert(bool_false);
        }
        #endif
        assert(mem->base != RBP && "TODO: Implement RBP base for sib byte. RBP is special and requires different logic");
        sib_offset = (mem->scale << 5) + 8*mem->index + mem->base;

        *self->code_it++ = rm_byte;
        *self->code_it++ = sib_offset;
    } else {
        if(mem->disp == 0) {
            if(mem->base == RBP) {
                rm_byte = 0x45;
                disp_bytes = 1;
            } else {
                rm_byte = mem->base;
                disp_bytes = 0;
            }
        } else if(abs_i32(mem->disp) <= INT8_MAX) {
            rm_byte = 0x40 + mem->base;
            disp_bytes = 1;
        } else {
            rm_byte = 0x80 + mem->base;
            disp_bytes = 4;
        }
        *self->code_it++ = (reg << 3) | rm_byte;
    }

    am_memcpy(self->code_it, &mem->disp, disp_bytes);
    self->code_it += disp_bytes;
}

/* There has to be at least 1 byte left in the asm buffer before calling this function. */
static void asm_rr(Asm *self, Reg64 dst, Reg64 src) {
    assert(asm_get_capacity_left(self) >= 1);
    *self->code_it++ = 0xC0 + dst + 8*src;
}

/* TODO: Implement 1 and 2 byte immediate? */
int asm_mov_mi(Asm *self, AsmPtr *dst, i32 immediate) {
    /* 12 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 12));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0xC7;
    asm_rm(self, dst, 0);
    am_memcpy(self->code_it, &immediate, sizeof(immediate));
    self->code_it += sizeof(immediate);
    return 0;
}

int asm_mov_mr(Asm *self, AsmPtr *dst, Reg64 src) {
    /* 8 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 8));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0x89;
    asm_rm(self, dst, src);
    return 0;
}

int asm_mov_rm(Asm *self, Reg64 dst, AsmPtr *src) {
    /* 8 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 8));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0x8B;
    asm_rm(self, src, dst);
    return 0;
}

int asm_mov_ri(Asm *self, Reg64 dst, i64 immediate) {
    /* 10 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 10));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0xB8 + dst;
    am_memcpy(self->code_it, &immediate, sizeof(immediate));
    self->code_it += sizeof(immediate);
    return 0;
}

int asm_mov_rr(Asm *self, Reg64 dst, Reg64 src) {
    /* 3 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 3));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0x89;
    asm_rr(self, dst, src);
    return 0;
}

int asm_add_rr(Asm *self, Reg64 dst, Reg64 src) {
    /* 3 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 3));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0x01;
    asm_rr(self, dst, src);
    return 0;
}

int asm_sub_rr(Asm *self, Reg64 dst, Reg64 src) {
    /* 3 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 3));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0x29;
    asm_rr(self, dst, src);
    return 0;
}

int asm_imul_rr(Asm *self, Reg64 dst, Reg64 src) {
    /* 3 bytes is the maximum size of the instruction. We don't how how large it will be so we prepare for the largest size */
    return_if_error(asm_ensure_capacity(self, 4));
    *self->code_it++ = REX_W;
    *self->code_it++ = 0x0F;
    *self->code_it++ = 0xAF;
    asm_rr(self, dst, src);
    return 0;
}

int asm_pushr(Asm *self, Reg64 reg) {
    return_if_error(asm_ensure_capacity(self, 1));
    *self->code_it++ = 0x50 + reg;
    return 0;
}

int asm_popr(Asm *self, Reg64 reg) {
    return_if_error(asm_ensure_capacity(self, 1));
    *self->code_it++ = 0x58 + reg;
    return 0;
}

/* /r */
#define DEFINE_INS_RM(mnemonic, opcode)                         \
int asm_##mnemonic##_rm32(Asm *self, Reg32 dst, Reg32 src) {    \
    return_if_error(asm_ensure_capacity(self, 2));              \
    *self->code_it++ = opcode;                             \
    *self->code_it++ = 0xC0 + 8*dst + src;                 \
    return 0;                                                   \
}                                                               \
                                                                \
int asm_##mnemonic##_rm64(Asm *self, Reg64 dst, Reg64 src) {    \
    return_if_error(asm_ensure_capacity(self, 1));              \
    *self->code_it++ = REX_W;                               \
    return asm_##mnemonic##_rm32(self, (Reg32)dst, (Reg32)src); \
}

DEFINE_INS_RM(mov, 0x8B)
DEFINE_INS_RM(add, 0x03)
DEFINE_INS_RM(sub, 0x2B)
DEFINE_INS_RM(and, 0x23)
DEFINE_INS_RM(or, 0x0B)
DEFINE_INS_RM(xor, 0x33)
DEFINE_INS_RM(cmp, 0x3B)

/*
    /number
    The number is called the extension, a number from 0 to 7;
    It's a number used to extend the opcode type, since the instruction only uses
    one register the other register can be encoded for that.
*/
#define DEFINE_INS_EXT_IMM(mnemonic, extension)                      \
int asm_##mnemonic##_rm32_imm(Asm *self, Reg32 reg, i32 immediate) { \
    if(abs_i32(immediate) <= INT8_MAX) {                                 \
        return_if_error(asm_ensure_capacity(self, 3));               \
        *self->code_it++ = 0x83;                                \
        *self->code_it++ = 0xC0 + 8*extension + reg;            \
        *self->code_it++ = (u8)immediate;                       \
    } else {                                                         \
        return_if_error(asm_ensure_capacity(self, 6));               \
        *self->code_it++ = 0x81;                                \
        *self->code_it++ = 0xC0 + 8*extension + reg;            \
        am_memcpy(self->code_it, &immediate, sizeof(immediate));     \
        self->code_it += sizeof(immediate);                          \
    }                                                                \
    return 0;                                                        \
}                                                                    \
                                                                     \
int asm_##mnemonic##_rm64_imm(Asm *self, Reg64 reg, i32 immediate) { \
    return_if_error(asm_ensure_capacity(self, 1));                   \
    *self->code_it++ = REX_W;                                    \
    return asm_##mnemonic##_rm32_imm(self, (Reg32)reg, immediate);   \
}

DEFINE_INS_EXT_IMM(add, 0)
DEFINE_INS_EXT_IMM(or, 1)
DEFINE_INS_EXT_IMM(adc, 2)
DEFINE_INS_EXT_IMM(sbb, 3)
DEFINE_INS_EXT_IMM(and, 4)
DEFINE_INS_EXT_IMM(sub, 5)
DEFINE_INS_EXT_IMM(xor, 6)
DEFINE_INS_EXT_IMM(cmp, 7)

/*
    /number
    The number is called the extension, a number from 0 to 7;
    It's a number used to extend the opcode type, since the instruction only uses
    one register the other register can be encoded for that.
*/
#define DEFINE_INS_SHIFT_IMM8(mnemonic, extension)                  \
int asm_##mnemonic##_rm32_imm(Asm *self, Reg32 reg, i8 immediate) { \
    if(immediate == 1) {                                            \
        return_if_error(asm_ensure_capacity(self, 2));              \
        *self->code_it++ = 0xC1;                               \
        *self->code_it++ = 0xC0 + 8*reg + extension;           \
    } else {                                                        \
        return_if_error(asm_ensure_capacity(self, 3));              \
        *self->code_it++ = 0xD1;                               \
        *self->code_it++ = 0xC0 + 8*reg + extension;           \
        *self->code_it++ = immediate;                          \
    }                                                               \
    return 0;                                                       \
}                                                                   \
                                                                    \
int asm_##mnemonic##_rm64_imm(Asm *self, Reg64 reg, i8 immediate) { \
    return_if_error(asm_ensure_capacity(self, 1));                  \
    *self->code_it++ = REX_W;                                   \
    return asm_##mnemonic##_rm32_imm(self, (Reg32)reg, immediate);  \
}

DEFINE_INS_SHIFT_IMM8(rol, 0)
DEFINE_INS_SHIFT_IMM8(ror, 1)
DEFINE_INS_SHIFT_IMM8(rcl, 2)
DEFINE_INS_SHIFT_IMM8(rcr, 3)
DEFINE_INS_SHIFT_IMM8(shl, 4)
DEFINE_INS_SHIFT_IMM8(shr, 5)
/*DEFINE_INS_SHIFT_IMM8(shl, 6)*/
DEFINE_INS_SHIFT_IMM8(sar, 7)

int asm_ret(Asm *self, u16 bytes) {
    if(bytes == 0) {
        return_if_error(asm_ensure_capacity(self, 1));
        *self->code_it++ = 0xC3;
    } else {
        return_if_error(asm_ensure_capacity(self, 3));
        *self->code_it++ = 0xC2;
        am_memcpy(self->code_it, &bytes, sizeof(bytes));
    }
    return 0;
}