#include "../../include/ssa/ssa.h"
#include "../../include/std/mem.h"
#include "../../include/std/log.h"
#include "../../include/ast.h"
#include <assert.h>

static int compare_number(const void *a, const void *b) {
    const SsaNumber *lhs;
    const SsaNumber *rhs;
    lhs = a;
    rhs = b;
    if(rhs->type == lhs->type && rhs->value == lhs->value)
        return 0;
    return 1;
}

static usize hash_number(const u8 *data, usize size) {
    SsaNumber number;
    assert(size == sizeof(SsaNumber));
    am_memcpy(&number, data, size);
    return number.value;
}

SsaNumber create_ssa_number(i64 value, SsaNumberType type) {
    SsaNumber result;
    result.value = value;
    result.type = type;
    return result;
}

int ssa_init(Ssa *self, ScopedAllocator *allocator) {
    return_if_error(buffer_init(&self->instructions, allocator));
    return_if_error(hash_map_init(&self->intermediates, allocator, sizeof(SsaIntermediateIndex), compare_number, hash_number));
    self->intermediate_counter = 0;
    self->reg_counter = 0;
    self->func_counter = 0;
    return 0;
}

int ssa_get_unique_reg(Ssa *self, SsaRegister *result) {
    /* Overflow */
    if(self->reg_counter + 1 < self->reg_counter)
        return -1;
    *result = self->reg_counter++;
    return 0;
}

static CHECK_RESULT int ssa_try_add_intermediate(Ssa *self, i64 intermediate, SsaNumberType number_type, SsaIntermediateIndex *result_index) {
    SsaNumber number;
    bool exists;
    BufferView key;

    assert(result_index);
    number = create_ssa_number(intermediate, number_type);
    key = create_buffer_view((const char*)&number, sizeof(number));

    exists = hash_map_get(&self->intermediates, key, result_index);
    if(exists)
        return 0;
    
    /* Overflow */
    if(self->intermediate_counter + 1 < self->intermediate_counter)
        return -1;
    
    *result_index = self->intermediate_counter;
    ++self->intermediate_counter;
    return hash_map_insert(&self->intermediates, key, result_index);
}

static CHECK_RESULT int ssa_add_ins_form1(Ssa *self, SsaInstructionType ins_type, SsaRegister lhs, u16 rhs) {
    usize index;
    index = self->instructions.size;

    return_if_error(buffer_append(&self->instructions, NULL, sizeof(u8) + sizeof(SsaRegister) + sizeof(u16)));
    self->instructions.data[index + 0] = ins_type;
    *(SsaRegister*)&self->instructions.data[index + 1] = lhs;
    *(u16*)&self->instructions.data[index + 3] = rhs;
    return 0;
}

static CHECK_RESULT int ssa_add_ins_form2(Ssa *self, SsaInstructionType ins_type, SsaRegister lhs, SsaRegister rhs, SsaRegister *result) {
    usize index;
    index = self->instructions.size;

    /* Overflow */
    if(self->reg_counter + 1 < self->reg_counter)
        return -1;

    assert(result);
    return_if_error(buffer_append(&self->instructions, NULL, sizeof(u8) + sizeof(SsaRegister) + sizeof(SsaRegister) + sizeof(SsaRegister)));
    *result = self->reg_counter++;
    self->instructions.data[index + 0] = ins_type;
    *(SsaRegister*)&self->instructions.data[index + 1] = *result;
    *(SsaRegister*)&self->instructions.data[index + 3] = lhs;
    *(SsaRegister*)&self->instructions.data[index + 5] = rhs;
    amal_log_debug("r%u = r%u + r%u", *result, lhs, rhs);
    return 0;
}

int ssa_ins_assign_inter(Ssa *self, SsaRegister dest, SsaNumber number) {
    SsaIntermediateIndex index;
    return_if_error(ssa_try_add_intermediate(self, number.value, number.type, &index));
    amal_log_debug("r%u = i%u", dest, index);
    return ssa_add_ins_form1(self, SSA_ASSIGN_INTER, dest, index);
}

int ssa_ins_assign_reg(Ssa *self, SsaRegister dest, SsaRegister src) {
    amal_log_debug("r%u = r%u", dest, src);
    return ssa_add_ins_form1(self, SSA_ASSIGN_INTER, dest, src);
}

int ssa_ins_binop(Ssa *self, SsaInstructionType binop_type, SsaRegister lhs, SsaRegister rhs, SsaRegister *result) {
    assert(binop_type >= SSA_ADD && binop_type <= SSA_DIV);
    return ssa_add_ins_form2(self, binop_type, lhs, rhs, result);
}

int ssa_ins_func_start(Ssa *self, u8 num_args, SsaFuncIndex *result) {
    usize index;
    index = self->instructions.size;

    /* Overflow */
    if(self->func_counter + 1 < self->func_counter)
        return -1;

    return_if_error(buffer_append(&self->instructions, NULL, sizeof(u8) + sizeof(SsaFuncIndex) + sizeof(u8)));
    *result = self->func_counter++;
    self->instructions.data[index + 0] = SSA_FUNC_START;
    *(SsaFuncIndex*)&self->instructions.data[index + 1] = *result;
    self->instructions.data[index + 3] = num_args;
    amal_log_debug("FUNC_START f%u(%u)", *result, num_args);
    return 0;
}

int ssa_ins_func_end(Ssa *self) {
    SsaInstructionType ins;
    ins = SSA_FUNC_END;
    return buffer_append(&self->instructions, &ins, 1);
}

int ssa_ins_push(Ssa *self, SsaRegister reg) {
    usize index;
    index = self->instructions.size;

    return_if_error(buffer_append(&self->instructions, NULL, sizeof(u8) + sizeof(SsaRegister)));
    self->instructions.data[index + 0] = SSA_PUSH;
    *(SsaRegister*)&self->instructions.data[index + 1] = reg;
    amal_log_debug("PUSH r%u", reg);
    return 0;
}

int ssa_ins_call(Ssa *self, SsaFuncIndex func, SsaRegister *result) {
    usize index;
    index = self->instructions.size;

    /* Overflow */
    if(self->reg_counter + 1 < self->reg_counter)
        return -1;

    return_if_error(buffer_append(&self->instructions, NULL, sizeof(u8) + sizeof(SsaFuncIndex) + sizeof(SsaRegister)));
    *result = self->reg_counter++;
    self->instructions.data[index + 0] = SSA_CALL;
    *(SsaFuncIndex*)&self->instructions.data[index + 1] = func;
    *(SsaRegister*)&self->instructions.data[index + 3] = *result;
    amal_log_debug("r%u = CALL f%u", *result, func);
    return 0;
}

static void ast_generate_ssa(Ast *self, SsaCompilerContext *context) {
    /* TODO: Implement */
    (void)self;
    (void)context;
}

void scope_generate_ssa(Scope *self, SsaCompilerContext *context) {
    Ast *ast;
    Ast *ast_end;
    ast = buffer_start(&self->ast_objects);
    ast_end = buffer_end(&self->ast_objects);
    for(; ast != ast_end; ++ast) {
        ast_generate_ssa(ast, context);
    }
}