#include "../include/program.h"

#include "../include/std/mem.h"
#include "../include/std/hash.h"
#include "../include/std/alloc.h"
#include "../include/std/log.h"
#include "../include/std/buffer_view.h"
#include <stdio.h>
#include <errno.h>
#include <assert.h>
#include <string.h>

/* One gigabyte */
#define PROGRAM_MAX_SIZE 1024*1024*1024

/* TODO: If system is big-endian, then do endian conversion for all reads */

/* This matches SsaNumberType */
typedef enum {
    NUMBER_TYPE_INTEGER,
    NUMBER_TYPE_FLOAT
} NumberType;

typedef union {
    i64 integer;
    f64 floating;
} NumberUnion;

typedef struct {
    NumberType type;
    NumberUnion value;
} Number;

static int hash_map_compare_u64(const void *a, const void *b) {
    return *(u64*)a == *(u64*)b;
}

static usize hash_u64(const u8 *data, usize size) {
    (void)size;
    return *(u64*)data;
}

int amal_program_init(amal_program *self) {
    ignore_result_int(buffer_init(&self->data, NULL));
    self->string_indices                = NULL;
    self->intermediates_start           = NULL;
    self->strings_start                 = NULL;
    self->extern_funcs_start            = NULL;
    self->exported_funcs                = NULL;
    self->exported_funcs_end            = NULL;
    self->imports_start                 = NULL;
    self->read_index                    = 0;
    self->main_func_instruction_offset  = ~(u32)0U;
    self->num_intermediates             = 0;
    self->num_strings                   = 0;
    self->num_functions                 = 0;
    self->num_extern_functions          = 0;
    self->num_exported_functions        = 0;
    self->num_imports                   = 0;
    self->return_value_index            = 0;

    cleanup_if_error(arena_allocator_init(&self->allocator));
    cleanup_if_error(hash_map_init(&self->extern_funcs_map, &self->allocator, sizeof(ProgramExternFunc), hash_map_compare_string, amal_hash_string));
    cleanup_if_error(hash_map_init(&self->deferred_func_calls, &self->allocator, sizeof(Buffer), hash_map_compare_u64, hash_u64));
    cleanup_if_error(buffer_append_header(&self->data));
    return 0;

    cleanup:
    amal_program_deinit(self);
    return -1;
}

void amal_program_deinit(amal_program *self) {
    arena_allocator_deinit(&self->allocator);
    am_free(self->string_indices);
    self->string_indices = NULL;
    if(self->data.data)
        buffer_deinit(&self->data);
}

int amal_program_register_extern_func(amal_program *self, BufferView name, void *func_ptr, int args_byte_size) {
    ProgramExternFunc extern_func;
    extern_func.func = func_ptr;
    extern_func.args_byte_size = args_byte_size;

    if(hash_map_contains(&self->extern_funcs_map, name))
        return AMAL_PROGRAM_EXTERN_FUNC_ALREADY_EXISTS;

    return hash_map_insert(&self->extern_funcs_map, name, &extern_func);
}

static u8* amal_program_get_extern_funcs_start_by_import_index(amal_program *self, u8 import_index) {
    BytecodeHeaderImport *header_import = (BytecodeHeaderImport*)self->imports_start;
    u32 extern_function_index = 0;
    header_import += import_index;
    assert(sizeof(extern_function_index) == sizeof(header_import->extern_function_index));
    am_memcpy(&extern_function_index, &header_import->extern_function_index, sizeof(extern_function_index));
    return (u8*)self->data.data + extern_function_index;
}

static u32 extern_func_get_total_param_size(BytecodeHeaderExternFunction *self) {
    return self->params_fixed_size + self->params_num_pointers * sizeof(void*);
}

typedef struct {
    BytecodeHeaderExternFunction extern_func;
    BufferView name;
} BytecodeHeaderExternFunctionFull;

/* TODO: Optimize this */
static void amal_program_get_header_extern_function_by_index(amal_program *self, u8 import_index, u16 index, BytecodeHeaderExternFunctionFull *result) {
    u32 i;
    u8 *extern_funcs_start = amal_program_get_extern_funcs_start_by_import_index(self, import_index);
    extern_funcs_start += sizeof(u16) + sizeof(u32);
    for(i = 0; i < (u32)index; ++i) {
        u8 name_len;
        extern_funcs_start += sizeof(BytecodeHeaderExternFunction);
        name_len = *extern_funcs_start;
        extern_funcs_start += 1 + name_len + 1; /* +1 for skipping length byte and the null-terminated character */
    }
    am_memcpy(&result->extern_func, extern_funcs_start, sizeof(result->extern_func));
    result->name.size = *(extern_funcs_start + sizeof(result->extern_func));
    result->name.data = (const char*)extern_funcs_start + sizeof(result->extern_func) + sizeof(u8);
}


static CHECK_RESULT int amal_program_get_extern_func_by_index(amal_program *self, u8 import_index, u16 index, ProgramExternFunc *result) {
    BytecodeHeaderExternFunctionFull extern_func;
    amal_program_get_header_extern_function_by_index(self, import_index, index, &extern_func);
    
    if(!hash_map_get(&self->extern_funcs_map, extern_func.name, result)) {
        amal_log_error("Extern function \"%.*s\" has not been registered", extern_func.name.size, extern_func.name.data);
        return AMAL_PROGRAM_NO_SUCH_EXTERNAL_FUNCTION;
    }
    
    /* TODO: This assumes all arguments are of size sizeof(isize) */
    if(result->args_byte_size != -1 && result->args_byte_size != (int)extern_func_get_total_param_size(&extern_func.extern_func)) {
        amal_log_error("Extern function %.*s was registered to take %d byte(s), but the program says it takes %d byte(s)",
            extern_func.name.size, extern_func.name.data, result->args_byte_size, (int)extern_func_get_total_param_size(&extern_func.extern_func));
        return AMAL_PROGRAM_NO_SUCH_EXTERNAL_FUNCTION;
    }
    return 0;
}

static CHECK_RESULT int amal_program_set_exported_function_instruction_offset_advance(amal_program *self, u32 instruction_offset) {
    if(self->exported_funcs >= self->exported_funcs_end) {
        amal_log_error("The number of exported functions in the instructions is more than the number of exported instructions in the header");
        return AMAL_PROGRAM_INSTRUCTION_INVALID_EXPORTED_FUNC_INDEX;
    }

    {
        u8 num_args;
        u8 func_name_size;

        if((usize)(self->exported_funcs_end - self->exported_funcs) < sizeof(instruction_offset) + sizeof(num_args) + sizeof(func_name_size))
            return AMAL_PROGRAM_INSTRUCTION_INVALID_EXPORTED_FUNC_INDEX;

        am_memcpy(self->exported_funcs, &instruction_offset, sizeof(instruction_offset));
        num_args = self->exported_funcs[sizeof(instruction_offset)];
        func_name_size = self->exported_funcs[sizeof(instruction_offset) + sizeof(num_args)];
        self->exported_funcs += sizeof(instruction_offset) + sizeof(num_args) + sizeof(func_name_size);
        if(self->main_func_instruction_offset == ~(u32)0U && func_name_size == 4 && am_memeql(self->exported_funcs, "main", 4))
            self->main_func_instruction_offset = instruction_offset;

        /* +1 to skip null-termination character */
        if((usize)(self->exported_funcs_end - self->exported_funcs) < func_name_size + 1U)
            return AMAL_PROGRAM_INSTRUCTION_INVALID_EXPORTED_FUNC_INDEX;

        self->exported_funcs += func_name_size + 1; /* +1 to skip null-termination character */
    }

    return 0;
}

int amal_program_append_bytecode(amal_program *self, Bytecode *bytecode) {
    /* Sanity check for indices used later. A program shouldn't be more than 1gb */
    if(self->data.size + bytecode->data.size > PROGRAM_MAX_SIZE) {
        amal_log_error("Program is too large. Max size is 1GB");
        return -1;
    }
    return buffer_append(&self->data, bytecode->data.data, bytecode->data.size);
}

static usize bytes_left_to_read(amal_program *self) {
    assert(self->read_index <= self->data.size);
    return self->data.size - self->read_index;
}

static CHECK_RESULT int amal_program_read_header(amal_program *self) {
    BytecodeHeader header;

    if(bytes_left_to_read(self) < sizeof(header))
        return AMAL_PROGRAM_INVALID_HEADER;

    am_memcpy(&header, self->data.data + self->read_index, sizeof(header));
    self->read_index += sizeof(header);
#ifdef AMAL_BIG_ENDIAN
    #error TODO: Support big endian for program decoding
#endif

    if(memcmp(header.magic_number, AMAL_BYTECODE_MAGIC_NUMBER, AMAL_BYTECODE_MAGIC_NUMBER_SIZE) != 0)
        return AMAL_PROGRAM_INVALID_MAGIC_NUMBER;

    /*
        A program is only incompatible if the major version is newer than the version that is used to run it.
        TODO: Implement backwards compatible reads, starting from when the program bytecode breaks backwards compatibility
    */
    if(header.major_version > AMAL_BYTECODE_MAJOR_VERSION)
        return AMAL_PROGRAM_INCOMPATIBLE;

    return AMAL_PROGRAM_OK;
}

static CHECK_RESULT bool amal_program_read_advance(amal_program *self, void *output, usize bytes_to_read) {
    if(bytes_left_to_read(self) < bytes_to_read)
        return bool_false;
    am_memcpy(output, self->data.data + self->read_index, bytes_to_read);
    self->read_index += bytes_to_read;
    return bool_true;
}

static CHECK_RESULT int amal_program_advance_section_magic_number(amal_program *self) {
    u32 magic_number;
    if(bytes_left_to_read(self) < sizeof(u32))
        return AMAL_PROGRAM_SECTION_ERROR;
    
    am_memcpy(&magic_number, self->data.data + self->read_index, sizeof(magic_number));
    if(magic_number != AMAL_BYTECODE_SECTION_MAGIC_NUMBER)
        return AMAL_PROGRAM_SECTION_ERROR;
    
    self->read_index += sizeof(magic_number);
    return 0;
}

static CHECK_RESULT int amal_program_read_intermediates(amal_program *self) {
    u32 intermediates_size;

    if(bytes_left_to_read(self) < sizeof(intermediates_size)) {
        amal_log_error("Not enough space in program to intermediates size");
        return AMAL_PROGRAM_INVALID_INTERMEDIATES;
    }

    am_memcpy(&intermediates_size, &self->data.data[self->read_index], sizeof(intermediates_size));
    self->read_index += sizeof(intermediates_size);

    if(bytes_left_to_read(self) < intermediates_size) {
        amal_log_error("Not enough space in program to read all intermediates");
        return AMAL_PROGRAM_INVALID_INTERMEDIATES_SIZE;
    }

    self->intermediates_start = (u8*)(self->data.data + self->read_index);
    self->num_intermediates = intermediates_size / (sizeof(u8) + sizeof(u64));
    self->read_index += intermediates_size;

    return AMAL_PROGRAM_OK;
}

static CHECK_RESULT int amal_program_read_strings(amal_program *self) {
    u32 strings_size;
    u32 *string_index_ptr;

    if(!amal_program_read_advance(self, &self->num_strings, sizeof(u16)))
        return AMAL_PROGRAM_INVALID_STRINGS;

    if(!amal_program_read_advance(self, &strings_size, sizeof(strings_size)))
        return AMAL_PROGRAM_INVALID_STRINGS;

    if(bytes_left_to_read(self) < strings_size)
        return AMAL_PROGRAM_INVALID_STRINGS_SIZE;

    am_free(self->string_indices);
    self->string_indices = NULL;
    if(am_malloc(sizeof(u32) * self->num_strings, (void**)&self->string_indices) != 0)
        return AMAL_PROGRAM_ALLOC_FAILURE;
    string_index_ptr = self->string_indices;

    {
        const u32 read_start = self->read_index;
        const u32 read_end = read_start + strings_size;
        self->strings_start = (u8*)(self->data.data + self->read_index);
        while(self->read_index < read_end) {
            u16 string_size;

            if(bytes_left_to_read(self) < sizeof(string_size))
                return AMAL_PROGRAM_INVALID_STRINGS;

            *string_index_ptr = self->read_index - read_start;
            ++string_index_ptr;
            am_memcpy(&string_size, &self->data.data[self->read_index], sizeof(string_size));
            self->read_index += sizeof(string_size);

            /* +1 to skip null-termination character */
            if(bytes_left_to_read(self) < string_size + 1U)
                return AMAL_PROGRAM_INVALID_STRINGS;

            self->read_index += string_size + 1; /* +1 to skip null-termination character */
        }
        assert(self->read_index == read_end);
    }

    return AMAL_PROGRAM_OK;
}

static CHECK_RESULT int amal_program_read_functions(amal_program *self) {
    u32 funcs_size;

    if(!amal_program_read_advance(self, &self->num_functions, sizeof(u16)))
        return AMAL_PROGRAM_INVALID_FUNCTIONS;

    if(!amal_program_read_advance(self, &funcs_size, sizeof(funcs_size)) || bytes_left_to_read(self) < funcs_size)
        return AMAL_PROGRAM_INVALID_FUNCTIONS;

    self->funcs_start = (u8*)(self->data.data + self->read_index);
    self->read_index += funcs_size;
    return AMAL_PROGRAM_OK;
}

static CHECK_RESULT int amal_program_read_external_functions(amal_program *self) {
    u32 extern_funcs_size;

    if(!amal_program_read_advance(self, &self->num_extern_functions, sizeof(u16)))
        return AMAL_PROGRAM_INVALID_EXTERNAL_FUNCTIONS;

    if(!amal_program_read_advance(self, &extern_funcs_size, sizeof(extern_funcs_size)))
        return AMAL_PROGRAM_INVALID_EXTERNAL_FUNCTIONS;

    if(bytes_left_to_read(self) < extern_funcs_size)
        return AMAL_PROGRAM_INVALID_EXTERNAL_FUNCTIONS_SIZE;

    self->read_index += extern_funcs_size;
    return AMAL_PROGRAM_OK;
}

static CHECK_RESULT int amal_program_read_exported_functions(amal_program *self) {
    u32 export_funcs_size;

    if(!amal_program_read_advance(self, &self->num_exported_functions, sizeof(u16)))
        return AMAL_PROGRAM_INVALID_EXPORTED_FUNCTIONS;

    if(!amal_program_read_advance(self, &export_funcs_size, sizeof(export_funcs_size)))
        return AMAL_PROGRAM_INVALID_EXPORTED_FUNCTIONS;

    if(bytes_left_to_read(self) < export_funcs_size)
        return AMAL_PROGRAM_INVALID_EXPORTED_FUNCTIONS_SIZE;

    /*
        Exported functions doesn't need list of indices to the data, since exported functions
        are always sorted in the instructions in the same order as list of exported functions
    */
    self->exported_funcs = (u8*)(self->data.data + self->read_index);
    self->exported_funcs_end = self->exported_funcs + export_funcs_size;
    self->read_index += export_funcs_size;
    return AMAL_PROGRAM_OK;
}

static CHECK_RESULT int amal_program_read_imports(amal_program *self) {
    u32 imports_size;

    if(!amal_program_read_advance(self, &self->num_imports, sizeof(u8)))
        return AMAL_PROGRAM_INVALID_IMPORTS;

    if(!amal_program_read_advance(self, &imports_size, sizeof(imports_size)) || bytes_left_to_read(self) < imports_size)
        return AMAL_PROGRAM_INVALID_IMPORTS;

    self->imports_start = (u8*)(self->data.data + self->read_index);
    self->read_index += imports_size;
    return AMAL_PROGRAM_OK;
}

static CHECK_RESULT int amal_program_get_intermediate_by_index(amal_program *self, u16 index, Number *result) {
    if(index >= self->num_intermediates)
        return AMAL_PROGRAM_INSTRUCTION_INVALID_INTERMEDIATE_INDEX;
    
    am_memcpy(&result->type, &self->intermediates_start[(sizeof(u8) + sizeof(u64)) * (usize)index], sizeof(u8));
    am_memcpy(&result->value, &self->intermediates_start[(sizeof(u8) + sizeof(u64)) * (usize)index + sizeof(u8)], sizeof(u64));
    return 0;
}

static CHECK_RESULT int amal_program_get_data_by_index(amal_program *self, u16 index, BufferView *result) {
    u8 *str_ptr;

    if(index >= self->num_strings) {
        amal_log_error("Data index %ld is out of range (%ld)", index, self->num_strings);
        return AMAL_PROGRAM_INSTRUCTION_INVALID_DATA_INDEX;
    }
    
    str_ptr = self->strings_start + self->string_indices[index];
    am_memcpy(&result->size, str_ptr, sizeof(u16));
    result->data = (const char*)(str_ptr + sizeof(u16));
    return 0;
}

static u8* amal_program_get_funcs_start_by_import_index(amal_program *self, u8 import_index) {
    BytecodeHeaderImport *header_import = (BytecodeHeaderImport*)self->imports_start;
    u32 function_index = 0;
    header_import += import_index;
    assert(sizeof(function_index) == sizeof(header_import->function_index));
    am_memcpy(&function_index, &header_import->function_index, sizeof(function_index));
    return (u8*)self->data.data + function_index;
}

static void amal_program_get_header_function_by_index(amal_program *self, u8 import_index, u16 index, BytecodeHeaderFunction *result) {
    u8 *funcs_start = amal_program_get_funcs_start_by_import_index(self, import_index);
    BytecodeHeaderFunction *header_func = (BytecodeHeaderFunction*)(funcs_start + sizeof(u16) + sizeof(u32));
    am_memcpy(result, header_func + index, sizeof(BytecodeHeaderFunction));
}

static u64 deferred_func_call_get_key(amal_program *self, u8 import_index, u16 func_index) {
    BytecodeHeaderImport *header_import = (BytecodeHeaderImport*)self->imports_start;
    u32 function_index = 0;
    header_import += import_index;
    assert(sizeof(function_index) == sizeof(header_import->function_index));
    am_memcpy(&function_index, &header_import->function_index, sizeof(function_index));
    return ((u64)func_index << 32) | (u64)function_index;
}

static CHECK_RESULT int resolve_deferred_func_calls(amal_program *self, amal_executor *executor, u16 func_index) {
    u64 key = deferred_func_call_get_key(self, 0, func_index);
    BufferView key_mem = create_buffer_view((char*)&key, sizeof(key));
    u32 current_code_offset = amal_exec_get_code_offset(executor);

    Buffer* deferred_func_call_list;
    if(!hash_map_get_ref(&self->deferred_func_calls, key_mem, (void**)&deferred_func_call_list))
        return 0;
    
    {
        u32 *code_offset = buffer_begin(deferred_func_call_list);
        u32 *code_offset_end = buffer_end(deferred_func_call_list);
        for(; code_offset != code_offset_end; ++code_offset) {
            amal_exec_call_overwrite(executor, *code_offset, current_code_offset - *code_offset);
        }
        return buffer_set_capacity(deferred_func_call_list, 0);
    }
}

static void header_func_set_offset(amal_program *self, u16 func_index, u32 code_offset) {
    BytecodeHeaderFunction *header_func = ((BytecodeHeaderFunction*)self->funcs_start) + func_index;
    assert(sizeof(header_func->func_offset) == sizeof(code_offset));
    am_memcpy(&header_func->func_offset, &code_offset, sizeof(code_offset));
}

static CHECK_RESULT int amal_program_read_instructions(amal_program *self, amal_executor *executor) {
    u32 instructions_size;
    u32 read_start;
    u32 read_end;
    u16 func_counter;
    bool inside_func;

    inside_func = bool_false;
    (void)inside_func;
    func_counter = 0;
    self->return_value_index = 0;

    if(!amal_program_read_advance(self, &instructions_size, sizeof(instructions_size)))
        return AMAL_PROGRAM_INVALID_INSTRUCTIONS_SIZE;

    if(bytes_left_to_read(self) < instructions_size)
        return AMAL_PROGRAM_INVALID_INSTRUCTIONS_SIZE;

    /*
        TODO: self->reg should be of type Number and each arithmetic operation should operate
        on the type of the register.
    */

    return_if_error(amal_executor_instructions_start(executor, self->num_functions));

    read_start = self->read_index;
    read_end = read_start + instructions_size;
    while(self->read_index < read_end) {
        AmalOpcode opcode;
        opcode = self->data.data[self->read_index];
        self->read_index += sizeof(AmalOpcodeType);
        /* TODO: Check instruction length and check if that amount of bytes can be read */
        switch(opcode) {
            case AMAL_OP_NOP: {
                return_if_error(amal_exec_nop(executor));
                break;
            }
            case AMAL_OP_SETZ: {
                return_if_error(amal_exec_setz(executor, self->data.data[self->read_index]));
                self->read_index += 1;
                break;
            }
            case AMAL_OP_MOV: {
                return_if_error(amal_exec_mov(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1]));
                self->read_index += 2;
                break;
            }
            case AMAL_OP_MOVI: {
                u16 intermediate_index;
                Number number;

                am_memcpy(&intermediate_index, &self->data.data[self->read_index + sizeof(i8)], sizeof(intermediate_index));
                return_if_error(amal_program_get_intermediate_by_index(self, intermediate_index, &number));
                return_if_error(amal_exec_movi(executor, self->data.data[self->read_index], number.value.integer));

                self->read_index += 3;
                break;
            }
            case AMAL_OP_MOVD: {
                u16 data_index;
                BufferView data_ptr;

                am_memcpy(&data_index, &self->data.data[self->read_index + sizeof(i8)], sizeof(data_index));

                return_if_error(amal_program_get_data_by_index(self, data_index, &data_ptr));
                return_if_error(amal_exec_movd(executor, self->data.data[self->read_index], data_ptr));

                self->read_index += 3;
                break;
            }
            case AMAL_OP_ADD: {
                return_if_error(amal_exec_add(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1], self->data.data[self->read_index + 2]));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_SUB: {
                return_if_error(amal_exec_sub(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1], self->data.data[self->read_index + 2]));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_IMUL: {
                return_if_error(amal_exec_imul(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1], self->data.data[self->read_index + 2]));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_MUL: {
                return_if_error(amal_exec_mul(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1], self->data.data[self->read_index + 2]));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_IDIV: {
                return_if_error(amal_exec_idiv(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1], self->data.data[self->read_index + 2]));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_DIV: {
                return_if_error(amal_exec_div(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1], self->data.data[self->read_index + 2]));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_PUSH: {
                return_if_error(amal_exec_push(executor, self->data.data[self->read_index]));
                self->read_index += 1;
                break;
            }
            case AMAL_OP_PUSHI: {
                u16 intermediate_index;
                Number number;

                am_memcpy(&intermediate_index, &self->data.data[self->read_index], sizeof(intermediate_index));
                return_if_error(amal_program_get_intermediate_by_index(self, intermediate_index, &number));
                return_if_error(amal_exec_pushi(executor, number.value.integer));
                self->read_index += 2;
                break;
            }
            case AMAL_OP_PUSHD: {
                u16 data_index;
                BufferView data_ptr;

                am_memcpy(&data_index, &self->data.data[self->read_index], sizeof(data_index));
                return_if_error(amal_program_get_data_by_index(self, data_index, &data_ptr));
                return_if_error(amal_exec_pushd(executor, data_ptr));
                self->read_index += 2;
                break;
            }
            case AMAL_OP_PUSH_RET: {
                /* TODO: Validate return value index doesn't go out of bounds? */
                self->return_values_stack[self->return_value_index++] = self->data.data[self->read_index];
                self->read_index += 1;
                break;
            }
            case AMAL_OP_CALL_START: {
                u8 num_args = self->data.data[self->read_index];
                return_if_error(amal_exec_call_start(executor, num_args));
                self->read_index += 1;
                break;
            }
            case AMAL_OP_CALL: {
                u8 import_index;
                u16 func_index;
                BytecodeHeaderFunction func_def;
                i8 dst_reg;

                am_memcpy(&import_index, self->data.data + self->read_index, sizeof(import_index));
                am_memcpy(&func_index, self->data.data + self->read_index + sizeof(import_index), sizeof(func_index));
                amal_program_get_header_function_by_index(self, import_index, func_index, &func_def);
                assert(func_def.num_return_types == 1 && "TODO: Support 0 and more than 1 return values");
                assert(self->return_value_index == 1);
                dst_reg = self->return_values_stack[0];
                self->return_value_index = 0;

                /* func_offset will only be non-zero when the function has been decoded (FUNC_START) */
                if(func_def.func_offset != 0) {
                    /* TODO: Instead of pushing num args, push the sum of sizeof the last num_args */
                    return_if_error(amal_exec_call(executor, func_def.func_offset, dst_reg));
                } else {
                    /*
                        The code for the function has not been generated yet (the function is defined after the current location).
                        Make a dummy call and replace the call target after the function has been generated
                    */
                    u64 key = deferred_func_call_get_key(self, import_index, func_index);
                    BufferView key_mem = create_buffer_view((char*)&key, sizeof(key));
                    u32 code_offset = amal_exec_get_code_offset(executor);

                    Buffer* deferred_func_call_list;
                    if(hash_map_get_ref(&self->deferred_func_calls, key_mem, (void**)&deferred_func_call_list))
                        return_if_error(buffer_append(deferred_func_call_list, &code_offset, sizeof(code_offset)));
                    else {
                        Buffer new_deferred_call_list;
                        return_if_error(buffer_init(&new_deferred_call_list, &self->allocator));
                        return_if_error(buffer_append(&new_deferred_call_list, &code_offset, sizeof(code_offset)));
                        return_if_error(hash_map_insert(&self->deferred_func_calls, key_mem, &new_deferred_call_list));
                    }
                    /* Dummy call to offset 0, offset will be replace later when the target function hits AMAL_OP_FUNC_START */
                    return_if_error(amal_exec_call(executor, 0, dst_reg));
                }
                self->read_index += 3;
                break;
            }
            case AMAL_OP_CALLR: {
                assert(bool_false && "TODO: Implement CALLR");
                self->read_index += 1;
                break;
            }
            case AMAL_OP_CALLE: {
                u8 import_index;
                u16 extern_func_index;
                i8 dst_reg;

                am_memcpy(&import_index, self->data.data + self->read_index, sizeof(import_index));
                am_memcpy(&extern_func_index, self->data.data + self->read_index + sizeof(import_index), sizeof(extern_func_index));
                assert(self->return_value_index == 1 && "TODO: Support extern functions that don't return any value");
                dst_reg = self->return_values_stack[0];
                self->return_value_index = 0;

                {
                    ProgramExternFunc extern_func;
                    return_if_error(amal_program_get_extern_func_by_index(self, import_index, extern_func_index, &extern_func));
                    return_if_error(amal_exec_calle(executor, extern_func.func, dst_reg));
                }
                self->read_index += 3;
                break;
            }
            case AMAL_OP_CMP: {
                return_if_error(amal_exec_cmp(executor, self->data.data[self->read_index], self->data.data[self->read_index + 1], self->data.data[self->read_index + 2]));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_JZ: {
                i8 reg;
                u16 target_label;
                reg = self->data.data[self->read_index];
                am_memcpy(&target_label, self->data.data + self->read_index + sizeof(reg), sizeof(target_label));
                return_if_error(amal_exec_jz(executor, reg, target_label));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_JMP: {
                u16 target_label;
                am_memcpy(&target_label, self->data.data + self->read_index, sizeof(target_label));
                return_if_error(amal_exec_jmp(executor, target_label));
                self->read_index += 2;
                break;
            }
            case AMAL_OP_RET: {
                const i8 reg = self->data.data[self->read_index];
                return_if_error(amal_exec_ret(executor, reg));
                self->read_index += 1;
                break;
            }
            case AMAL_OP_FUNC_START: {
                u8 func_flags;
                u16 func_num_local_var_regs;

                assert(!inside_func);
                inside_func = bool_true;
                assert(func_counter < self->num_functions);

                header_func_set_offset(self, func_counter, amal_exec_get_code_offset(executor));
                return_if_error(resolve_deferred_func_calls(self, executor, func_counter));
                ++func_counter;

                func_flags = self->data.data[self->read_index];
                am_memcpy(&func_num_local_var_regs, self->data.data + self->read_index + sizeof(func_flags), sizeof(func_num_local_var_regs));
                if(func_flags & FUNC_FLAG_EXPORTED)
                    return_if_error(amal_program_set_exported_function_instruction_offset_advance(self, amal_exec_get_code_offset(executor)));
                return_if_error(amal_exec_func_start(executor, func_num_local_var_regs));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_FUNC_END: {
                assert(inside_func);
                inside_func = bool_false;
                /* TODO: Validate FUNC_END is called for every FUNC_START, otherwise stack will be corrupted */
                return_if_error(amal_exec_func_end(executor));
                break;
            }
            case AMAL_OP_LABEL: {
                return_if_error(amal_exec_label(executor));
                break;
            }
        }
    }

    assert(self->read_index == read_end);
    return_if_error(amal_executor_instructions_end(executor));
    return AMAL_PROGRAM_OK;
}

int amal_program_run(amal_program *self) {
    int result;
    amal_executor *executor;

    if(self->data.size > PROGRAM_MAX_SIZE) {
        amal_log_error("Program is too large. Max size is 1GB");
        return AMAL_PROGRAM_ERR;
    }
    result = AMAL_PROGRAM_ERR;
    return_if_error(amal_executor_init(&executor));

    cleanup_if_error(amal_program_read_header(self));
    while(bytes_left_to_read(self) > 0) {
        cleanup_if_error(amal_program_advance_section_magic_number(self));
        cleanup_if_error(amal_program_read_intermediates(self));

        cleanup_if_error(amal_program_advance_section_magic_number(self));
        cleanup_if_error(amal_program_read_strings(self));

        cleanup_if_error(amal_program_advance_section_magic_number(self));
        cleanup_if_error(amal_program_read_functions(self));

        cleanup_if_error(amal_program_advance_section_magic_number(self));
        cleanup_if_error(amal_program_read_external_functions(self));

        cleanup_if_error(amal_program_advance_section_magic_number(self));
        cleanup_if_error(amal_program_read_exported_functions(self));

        cleanup_if_error(amal_program_advance_section_magic_number(self));
        cleanup_if_error(amal_program_read_imports(self));

        cleanup_if_error(amal_program_advance_section_magic_number(self));
        cleanup_if_error(amal_program_read_instructions(self, executor));
    }
    if(self->main_func_instruction_offset == ~(u32)0U) {
        amal_log_error("The program is missing a main function");
        result = AMAL_PROGRAM_NO_MAIN_FUNC;
        goto cleanup;
    }
    result = amal_executor_run(executor, self->main_func_instruction_offset);

    cleanup:
    amal_executor_deinit(executor);
    return result;
}

int amal_program_save(amal_program *self, const char *filepath) {
    int err = 0;
    FILE *file = fopen(filepath, "wb");
    if(!file) {
        err = -errno;
        goto cleanup;
    }

    if(fwrite(self->data.data, 1, self->data.size, file) != self->data.size) {
        err = -errno;
        goto cleanup;
    }

    cleanup:
    if(err != 0)
        perror(filepath);
    if(file)
        fclose(file);
    return err;
}