#include "../../include/ir/ir.h"
#include "../../include/std/mem.h"
#include "../../include/std/log.h"
#include "../../include/std/hash.h"
#include "../../include/std/thread.h"
#include "../../include/ast.h"
#include "../../include/compiler.h"
#include "../../include/parser.h"
#include <assert.h>

/*
    TODO: Instead of using memcpy to copy data to the ir, make it cleaner by
    defining all the data in structs and copying the structs. Even if it takes more space,
    it might even be faster.
*/

#define throw(result) do { longjmp(context->env, (result)); } while(0)
#define throw_if_error(result) \
do { \
    int return_if_result; \
    return_if_result = (result); \
    if((return_if_result) != 0) \
        throw(return_if_result); \
} while(0)

/* Max length of a string that fits in u16 */
#define MAX_STRING_LENGTH UINT16_MAX
#define FUNC_MAX_ARGS 128

static CHECK_RESULT IrRegister variable_generate_ir(Variable *self, IrCompilerContext *context);

static int compare_number(const void *a, const void *b) {
    const IrNumber *lhs = a;
    const IrNumber *rhs = b;
    return (rhs->type == lhs->type && rhs->value.integer == lhs->value.integer);
}

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

IrNumber create_ir_integer(i64 value) {
    IrNumber result;
    result.value.integer = value;
    result.type = IR_NUMBER_TYPE_INTEGER;
    return result;
}

IrNumber create_ir_float(f64 value) {
    IrNumber result;
    result.value.floating = value;
    result.type = IR_NUMBER_TYPE_FLOAT;
    return result;
}

int ir_init(Ir *self, Parser *parser) {
    return_if_error(buffer_init(&self->instructions, parser->allocator));
    return_if_error(hash_map_init(&self->intermediates_map, parser->allocator, sizeof(IrIntermediateIndex), compare_number, hash_number));
    return_if_error(buffer_init(&self->intermediates, parser->allocator));
    return_if_error(hash_map_init(&self->strings_map, parser->allocator, sizeof(IrStringIndex), hash_map_compare_string, amal_hash_string));
    return_if_error(buffer_init(&self->strings, parser->allocator));
    return_if_error(hash_map_init(&self->extern_funcs_map, parser->allocator, sizeof(IrExternFuncIndex), hash_map_compare_string, amal_hash_string));
    return_if_error(buffer_init(&self->extern_funcs, parser->allocator));
    return_if_error(buffer_init(&self->export_funcs, parser->allocator));
    return_if_error(buffer_init(&self->funcs, parser->allocator));
    self->intermediate_counter = 0;
    self->string_counter = 0;
    self->extern_func_counter = 0;
    self->export_func_counter = 0;
    self->func_counter = 0;
    self->reg_counter = 0;
    self->param_counter = 0;
    self->label_counter = 0;
    self->parser = parser;
    return 0;
}

static CHECK_RESULT int ir_get_unique_reg(Ir *self, IrRegister *result) {
    assert(result);
    /* Overflow */
    if((u16)self->reg_counter + 1 > INT16_MAX) {
        amal_log_error("Ir too many registers!");
        return -1;
    }
    assert(self->reg_counter <= INT8_MAX && "TODO: Implement usage of reg higher than 128");
    *result = self->reg_counter++;
    return 0;
}

static CHECK_RESULT int ir_get_unique_param_reg(Ir *self, IrRegister *result) {
    assert(result);
    /* Overflow */
    if((u16)self->param_counter + 1 > INT16_MAX) {
        amal_log_error("Ir too many param registers!");
        return -1;
    }
    assert(self->param_counter <= INT8_MAX && "TODO: Implement usage of reg higher than 128");
    *result = self->param_counter++ | REG_FLAG_PARAM;
    return 0;
}

IrNumber ir_get_intermediate(Ir *self, IrIntermediateIndex index) {
    IrNumber result;
    assert(index < buffer_get_size(&self->intermediates, IrNumber));
    am_memcpy(&result, buffer_get(&self->intermediates, index, sizeof(IrNumber)), sizeof(IrNumber));
    return result;
}

BufferView ir_get_string(Ir *self, IrStringIndex index) {
    BufferView result;
    assert(index < buffer_get_size(&self->strings, BufferView));
    am_memcpy(&result, buffer_get(&self->strings, index, sizeof(BufferView)), sizeof(BufferView));
    return result;
}

static CHECK_RESULT int ir_try_add_intermediate(Ir *self, IrNumber number, IrIntermediateIndex *result_index) {
    bool exists;
    BufferView key;

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

    exists = hash_map_get(&self->intermediates_map, key, result_index);
    if(exists)
        return 0;
    
    /* Overflow */
    if(self->intermediate_counter + 1 <= self->intermediate_counter) {
        amal_log_error("Ir too many intermediates!");
        return -1;
    }
    
    *result_index = self->intermediate_counter;
    ++self->intermediate_counter;
    switch(number.type) {
        case IR_NUMBER_TYPE_INTEGER: {
            amal_log_debug("i%u = %lld", *result_index, number.value.integer);
            break;
        }
        case IR_NUMBER_TYPE_FLOAT: {
            amal_log_debug("i%u = %f", *result_index, number.value.floating);
            break;
        }
    }

    return_if_error(buffer_append(&self->intermediates, &number, sizeof(number)));
    return hash_map_insert(&self->intermediates_map, key, result_index);
}

static CHECK_RESULT int ir_try_add_string(Ir *self, BufferView str, IrStringIndex *result_index) {
    bool exists;
    assert(result_index);

    exists = hash_map_get(&self->strings_map, str, result_index);
    if(exists)
        return 0;
    
    /* Overflow */
    if(self->string_counter + 1 <= self->string_counter) {
        amal_log_error("Ir too many strings!");
        return -1;
    }

    if(str.size > MAX_STRING_LENGTH) {
        amal_log_error("String \"%.*s\" is longer than %d\n", str.size, str.data, MAX_STRING_LENGTH);
        return -2;
    }
    
    *result_index = self->string_counter;
    ++self->string_counter;
    amal_log_debug("s%u = \"%.*s\"", *result_index, str.size, str.data);
    return_if_error(buffer_append(&self->strings, &str, sizeof(str)));
    return hash_map_insert(&self->strings_map, str, result_index);
}

/*
    TODO: Right now this has the same scope as a file. This should be global, otherwise you could define multiple
    extern func with the same name but different signature as long as they are defined in different files
*/
static CHECK_RESULT int ir_try_add_extern_func(Ir *self, FunctionSignature *func_sig, BufferView name, IrExternFuncIndex *result_index, BufferView *existing_func) {
    bool exists;
    assert(result_index);
    assert(existing_func);

    exists = hash_map_get(&self->extern_funcs_map, name, result_index);
    if(exists) {
        const IrExternFunc *existing_extern_func = buffer_get(&self->extern_funcs, *result_index, sizeof(IrExternFunc));
        *existing_func = existing_extern_func->name;
        if(!function_signature_equals(func_sig, existing_extern_func->func_sig))
            return IR_ERR_EXTERN_FUNC_SIG_MISMATCH;
        return 0;
    }
    
    /* Overflow */
    if(self->extern_func_counter + 1 <= self->extern_func_counter) {
        amal_log_error("Ir too many extern closures!");
        return -1;
    }
    
    *result_index = self->extern_func_counter;
    ++self->extern_func_counter;
    amal_log_debug("extern_func%u = %.*s", *result_index, name.size, name.data);
    {
        IrExternFunc extern_func;
        extern_func.func_sig = func_sig;
        extern_func.name = name;
        return_if_error(buffer_append(&self->extern_funcs, &extern_func, sizeof(extern_func)));
        return hash_map_insert(&self->extern_funcs_map, name, result_index);
    }
}

/*
    Exported functions with the same name are allowed to exist in different files, but not in the same file.
    There is no need to check if there are two exported functions with the same name in the same file
    as that is already checked in the AST stage, as you are not allowed to have multiple variables of the same name in the same scope,
    exported or not.
*/
static CHECK_RESULT int ir_try_add_export_func(Ir *self, FunctionSignature *func_sig, BufferView name) {
    /* Overflow */
    if(self->export_func_counter + 1 <= self->export_func_counter) {
        amal_log_error("Ir too many exported closures!");
        return -1;
    }
    
    amal_log_debug("exported_func%u = %.*s", self->export_func_counter, name.size, name.data);
    ++self->export_func_counter;
    {
        IrExportFunc export_func;
        export_func.func_sig = func_sig;
        export_func.name = name;
        return buffer_append(&self->export_funcs, &export_func, sizeof(export_func));
    }
}

static CHECK_RESULT int ir_add_ins_form1(Ir *self, IrInstruction ins_type, IrRegister lhs, u16 rhs) {
    IrInsForm1 ins_form_1;
    ins_form_1.lhs = lhs;
    ins_form_1.rhs = rhs;
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &ins_form_1, sizeof(ins_form_1));
}

static const char* binop_type_to_string(IrInstruction binop_type) {
    assert(binop_type >= IR_ADD && binop_type <= IR_GE);
    switch(binop_type) {
        case IR_ADD:     return "+";
        case IR_SUB:     return "-";
        case IR_MUL:     return "*";
        case IR_DIV:     return "/";
        case IR_EQUALS:  return "==";
        case IR_AND:     return "&&";
        case IR_ILT:     return "<";
        case IR_LT:      return "<";
        case IR_ILE:     return "<=";
        case IR_LE:      return "<=";
        case IR_IGT:     return ">";
        case IR_GT:      return ">";
        case IR_IGE:     return ">=";
        case IR_GE:      return ">=";
        default:          return "";
    }
}

static CHECK_RESULT int ir_add_ins_form2(Ir *self, IrInstruction ins_type, IrRegister lhs, IrRegister rhs, IrRegister *result) {
    IrInsForm2 ins_form_2;
    return_if_error(ir_get_unique_reg(self, result));
    ins_form_2.result = *result;
    ins_form_2.lhs = lhs;
    ins_form_2.rhs = rhs;
    amal_log_debug("r%d = r%d %s r%d", *result, lhs, binop_type_to_string(ins_type), rhs);
    /* TODO: Store ins_type as an uint8_t instead of int */
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &ins_form_2, sizeof(ins_form_2));
}

static CHECK_RESULT int ir_ins_assign_inter(Ir *self, IrRegister dest, IrNumber number) {
    IrIntermediateIndex index;
    return_if_error(ir_try_add_intermediate(self, number, &index));
    amal_log_debug("r%d = i%u", dest, index);
    return ir_add_ins_form1(self, IR_ASSIGN_INTER, dest, index);
}

static CHECK_RESULT int ir_ins_assign_string(Ir *self, IrRegister dest, BufferView str) {
    IrStringIndex index;
    return_if_error(ir_try_add_string(self, str, &index));
    amal_log_debug("r%d = s%u", dest, index);
    return ir_add_ins_form1(self, IR_ASSIGN_STRING, dest, index);
}

static CHECK_RESULT int ir_ins_assign_reg(Ir *self, IrRegister dest, IrRegister src) {
    amal_log_debug("r%d = r%d", dest, src);
    return ir_add_ins_form1(self, IR_ASSIGN_REG, dest, src);
}

static CHECK_RESULT int ir_ins_assign_func(Ir *self, IrRegister dest, IrFuncIndex func_index) {
    amal_log_debug("r%d = f%d", dest, func_index);
    return ir_add_ins_form1(self, IR_ASSIGN_FUNC, dest, func_index);
}

static CHECK_RESULT int ir_ins_binop(Ir *self, IrInstruction binop_type, IrRegister lhs, IrRegister rhs, IrRegister *result) {
    assert(binop_type >= IR_ADD && binop_type <= IR_GE);
    return ir_add_ins_form2(self, binop_type, lhs, rhs, result);
}

static CHECK_RESULT int ir_ins_func_start(Ir *self, u8 func_flags, FunctionSignature *func_sig, usize *func_metadata_index) {
    const u8 ins_type = IR_FUNC_START;
    IrInsFuncStart ins_func_start;
    IrFunc func;

    func.func_sig = func_sig;
    return_if_error(buffer_append(&self->funcs, &func, sizeof(func)));

    ins_func_start.flags = func_flags;
    /* Dont set number of local registers yet. That will be set by @func_metadata_index later when it's known */
    /*ins_func_start.num_local_vars_regs = ---*/
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return_if_error(buffer_append(&self->instructions, &ins_func_start, sizeof(ins_func_start)));
    *func_metadata_index = self->instructions.size - sizeof(ins_func_start.num_local_vars_regs);
    amal_log_debug("FUNC_START f%d(%d) %d",
        func_sig->func_decl ? func_sig->func_decl->ir_func_index : -1,
        buffer_get_size(&func_sig->parameters, FunctionParameter),
        buffer_get_size(&func_sig->return_types, FunctionReturnType));
    return 0;
}

static CHECK_RESULT int ir_inc_func_index(Ir *self, IrFuncIndex *result) {
    /* Overflow */
    if(self->func_counter + 1 <= self->func_counter) {
        amal_log_error("Ir too many closures!");
        return -1;
    }
    *result = self->func_counter++;
    return 0;
}

static CHECK_RESULT int ir_ins_func_end(Ir *self) {
    const u8 ins_type = IR_FUNC_END;
    amal_log_debug("FUNC_END");
    return buffer_append(&self->instructions, &ins_type, 1);
}

static CHECK_RESULT int ir_ins_push(Ir *self, IrRegister reg) {
    const u8 ins_type = IR_PUSH;
    amal_log_debug("PUSH r%d", reg);
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &reg, sizeof(reg));
}

static CHECK_RESULT int ir_ins_push_ret(Ir *self, IrRegister reg) {
    const u8 ins_type = IR_PUSH_RET;
    amal_log_debug("PUSH RET r%d", reg);
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &reg, sizeof(reg));
}

static CHECK_RESULT int ir_ins_call_start(Ir *self, u8 num_args) {
    const u8 ins_type = IR_CALL_START;
    IrInsCallStart ins_call_start;
    ins_call_start.num_args = num_args;
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &ins_call_start, sizeof(ins_call_start));
}

static CHECK_RESULT int ir_ins_call(Ir *self, int import_index, FunctionDecl *func_decl) {
    const u8 ins_type = IR_CALL;
    IrInsFuncCall ins_func_call;
    ins_func_call.func_decl = func_decl;
    ins_func_call.import_index = import_index;
    amal_log_debug("CALL f(%d,%p)", import_index, func_decl);
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &ins_func_call, sizeof(ins_func_call));
}

static CHECK_RESULT int ir_ins_call_extern(Ir *self, int import_index, LhsExpr *func_decl_lhs) {
    const u8 ins_type = IR_CALL_EXTERN;
    IrInsFuncCallExtern ins_func_call_extern;
    ins_func_call_extern.func_decl_lhs = func_decl_lhs;
    ins_func_call_extern.import_index = import_index;
    amal_log_debug("CALL_EXTERN ef(%d,%p)", import_index, func_decl_lhs);
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &ins_func_call_extern, sizeof(ins_func_call_extern));
}

static CHECK_RESULT int ir_ins_call_reg(Ir *self, IrRegister reg) {
    const u8 ins_type = IR_CALLR;
    amal_log_debug("CALLR r%d", reg);
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &reg, sizeof(reg));
}

static CHECK_RESULT int ir_ins_jumpzero(Ir *self, IrRegister condition_reg, IrLabelIndex target_label, usize *instruction_offset) {
    const u8 ins_type = IR_JUMP_ZERO;
    IrInsJumpZero ins_jump_zero;
    ins_jump_zero.condition_reg = condition_reg;
    ins_jump_zero.target_label = target_label;
    if(target_label == 0)
        amal_log_debug("JUMP_ZERO r%d, DUMMY", condition_reg);
    else
        amal_log_debug("JUMP_ZERO r%d, l%d", condition_reg, target_label);
    *instruction_offset = self->instructions.size;
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &ins_jump_zero, sizeof(ins_jump_zero));
}

static CHECK_RESULT int ir_ins_jump(Ir *self, IrLabelIndex target_label, usize *instruction_offset) {
    const u8 ins_type = IR_JUMP;
    IrInsJump ins_jump;
    ins_jump.target_label = target_label;

    if(target_label == 0)
        amal_log_debug("JUMP DUMMY");
    else
        amal_log_debug("JUMP l%d", target_label);

    if(instruction_offset)
        *instruction_offset = self->instructions.size;

    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &ins_jump, sizeof(ins_jump));
}

static CHECK_RESULT int ir_ins_return(Ir *self, IrRegister reg) {
    const u8 ins_type = IR_RET;
    return_if_error(buffer_append(&self->instructions, &ins_type, 1));
    return buffer_append(&self->instructions, &reg, sizeof(reg));
}

static CHECK_RESULT int ir_ins_label(Ir *self, u16 *label_index) {
    const u8 ins_type = IR_LABEL;
    /* Overflow */
    if(self->label_counter + 1 <= self->label_counter) {
        amal_log_error("Ir too many labels!");
        return -1;
    }
    *label_index = self->label_counter;
    ++self->label_counter;
    return buffer_append(&self->instructions, &ins_type, 1);
}

static CHECK_RESULT int ir_set_jump_label(Ir *self, usize jump_ins_index, IrLabelIndex label) {
    switch(self->instructions.data[jump_ins_index]) {
        case IR_JUMP_ZERO: {
            /* +1 to skip instruction opcode */
            am_memcpy(self->instructions.data + jump_ins_index + 1 + offsetof(IrInsJumpZero, target_label), &label, sizeof(IrLabelIndex));
            break;
        }
        case IR_JUMP: {
            /* +1 to skip instruction opcode */
            am_memcpy(self->instructions.data + jump_ins_index + 1 + offsetof(IrInsJump, target_label), &label, sizeof(IrLabelIndex));
            break;
        }
        default:
            assert(bool_false && "Unexpected error... jump_ins_index doesn't point to a valid index to a jump instruction");
            break;
    }
    return 0;
}

static CHECK_RESULT IrRegister ast_generate_ir(Ast *self, IrCompilerContext *context);
static CHECK_RESULT IrRegister scope_named_object_generate_ir(ScopeNamedObject *self, IrCompilerContext *context);

#if 0
static bool ast_resolved_type_is_decl(AstResolvedType *self) {
    /* TODO: Add more types as they are introduced */
    LhsExpr *lhs_expr;
    if(self->type != RESOLVED_TYPE_LHS_EXPR)
        return bool_false;

    lhs_expr = self->value.lhs_expr;
    switch(lhs_expr->type.type) {
        case VARIABLE_TYPE_NONE:
            break;
        case VARIABLE_TYPE_VARIABLE:
            return bool_false;
        case VARIABLE_TYPE_SIGNATURE:
            /* TODO: This should return bool_false when it's possible to use signature in expressions */
            return bool_true;
    }

    assert(lhs_expr->rhs_expr);
    return lhs_expr->rhs_expr->type == AST_FUNCTION_DECL || lhs_expr->rhs_expr->type == AST_STRUCT_DECL;
}
#endif

static CHECK_RESULT IrRegister number_generate_ir(Number *self, IrCompilerContext *context) {
    IrRegister reg;
    IrNumber number;
    switch(self->value.type) {
        case AMAL_NUMBER_SIGNED_INTEGER:
        case AMAL_NUMBER_UNSIGNED_INTEGER:
            number = create_ir_integer(self->value.value.integer);
            break;
        case AMAL_NUMBER_FLOAT:
            number = create_ir_float(self->value.value.floating);
            break;
    }
    throw_if_error(ir_get_unique_reg(context->ir, &reg));
    throw_if_error(ir_ins_assign_inter(context->ir, reg, number));
    return reg;
}

static CHECK_RESULT IrRegister ast_bool_generate_ir(AstBool *self, IrCompilerContext *context) {
    IrRegister reg;
    IrNumber number;
    number.type = IR_NUMBER_TYPE_INTEGER;
    number.value.integer = self->value;
    throw_if_error(ir_get_unique_reg(context->ir, &reg));
    /* TODO: Maybe bool shouldn't be a regular number? */
    throw_if_error(ir_ins_assign_inter(context->ir, reg, number));
    return reg;
}

static CHECK_RESULT IrRegister lhsexpr_extern_generate_ir(LhsExpr *self, IrCompilerContext *context) {
    /* TODO: IrRegister should be extended to include static and extern data */
    if(self->type.type == VARIABLE_TYPE_SIGNATURE) {
        int err;
        BufferView existing_func;

        err = ir_try_add_extern_func(context->ir, self->type.value.signature, self->var_name, &self->extern_index, &existing_func);
        if(err == IR_ERR_EXTERN_FUNC_SIG_MISMATCH) {
            Tokenizer *tokenizer;
            tokenizer = &context->ir->parser->tokenizer;
            tokenizer_print_error(tokenizer, tokenizer_get_code_reference_index(tokenizer, self->var_name.data),
                "Extern closure defined here with the name %.*s doesn't match extern closure with the same name defined in another location",
                    self->var_name.size, self->var_name.data);
            /*
                TODO: This wont work right now since the other location might belong to another parser.
                There should be a function to get a parser from a code reference (loop all tokens and check code range).
                Then the parsers that belong to scopes can also be removed. This is fine, since the lookup is only done on error.
            */
            #if 0
            tokenizer_print_error(tokenizer, tokenizer_get_code_reference_index(tokenizer, self->var_name.data),
                "Extern closure defined here with the name %.*s doesn't match extern closure with the same name defined in another location",
                    self->var_name.size, self->var_name.data);
            #endif
            throw(err);
        }
        if(err != 0)
            throw(err);
    } else {
        assert(bool_false && "TODO: Implement lhsexpr_extern_generate_ir for other data than functions");
    }
    return 0;
}

#if 0
static CHECK_RESULT IrRegister lhsexpr_export_generate_ir(LhsExpr *self, IrCompilerContext *context) {
    /* TODO: IrRegister should be extended to include static and export data */
    if(self->rhs_expr->type == AST_FUNCTION_DECL) {
        throw_if_error(ir_try_add_export_func(context->ir, self->rhs_expr->value.func_decl->signature, self->var_name));
    } else {
        assert(bool_false && "TODO: Implement lhsexpr_export_generate_ir for other data than functions");
    }
    return 0;
}
#endif

static CHECK_RESULT IrRegister lhsexpr_generate_ir(LhsExpr *self, AstResolveData *resolve_data, IrCompilerContext *context) {
    IrRegister reg;

    if(LHS_EXPR_IS_EXTERN(self))
        return lhsexpr_extern_generate_ir(self, context);

    if(self->rhs_expr) {
        Ast *rhs_expr = self->rhs_expr;
        IrRegister rhs_reg = ast_generate_ir(rhs_expr, context);

#if 0
        if(LHS_EXPR_IS_EXPORT(self))
            return lhsexpr_export_generate_ir(self, context);
#endif
        /*
            Import expression also has no meaning in IR until it's used.
            TODO: Shouldn't lhsexpr that have struct/function declaration as rhs be different ast expression types?
        */
        if(rhs_expr->type == AST_IMPORT) {
            /*assert(bool_false);*/
            return 0;
        }
        return rhs_reg;
    } else {
        /* TODO: Do not assign if we dont want default value */
        if(resolve_data->type.type == RESOLVED_TYPE_LHS_EXPR) {
            IrNumber number;
            if(resolve_data->type.value.lhs_expr == (LhsExpr*)context->compiler->default_types.i64)
                number = create_ir_integer(0);
            else if(resolve_data->type.value.lhs_expr == (LhsExpr*)context->compiler->default_types.f64)
                number = create_ir_float(0.0);
            else
                assert(bool_false && "TODO: assign default value to reg depending on LhsExpr type");
            throw_if_error(ir_get_unique_reg(context->ir, &reg));
            throw_if_error(ir_ins_assign_inter(context->ir, reg, number));
        } else if(resolve_data->type.type == RESOLVED_TYPE_FUNC_SIG) {
            assert(bool_false && "TODO: Implement this when variable can be null. Then the function pointer should be null");
        } else {
            assert(bool_false);
        }
    }
    return reg;
}

static CHECK_RESULT IrRegister assignmentexpr_generate_ir(AssignmentExpr *self, IrCompilerContext *context) {
    IrRegister lhs_reg, rhs_reg;
    lhs_reg = ast_generate_ir(self->lhs_expr, context);
    rhs_reg = ast_generate_ir(self->rhs_expr, context);
    throw_if_error(ir_ins_assign_reg(context->ir, lhs_reg, rhs_reg));
    return lhs_reg;
}

static CHECK_RESULT IrRegister function_parameter_generate_ir(FunctionParameter *self, IrCompilerContext *context) {
    IrRegister reg;
    if(self->resolve_data.status == AST_IR_RESOLVED)
        return self->resolve_data.ir_reg;

    throw_if_error(ir_get_unique_param_reg(context->ir, &reg));
    self->resolve_data.status = AST_IR_RESOLVED;
    self->resolve_data.ir_reg = reg;
    return reg;
}

static CHECK_RESULT void function_signature_generate_params_ir(FunctionSignature *self, IrCompilerContext *context) {
    FunctionParameter *param, *param_end;
    param = buffer_begin(&self->parameters);
    param_end = buffer_end(&self->parameters);
    for(; param != param_end; ++param) {
        IrRegister reg;
        reg = function_parameter_generate_ir(param, context);
        (void)reg;
    }
}

static CHECK_RESULT IrRegister funcdecl_generate_ir(FunctionDecl *self, IrCompilerContext *context) {
    /*
        Reset reg counter in each function, because each function has a separate register context
        that is reset after function end
    */
    usize func_metadata_index;
    u8 func_flags = 0;
    context->ir->reg_counter = 0;
    context->ir->param_counter = 0;
    context->ir->label_counter = 0;

    /* All parameters need to be generated, so that the parameter matches its index... */
    function_signature_generate_params_ir(self->signature, context);

    amal_log_debug("IR funcdecl %p", self);
    /* Anonymous closure doesn't have lhs_expr, and neither can it have any flags (extern, export etc) */
    if(self->lhs_expr) {
        assert(!LHS_EXPR_IS_EXTERN(self->lhs_expr));
        if(LHS_EXPR_IS_EXPORT(self->lhs_expr))
            func_flags |= FUNC_FLAG_EXPORTED;
    }
    throw_if_error(ir_ins_func_start(context->ir, func_flags, self->signature, &func_metadata_index));
    scope_generate_ir(&self->body, context);
    throw_if_error(ir_ins_func_end(context->ir));

    /* Add the number of registers used to the function metadata (FUNC_START) */
    am_memcpy(context->ir->instructions.data + func_metadata_index, &context->ir->reg_counter, sizeof(u16));
    return 0;
}

static CHECK_RESULT IrRegister funcdecl_ref_generate_ir(FunctionDecl *self, IrCompilerContext *context) {
    IrRegister reg;
    if(self->lhs_expr && LHS_EXPR_IS_EXTERN(self->lhs_expr)) {
        assert(bool_false && "TODO: Implement assign func for extern closures");
        throw(-1);
    }
    throw_if_error(ir_get_unique_reg(context->ir, &reg));
    throw_if_error(ir_ins_assign_func(context->ir, reg, self->ir_func_index));
    return reg;
}


static CHECK_RESULT IrRegister funccall_generate_ir(FunctionCall *self, IrCompilerContext *context) {
    IrRegister reg;
    FunctionSignature *func_sig;
    FunctionDecl *func_decl;
    LhsExpr *func_lhs_expr;
    FunctionParameter *func_param_expr;
    int import_index = context->import_index;

    func_lhs_expr = NULL;
    func_param_expr = NULL;
    context->import_index = 0;
    throw_if_error(ir_get_unique_reg(context->ir, &reg));

    switch(self->func.resolved_var.type) {
        case NAMED_OBJECT_NONE:
            assert(bool_false);
            break;
        case NAMED_OBJECT_LHS_EXPR: {
            func_lhs_expr = self->func.resolved_var.value.lhs_expr;
            if(func_lhs_expr->type.type == VARIABLE_TYPE_SIGNATURE) {
                func_sig = func_lhs_expr->type.value.signature;
            } else if(func_lhs_expr->type.type == VARIABLE_TYPE_VARIABLE) {
                AstResolveData *resolve_data = func_lhs_expr->type.value.variable->resolved_var.resolve_data;
                assert(resolve_data->type.type == RESOLVED_TYPE_FUNC_SIG);
                func_sig = resolve_data->type.value.func_sig;
            } else {
                assert(func_lhs_expr->rhs_expr && func_lhs_expr->rhs_expr->resolve_data.type.type == RESOLVED_TYPE_FUNC_SIG);
                func_sig = func_lhs_expr->rhs_expr->resolve_data.type.value.func_sig;
            }
            break;
        }
        case NAMED_OBJECT_FUNC_PARAM: {
            func_param_expr = self->func.resolved_var.value.func_param;
            assert(func_param_expr->type.type == VARIABLE_TYPE_SIGNATURE);
            func_sig = func_param_expr->type.value.signature;
            break;
        }
    }
    
    func_decl = func_sig->func_decl;

    /* Push return arguments */
    {
        /*
            TODO: When amalgam supports multiple return types in assignment/declaration, update this to take
            all of them into account. Right now it only uses one return type.
            It should also take into account the size of the type.
        */
        assert(buffer_get_size(&func_sig->return_types, FunctionReturnType) <= 1);
        throw_if_error(ir_ins_push_ret(context->ir, reg));
    }

    /* Push parameter arguments */
    assert(buffer_get_size(&self->args, Ast*) <= FUNC_MAX_ARGS);
    {
        IrRegister arg_regs[FUNC_MAX_ARGS];

        Ast **arg = buffer_begin(&self->args);
        Ast **arg_end = buffer_end(&self->args);
        for(; arg != arg_end; ++arg) {
            arg_regs[arg_end - arg] = ast_generate_ir(*arg, context);
        }

        /*
            This is done in two steps since first we want the argument expressions to be generated
            and then at the end, push the registers.
            This allows Amalgam to push arguments directly to registers on abi that uses registers as function arguments (system-v x86_64)
            instead of first pushing the registers to stack and then moving them to registers.
        */

        arg = buffer_begin(&self->args);
        throw_if_error(ir_ins_call_start(context->ir, arg_end - arg));
        for(; arg != arg_end; ++arg) {
            throw_if_error(ir_ins_push(context->ir, arg_regs[arg_end - arg]));
        }
    }

    if(func_lhs_expr && LHS_EXPR_IS_EXTERN(func_lhs_expr)) {
        throw_if_error(ir_ins_call_extern(context->ir, import_index, func_lhs_expr));
    } else if(func_decl) {
        throw_if_error(ir_ins_call(context->ir, import_index, func_decl));
    } else if(func_param_expr) {
        throw_if_error(ir_ins_call_reg(context->ir, function_parameter_generate_ir(func_param_expr, context)));
    }

    return reg;
}

#if 0
static CHECK_RESULT IrRegister structdecl_generate_ir(StructDecl *self, IrCompilerContext *context) {
    /* TODO: Implement */
    /*assert(bool_false);*/
    scope_generate_ir(&self->body, context);
    return 0;
}

static CHECK_RESULT IrRegister structfield_generate_ir(StructField *self, IrCompilerContext *context) {
    /* TODO: Implement */
    assert(bool_false);
    (void)self;
    (void)context;
    return 0;
}
#endif

static CHECK_RESULT IrRegister string_generate_ir(String *self, IrCompilerContext *context) {
    IrRegister reg;
    throw_if_error(ir_get_unique_reg(context->ir, &reg));
    throw_if_error(ir_ins_assign_string(context->ir, reg, self->str));
    return reg;
}

IrRegister variable_generate_ir(Variable *self, IrCompilerContext *context) {
    /* TODO: If resolved_var refers to a variable in another file, use a cross file reference that requires no locking (not yet implemented) */
    /* This is not thread-safe:*/
    assert(self->resolved_var.type != NAMED_OBJECT_NONE);
    return scope_named_object_generate_ir(&self->resolved_var, context);
}

static IrInstruction binop_type_to_ir_type(BinopType binop_type, amal_default_type *type) {
    switch(binop_type) {
        case BINOP_ADD:
            return IR_ADD;
        case BINOP_SUB:
            return IR_SUB;
        case BINOP_MUL:
            return amal_default_type_is_signed(type) ? IR_IMUL : IR_MUL;
        case BINOP_DIV:
            return amal_default_type_is_signed(type) ? IR_IDIV : IR_DIV;
        case BINOP_DOT:
            assert(bool_false && "Binop dot not valid for arithmetic operation and requires special functionality");
            return 0;
        case BINOP_EQUALS:
            return IR_EQUALS;
        case BINOP_NOT_EQUAL:
            return IR_NOT_EQUAL;
        case BINOP_AND:
            return IR_AND;
        case BINOP_LESS:
            return amal_default_type_is_signed(type) ? IR_ILT : IR_LT;
        case BINOP_LESS_EQUAL:
            return amal_default_type_is_signed(type) ? IR_ILE : IR_LE;
        case BINOP_GREATER:
            return amal_default_type_is_signed(type) ? IR_IGT : IR_GT;
        case BINOP_GREATER_EQUAL:
            return amal_default_type_is_signed(type) ? IR_IGE : IR_GE;
    }
    return 0;
}

/* Returns the import statement for lhs of binop dot expression, where lhs is a variable name */
static Import* binop_lhs_get_import_or_null(Binop *self) {
    if(self->lhs->type == AST_VARIABLE) {
        ScopeNamedObject *resolved_var = &self->lhs->value.variable->resolved_var;
        if(resolved_var->type == NAMED_OBJECT_LHS_EXPR && resolved_var->value.lhs_expr->rhs_expr && resolved_var->value.lhs_expr->rhs_expr->type == AST_IMPORT)
            return resolved_var->value.lhs_expr->rhs_expr->value.import;
    }
    return NULL;
}

static CHECK_RESULT IrRegister binop_generate_ir(Binop *self, IrCompilerContext *context) {
    IrRegister reg;

    /*
        Syntax example for binop dot + func_decl expr
        const std = @import("std.amal");
        std.printf
    */
    if(self->type == BINOP_DOT && self->rhs->type == AST_FUNCTION_CALL) {
        Import *lhs_import = binop_lhs_get_import_or_null(self);
        if(lhs_import)
            context->import_index = 1 + lhs_import->file_scope->import_index;
        reg = ast_generate_ir(self->rhs, context);
        context->import_index = 0;
    } else {
        const IrRegister lhs_reg = ast_generate_ir(self->lhs, context);
        const IrRegister rhs_reg = ast_generate_ir(self->rhs, context);
        assert(self->lhs->resolve_data.type.type == RESOLVED_TYPE_LHS_EXPR && "TODO: Implement binop_generate_ir for function signature");
        throw_if_error(ir_ins_binop(context->ir, binop_type_to_ir_type(self->type, (amal_default_type*)self->lhs->resolve_data.type.value.lhs_expr), lhs_reg, rhs_reg, &reg));
    }
    return reg;
}

static void else_if_statement_generate_ir(ElseIfStatement *else_if_stmt, IrCompilerContext *context, IrLabelIndex *skip_other_else_statements_label) {
    if(else_if_stmt->condition) {
        usize jump_ins_index;
        usize jump_skip_else_index;
        IrLabelIndex skip_body_label;
        IrRegister condition_reg =  ast_generate_ir(else_if_stmt->condition, context);
        throw_if_error(ir_ins_jumpzero(context->ir, condition_reg, 0, &jump_ins_index));
        scope_generate_ir(&else_if_stmt->body, context);
        if(else_if_stmt->next_else_if_stmt)
            throw_if_error(ir_ins_jump(context->ir, 0, &jump_skip_else_index));

        throw_if_error(ir_ins_label(context->ir, &skip_body_label));
        throw_if_error(ir_set_jump_label(context->ir, jump_ins_index, skip_body_label));
        if(else_if_stmt->next_else_if_stmt) {
            else_if_statement_generate_ir(else_if_stmt->next_else_if_stmt, context, skip_other_else_statements_label);
            /* Skip over all other else if statements, since else_if_statement_generate_ir is recursive */
            throw_if_error(ir_set_jump_label(context->ir, jump_skip_else_index, *skip_other_else_statements_label));
            return;
        }
    } else {
        assert(!else_if_stmt->next_else_if_stmt);
        scope_generate_ir(&else_if_stmt->body, context);
    }

    /* Note: The last else if statement doesn't need a jump */
    throw_if_error(ir_ins_label(context->ir, skip_other_else_statements_label));
}

static void if_statement_generate_ir(IfStatement *if_stmt, IrCompilerContext *context) {
    usize jump_ins_index;
    usize jump_skip_else_index;
    IrLabelIndex skip_body_label;
    IrLabelIndex skip_else_statements_label;

    IrRegister condition_reg = ast_generate_ir(if_stmt->condition, context);
    throw_if_error(ir_ins_jumpzero(context->ir, condition_reg, 0, &jump_ins_index));
    scope_generate_ir(&if_stmt->body, context);
    if(if_stmt->else_if_stmt)
        throw_if_error(ir_ins_jump(context->ir, 0, &jump_skip_else_index));

    throw_if_error(ir_ins_label(context->ir, &skip_body_label));
    throw_if_error(ir_set_jump_label(context->ir, jump_ins_index, skip_body_label));
    if(if_stmt->else_if_stmt) {
        else_if_statement_generate_ir(if_stmt->else_if_stmt, context, &skip_else_statements_label);
        /*
            Skip over all else if statements, since else_if_statement_generate_ir is recursive.
            We want to jump since we want to skip the else if statements if we are inside the first if-statement
        */
        throw_if_error(ir_set_jump_label(context->ir, jump_skip_else_index, skip_else_statements_label));
    }
}

static void while_statement_generate_ir(WhileStatement *while_stmt, IrCompilerContext *context) {
    IrLabelIndex before_condition_label;
    IrLabelIndex skip_body_label;
    usize jump_after_condition_index;
    IrRegister condition_reg;

    throw_if_error(ir_ins_label(context->ir, &before_condition_label));
    condition_reg = ast_generate_ir(while_stmt->condition, context);
    throw_if_error(ir_ins_jumpzero(context->ir, condition_reg, 0, &jump_after_condition_index));
    scope_generate_ir(&while_stmt->body, context);

    throw_if_error(ir_ins_jump(context->ir, before_condition_label, NULL));
    throw_if_error(ir_ins_label(context->ir, &skip_body_label));
    throw_if_error(ir_set_jump_label(context->ir, jump_after_condition_index, skip_body_label));
}

static void return_expr_generate_ir(ReturnExpr *self, IrCompilerContext *context) {
    const IrRegister reg = ast_generate_ir(self->rhs_expr, context);
    throw_if_error(ir_ins_return(context->ir, reg));
}

static CHECK_RESULT IrRegister ast_generate_ir_resolve_data(void *ast_data, AstType ast_type, AstResolveData *resolve_data, IrCompilerContext *context) {
    /*if(resolve_data->status == AST_IR_RESOLVED)
        return resolve_data->ir_reg;*/

    switch(ast_type) {
        case AST_NUMBER:
            resolve_data->ir_reg = number_generate_ir(ast_data, context);
            break;
        case AST_BOOL:
            resolve_data->ir_reg = ast_bool_generate_ir(ast_data, context);
            break;
        case AST_FUNCTION_DECL:
            /* The IR for function declarations is done separately, in @scope_generate_functions_ir */
            resolve_data->ir_reg = funcdecl_ref_generate_ir(ast_data, context);
            break;
        case AST_FUNCTION_CALL:
            resolve_data->ir_reg = funccall_generate_ir(ast_data, context);
            break;
        case AST_STRUCT_DECL:
            resolve_data->ir_reg = 0;/*structdecl_generate_ir(ast_data, context);*/
            break;
        case AST_STRUCT_FIELD:
            assert(bool_false);
            resolve_data->ir_reg = 0;/*structfield_generate_ir(ast_data, context);*/
            break;
        case AST_LHS:
            resolve_data->ir_reg = lhsexpr_generate_ir(ast_data, resolve_data, context);
            break;
        case AST_ASSIGN:
            resolve_data->ir_reg = assignmentexpr_generate_ir(ast_data, context);
            break;
        case AST_IMPORT:
            /* TODO: Implement cross file references */
            resolve_data->ir_reg = 0;
            break;
        case AST_STRING:
            resolve_data->ir_reg = string_generate_ir(ast_data, context);
            break;
        case AST_VARIABLE:
            resolve_data->ir_reg = variable_generate_ir(ast_data, context);
            break;
        case AST_BINOP:
            resolve_data->ir_reg = binop_generate_ir(ast_data, context);
            break;
        case AST_IF_STATEMENT:
            if_statement_generate_ir(ast_data, context);
            break;
        case AST_WHILE_STATEMENT:
            while_statement_generate_ir(ast_data, context);
            break;
        case AST_RETURN:
            return_expr_generate_ir(ast_data, context);
            break;
    }

    resolve_data->status = AST_IR_RESOLVED;
    return resolve_data->ir_reg;
}

CHECK_RESULT IrRegister ast_generate_ir(Ast *self, IrCompilerContext *context) {
#ifdef DEBUG
    assert(self);
    if(self->resolve_data.status != AST_RESOLVED && self->resolve_data.status != AST_IR_RESOLVED) {
        amal_log_error("Ast type not resolved: %d", self->type);
        assert(bool_false);
    }
#endif
    return ast_generate_ir_resolve_data(self->value.data, self->type, &self->resolve_data, context);
}

CHECK_RESULT IrRegister scope_named_object_generate_ir(ScopeNamedObject *self, IrCompilerContext *context) {
    switch(self->type) {
        case NAMED_OBJECT_NONE:
            assert(bool_false);
            return 0;
        case NAMED_OBJECT_LHS_EXPR:
            return ast_generate_ir_resolve_data(self->value.lhs_expr, AST_LHS, self->resolve_data, context);
        case NAMED_OBJECT_FUNC_PARAM:
            return function_parameter_generate_ir(self->value.func_param, context);
    }
    return 0;
}

void scope_generate_ir(Scope *self, IrCompilerContext *context) {
    Ast **ast = buffer_begin(&self->ast_objects);
    Ast **ast_end = buffer_end(&self->ast_objects);
    for(; ast != ast_end; ++ast) {
        ignore_result_int(ast_generate_ir(*ast, context));
    }
}

void scope_generate_function_ids(Scope *self, IrCompilerContext *context) {
    Ast **ast = buffer_begin(&self->ast_objects);
    Ast **ast_end = buffer_end(&self->ast_objects);
    for(; ast != ast_end; ++ast) {
        if((*ast)->type == AST_LHS && (*ast)->value.lhs_expr->rhs_expr->type == AST_FUNCTION_DECL) {
            LhsExpr *lhs_expr = (*ast)->value.lhs_expr;
            FunctionDecl *func_decl = lhs_expr->rhs_expr->value.func_decl;
            /*
                Going depth first will optimize scope private closures, so they are declared
                before the function they are defined in. This means calling them wont create a deferred function call
                in program.c
            */
            scope_generate_functions_ir(&func_decl->body, context);
            /* TODO: Should this not be done for extern closures? */
            throw_if_error(ir_inc_func_index(context->ir, &func_decl->ir_func_index));
        }
    }
}

void scope_generate_functions_ir(Scope *self, IrCompilerContext *context) {
    Ast **ast = buffer_begin(&self->ast_objects);
    Ast **ast_end = buffer_end(&self->ast_objects);
    for(; ast != ast_end; ++ast) {
        if((*ast)->type == AST_LHS && (*ast)->value.lhs_expr->rhs_expr->type == AST_FUNCTION_DECL) {
            LhsExpr *lhs_expr = (*ast)->value.lhs_expr;
            FunctionDecl *func_decl = lhs_expr->rhs_expr->value.func_decl;
            /*
                Going depth first will optimize scope private closures, so they are declared
                before the function they are defined in. This means calling them wont create a deferred function call
                in program.c
            */
            scope_generate_functions_ir(&func_decl->body, context);
            ignore_result_int(funcdecl_generate_ir(func_decl, context));
            if(LHS_EXPR_IS_EXPORT(lhs_expr))
                throw_if_error(ir_try_add_export_func(context->ir, func_decl->signature, lhs_expr->var_name));
        }
    }
}