#include "../include/program.h"
#include "../include/std/mem.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>

/* 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 CHECK_RESULT int amal_program_append_header(amal_program *self) {
    /*doc(Bytecode header)
        # Header layout
        |Type|Field        |Description                                                                 |
        |----|-------------|----------------------------------------------------------------------------|
        |u32 |Magic number |The magic number used to identify an amalgam bytecode file.                 |
        |u8  |Major version|The major version of the bytecode. Updates in this is a breaking change.    |
        |u8  |Minor version|The minor version of the bytecode. Updates in this are backwards compatible.|
        |u8  |Patch version|The patch version of the bytecode. Updates in this are only minor bug fixes.|

        The versions in the header only changes for every release, not every change.
    */

    const u32 magic_number = AMAL_PROGRAM_MAGIC_NUMBER;
    const u8 major_version = AMAL_PROGRAM_MAJOR_VERSION;
    const u8 minor_version = AMAL_PROGRAM_MINOR_VERSION;
    const u8 patch_version = AMAL_PROGRAM_PATCH_VERSION;

    return_if_error(buffer_append(&self->data, &magic_number, 4));
    return_if_error(buffer_append(&self->data, &major_version, 1));
    return_if_error(buffer_append(&self->data, &minor_version, 1));
    return_if_error(buffer_append(&self->data, &patch_version, 1));

    return 0;
}

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->read_index = 0;
    self->num_strings = 0;
    self->num_intermediates = 0;

    cleanup_if_error(amal_program_append_header(self));
    return 0;

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

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

int amal_program_append_bytecode(amal_program *self, Bytecode *bytecode) {
    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) {
    u32 magic_number;
    u8 major_version;
    u8 minor_version;
    u8 patch_version;

    if(bytes_left_to_read(self) < sizeof(u32) + sizeof(u8) * 3)
        return AMAL_PROGRAM_INVALID_HEADER;

    am_memcpy(&magic_number, &self->data.data[self->read_index], sizeof(magic_number));
    self->read_index += sizeof(u32);
    am_memcpy(&major_version, &self->data.data[self->read_index], sizeof(major_version));
    self->read_index += sizeof(u8);
    am_memcpy(&minor_version, &self->data.data[self->read_index], sizeof(minor_version));
    self->read_index += sizeof(u8);
    am_memcpy(&patch_version, &self->data.data[self->read_index], sizeof(patch_version));
    self->read_index += sizeof(u8);

    if(magic_number != AMAL_PROGRAM_MAGIC_NUMBER)
        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(major_version > AMAL_PROGRAM_MAJOR_VERSION)
        return AMAL_PROGRAM_INCOMPATIBLE;

    return AMAL_PROGRAM_OK;
}

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

    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 = &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) {
    u16 num_strings;
    u32 strings_size;
    u32 read_start;
    u32 read_end;
    u32 *string_index_ptr;

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

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

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

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

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

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

    read_start = self->read_index;
    read_end = read_start + strings_size;
    self->strings_start = &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);

        if(bytes_left_to_read(self) < string_size)
            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_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) {
    char *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 = str_ptr + sizeof(u16);
    return 0;
}

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

    inside_func = bool_false;
    (void)inside_func;

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

    am_memcpy(&instructions_size, &self->data.data[self->read_index], sizeof(instructions_size));
    self->read_index += sizeof(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.
    */

    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);
        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(u8)], 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(u8)], 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_CALL:
                /*assert(bool_false && "TODO: Implement CALL");*/
                self->read_index += 3;
                break;
            case AMAL_OP_CALLR:
                /*assert(bool_false && "TODO: Implement CALLR");*/
                self->read_index += 2;
                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: {
                i16 jump_offset;
                am_memcpy(&jump_offset, &self->data.data[self->read_index + sizeof(u8)], sizeof(jump_offset));
                return_if_error(amal_exec_jz(executor, self->data.data[self->read_index], jump_offset));
                self->read_index += 3;
                break;
            }
            case AMAL_OP_JMP: {
                i16 jump_offset;
                am_memcpy(&jump_offset, &self->data.data[self->read_index], sizeof(jump_offset));
                return_if_error(amal_exec_jmp(executor, jump_offset));
                self->read_index += 2;
                break;
            }
            case AMAL_OP_RET:
                return_if_error(amal_exec_ret(executor));
                break;
            case AMAL_OP_FUNC_START: {
                u16 func_num_registers;
                assert(!inside_func);
                inside_func = bool_true;
                am_memcpy(&func_num_registers, &self->data.data[self->read_index], sizeof(func_num_registers));
                return_if_error(amal_exec_func_start(executor, func_num_registers));
                self->read_index += 2;
                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;
            }
        }
    }

    return AMAL_PROGRAM_OK;
}

int amal_program_run(amal_program *self) {
    int result;
    amal_executor *executor;
    result = 0;
    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_read_intermediates(self));
        cleanup_if_error(amal_program_read_strings(self));
        cleanup_if_error(amal_program_read_instructions(self, executor));
    }
    result = amal_executor_run(executor);

    cleanup:
    amal_executor_deinit(executor);
    return result;
}

int amal_program_save(amal_program *self, const char *filepath) {
    FILE *file;
    file = fopen(filepath, "wb");
    if(!file) {
        int err;
        err = errno;
        perror(filepath);
        return -err;
    }
    if(fwrite(self->data.data, 1, self->data.size, file) != self->data.size) {
        int err;
        err = errno;
        perror(filepath);
        return -err;
    }
    fclose(file);
    return 0;
}