#include "jit_x86.h"
#include "codegen_x86.h"
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>

static uint8_t reg_enc(PhysReg r)
{
    switch (r) {
        case RAX:
            return 0;
        case RCX:
            return 1;
        case RDX:
            return 2;
        case RBX:
            return 3;
        case RSI:
            return 6;
        case RDI:
            return 7;
        case R8:
            return 8;
        case R9:
            return 9;
        default:
            return 0;
    }
}

typedef struct {
    uint8_t* buf;
    size_t cap;
    size_t len;
} AsmBuf;

static void emit8(AsmBuf* a, uint8_t v)
{
    if (a->len >= a->cap) {
        fprintf(stderr, "JIT buffer overflow\n");
        abort();
    }
    a->buf[a->len++] = v;
}

[[maybe_unused]]
static void emit32(AsmBuf* a, uint32_t v)
{
    if (a->len >= a->cap) {
        fprintf(stderr, "JIT buffer overflow\n");
        abort();
    }
    memcpy(&a->buf[a->len], &v, 4);
    a->len += 4;
}

static void emit64(AsmBuf* a, uint64_t v)
{
    if (a->len >= a->cap) {
        fprintf(stderr, "JIT buffer overflow\n");
        abort();
    }
    memcpy(&a->buf[a->len], &v, 8);
    a->len += 8;
}

static void emit_mov_imm64(AsmBuf* a, PhysReg dst, uint64_t imm)
{
    uint8_t r = reg_enc(dst);

    // mov r64, imm64 = 48 B8+rd imm64
    emit8(a, 0x48);
    emit8(a, 0xB8 + r);
    emit64(a, imm);
}

static uint8_t rex_enc(uint8_t dst, uint8_t src)
{
    return 0x40 | ((dst & 8) ? 0x01 : 0) | // B
        ((src & 8) ? 0x04 : 0); // R
}

static void emit_alu_rr(
    AsmBuf* a, uint8_t rex, uint8_t op, PhysReg dst, PhysReg src)
{
    (void)rex;

    uint8_t d = reg_enc(dst);
    uint8_t s = reg_enc(src);

    emit8(a, rex_enc(d, s));
    emit8(a, op);

    uint8_t modrm = 0xC0 | ((s & 7) << 3) | (d & 7);
    emit8(a, modrm);
}

static void emit_add(AsmBuf* a, PhysReg d, PhysReg s)
{
    emit_alu_rr(a, 0x48, 0x01, d, s);
}

static void emit_sub(AsmBuf* a, PhysReg d, PhysReg s)
{
    emit_alu_rr(a, 0x48, 0x29, d, s);
}

static void emit_mul(AsmBuf* a, PhysReg d, PhysReg s)
{
    emit8(a, 0x48);
    emit8(a, 0x0F);
    emit8(a, 0xAF);
    emit8(a, 0xC0 | (reg_enc(s) << 3) | reg_enc(d));
}

static void emit_ret(AsmBuf* a)
{
    emit8(a, 0xC3);
}

static size_t align_page(size_t n)
{
    size_t page = sysconf(_SC_PAGESIZE);
    return (n + page - 1) & ~(page - 1);
}

JitFn cg_block_emit_x86_machine_code(CgBlock* block)
{
    AsmBuf a = { 0 };

    a.cap = align_page(1024);
    a.len = 0;
    a.buf = mmap(NULL,
        a.cap,
        PROT_READ | PROT_WRITE,
        MAP_PRIVATE | MAP_ANONYMOUS,
        -1,
        0);

    if (a.buf == MAP_FAILED) {
        return NULL;
    }

    for (size_t i = 0; i < block->count; i++) {
        CgInst* inst = block->insts[i];

        switch (inst->op) {

            case CG_IMM64:
                emit_mov_imm64(&a, inst->dst, inst->imm);
                break;

            case CG_ADD8:
                emit_add(&a, inst->dst, inst->binop.rhs);
                break;

            case CG_SUB8:
                emit_sub(&a, inst->dst, inst->binop.rhs);
                break;

            case CG_MUL8:
                emit_mul(&a, inst->dst, inst->binop.rhs);
                break;

            default:
                break;
        }
    }

    // ensure result is in rax (ABI return register)
    emit_ret(&a);

    // make executable
    mprotect(a.buf, align_page(a.len), PROT_READ | PROT_EXEC);

    for (size_t i = 0; i < a.len; i++) {
        printf("%02x ", a.buf[i]);
    }
    printf("\n");

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpedantic"
    return (JitFn)a.buf;
#pragma GCC diagnostic pop
}