#include "../../include/bytecode/bytecode.h"
#include "../../include/std/mem.h"
#include "../../include/ssa/ssa.h"
#include "../../include/parser.h"
#include "../../include/compiler.h"
#include <assert.h>
#include <stdio.h> /* TODO: Remove this */
#include <stdlib.h>
#include <string.h>

#define throw(result) do { throw_debug_msg; longjmp(self->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)

int bytecode_init(Bytecode *self, ScopedAllocator *allocator) {
    return buffer_init(&self->instructions, allocator);
}

static CHECK_RESULT usize ssa_extract_form1(u8 *instruction_data, SsaInsForm1 *result) {
    am_memcpy(&result->lhs, instruction_data, sizeof(result->lhs));
    am_memcpy(&result->rhs, instruction_data + sizeof(result->lhs), sizeof(result->rhs));
    return sizeof(result->lhs) + sizeof(result->rhs);
}

static CHECK_RESULT usize ssa_extract_form2(u8 *instruction_data, SsaInsForm2 *result) {
    am_memcpy(&result->result, instruction_data, sizeof(result->result));
    am_memcpy(&result->lhs, instruction_data + sizeof(result->result), sizeof(result->lhs));
    am_memcpy(&result->rhs, instruction_data + sizeof(result->result) + sizeof(result->lhs), sizeof(result->rhs));
    return sizeof(result->result) + sizeof(result->lhs) + sizeof(result->rhs);
}

static CHECK_RESULT usize ssa_extract_func_start(u8 *instruction_data, SsaInsFuncStart *result) {
    am_memcpy(&result->func_index, instruction_data, sizeof(result->func_index));
    am_memcpy(&result->num_args, instruction_data + sizeof(result->func_index), sizeof(result->num_args));
    return sizeof(result->func_index) + sizeof(result->num_args);
}

static CHECK_RESULT usize ssa_extract_func_call(u8 *instruction_data, SsaInsFuncCall *result) {
    am_memcpy(&result->result, instruction_data, sizeof(result->result));
    am_memcpy(&result->func_decl, instruction_data + sizeof(result->result), sizeof(result->func_decl));
    return sizeof(result->result) + sizeof(result->func_decl);
}

static void add_intermediates(BytecodeCompilerContext *self) {
    Ssa *ssa;
    Buffer *instructions;
    SsaNumber *intermediate;
    SsaNumber *intermediates_end;

    ssa = self->parser->ssa;
    instructions = &self->bytecode->instructions;
    intermediate = buffer_begin(&ssa->intermediates);
    intermediates_end = buffer_end(&ssa->intermediates);

    throw_if_error(buffer_expand(instructions,
                                 sizeof(u16) + (sizeof(u8) + sizeof(u64)) * ssa->intermediates.size));
    throw_if_error(buffer_append(instructions, &ssa->intermediates.size, sizeof(u16)));
    for(; intermediate != intermediates_end; ++intermediate) {
        throw_if_error(buffer_append(instructions, &intermediate->type, sizeof(u8)));
        throw_if_error(buffer_append(instructions, &intermediate->value.integer, sizeof(u64)));
    }
}

#define NUM_MAX_REGS 256

static void add_instructions(BytecodeCompilerContext *self) {
    Ssa *ssa;
    u8 *instruction;
    u8 *instructions_end;
    SsaInsForm1 ssa_ins_form1;
    SsaInsForm2 ssa_ins_form2;
    SsaInsFuncStart ssa_ins_func_start;
    SsaInsFuncCall ssa_ins_func_call;
    FILE *file;
    char *filename;
    LhsExpr *reg_types[NUM_MAX_REGS]; /* TODO: Remove this. Encode this data in the register itself */

    ssa = self->parser->ssa;
    instruction = buffer_begin(&ssa->instructions);
    instructions_end = buffer_end(&ssa->instructions);
    /*#warning "dont forget to remove this" */
    filename = malloc(self->parser->tokenizer.code_name.size + 3);
    filename[0] = '\0';
    strcat(filename, self->parser->tokenizer.code_name.data);
    strcat(filename, ".z");
    file = fopen(filename, "wb");
    free(filename);
    #ifdef DEBUG
    am_memset(reg_types, 0, sizeof(reg_types));
    #endif

    fputs("typedef i64 signed long long;\n", file);
    fputs("typedef f64 double;\n", file);

    while(instruction != instructions_end) {
        switch(*instruction++) {
            case SSA_ASSIGN_INTER: {
                SsaNumber number;
                instruction += ssa_extract_form1(instruction, &ssa_ins_form1);
                number = ssa_get_intermediate(ssa, ssa_ins_form1.rhs);
                assert(ssa_ins_form1.lhs < NUM_MAX_REGS);
                if(number.type == SSA_NUMBER_TYPE_INTEGER) {
                    fprintf(file, "i64 r%d = %zu;\n", ssa_ins_form1.lhs, number.value.integer);
                    reg_types[ssa_ins_form1.lhs] = self->parser->compiler->default_types.i64;
                } else if(number.type == SSA_NUMBER_TYPE_FLOAT) {
                    fprintf(file, "f64 r%d = %f;\n", ssa_ins_form1.lhs, number.value.floating);
                    reg_types[ssa_ins_form1.lhs] = self->parser->compiler->default_types.f64;
                } else {
                    assert(bool_false && "TODO: Implement");
                }
                break;
            }
            case SSA_ASSIGN_STRING: {
                BufferView str;
                instruction += ssa_extract_form1(instruction, &ssa_ins_form1);
                str = ssa_get_string(ssa, ssa_ins_form1.rhs);
                fprintf(file, "const char* r%d = \"%.*s\";\n", ssa_ins_form1.lhs, (int)str.size, str.data);
                assert(ssa_ins_form1.lhs < NUM_MAX_REGS);
                reg_types[ssa_ins_form1.lhs] = self->parser->compiler->default_types.str;
                break;
            }
            case SSA_ASSIGN_REG: {
                LhsExpr *rhs_type;
                const char *rhs_type_name;
                instruction += ssa_extract_form1(instruction, &ssa_ins_form1);
                assert(ssa_ins_form1.rhs < NUM_MAX_REGS);
                rhs_type = reg_types[ssa_ins_form1.rhs];
                if(rhs_type == self->parser->compiler->default_types.i64) {
                    rhs_type_name = "i64";
                } else if(rhs_type == self->parser->compiler->default_types.f64) {
                    rhs_type_name = "f64";
                } else if(rhs_type == self->parser->compiler->default_types.str) {
                    rhs_type_name = "const char*";
                } else {
                    amal_log_error("Invalid rhs type %p for reg %d", rhs_type, ssa_ins_form1.rhs);
                    assert(bool_false && "TODO: Implement");
                }
                fprintf(file, "%s r%d = r%d;\n", rhs_type_name, ssa_ins_form1.lhs, ssa_ins_form1.rhs);
                break;
            }
            case SSA_ADD: {
                instruction += ssa_extract_form2(instruction, &ssa_ins_form2);
                fprintf(file, "r%d = r%d + r%d;\n", ssa_ins_form2.result, ssa_ins_form2.lhs, ssa_ins_form2.rhs);
                assert(ssa_ins_form2.result < NUM_MAX_REGS);
                assert(ssa_ins_form2.rhs < NUM_MAX_REGS);
                reg_types[ssa_ins_form2.result] = reg_types[ssa_ins_form2.rhs];
                break;
            }
            case SSA_SUB: {
                instruction += ssa_extract_form2(instruction, &ssa_ins_form2);
                fprintf(file, "r%d = r%d - r%d;\n", ssa_ins_form2.result, ssa_ins_form2.lhs, ssa_ins_form2.rhs);
                assert(ssa_ins_form2.result < NUM_MAX_REGS);
                assert(ssa_ins_form2.rhs < NUM_MAX_REGS);
                reg_types[ssa_ins_form2.result] = reg_types[ssa_ins_form2.rhs];
                break;
            }
            case SSA_MUL: {
                instruction += ssa_extract_form2(instruction, &ssa_ins_form2);
                fprintf(file, "r%d = r%d * r%d;\n", ssa_ins_form2.result, ssa_ins_form2.lhs, ssa_ins_form2.rhs);
                assert(ssa_ins_form2.result < NUM_MAX_REGS);
                assert(ssa_ins_form2.rhs < NUM_MAX_REGS);
                reg_types[ssa_ins_form2.result] = reg_types[ssa_ins_form2.rhs];
                break;
            }
            case SSA_DIV: {
                instruction += ssa_extract_form2(instruction, &ssa_ins_form2);
                fprintf(file, "r%d = r%d / r%d;\n", ssa_ins_form2.result, ssa_ins_form2.lhs, ssa_ins_form2.rhs);
                assert(ssa_ins_form2.result < NUM_MAX_REGS);
                assert(ssa_ins_form2.rhs < NUM_MAX_REGS);
                reg_types[ssa_ins_form2.result] = reg_types[ssa_ins_form2.rhs];
                break;
            }
            case SSA_FUNC_START: {
                int i;
                instruction += ssa_extract_func_start(instruction, &ssa_ins_func_start);
                fprintf(file, "void f%zu(", ssa_ins_func_start.func_index);
                fputs(") {\n", file);
                for(i = 0; i < ssa_ins_func_start.num_args; ++i) {
                    if(i > 0)
                        fputs(", ", file);
                    fprintf(file, "p%d", i);
                }
                break;
            }
            case SSA_FUNC_END:
                fputs("}\n", file);
                break;
            case SSA_PUSH: {
                SsaRegister reg;
                am_memcpy(&reg, instruction, sizeof(SsaRegister));
                instruction += sizeof(SsaRegister);
                fprintf(file, "PUSH r%d ***\n", reg);
                break;
            }
            case SSA_CALL:
                instruction += ssa_extract_func_call(instruction, &ssa_ins_func_call);
                fprintf(file, "r%d = CALL %p ***\n", ssa_ins_func_call.result, ssa_ins_func_call.func_decl);
                break;
        }
    }

    fclose(file);
}

void generate_bytecode_from_ssa(BytecodeCompilerContext *self) {
    add_intermediates(self);
    add_instructions(self);
}