#include "compiler.h"

void compiler_init(compiler_t* self, module_t* module)
{
    assert(self);
    err_init(&self->err);
    self->module = module;
    self->loc_counter = 0;
}

void compiler_free(compiler_t* self)
{
    assert(self);
    err_free(&self->err);
}

void compiler_compile(compiler_t* self,
                      node_t* node,
                      prog_t* prog)
{
    assert(self);
    assert(node);
    assert(prog);

    if (!err_is_ok(&self->err))
    {
        return;
    }

    ccm_t* ccm = &self->module->ccm;
    sym_t* sym = &self->module->sym;

    switch (node->kind)
    {
        case NODE_ASSIGN: {
            node_t* target = node->children.data[0];
            node_t* expr = node->children.data[1];

            if (target->kind == NODE_INDEX) {
                node_t* ident = target->children.data[0];

                sym_entry_t const* entry = sym_try_get_value(
                    sym,
                    ident->value
                );


                if (!entry) {
                    err_push(&self->err, ident->line,
                             "undefined array '%s'", ident->value);
                    return;
                }

                if (entry->is_const) {
                    err_push(&self->err, ident->line,
                             "cannot assign value to constant '%s'",
                             ident->value);
                    return;
                }
                compiler_compile(self, expr, prog);
                size_t indexes_count = target->children.size - 1;

                for (size_t i=1; i<target->children.size; i++) {
                    compiler_compile(
                        self,
                        target->children.data[i],
                        prog
                    );
                }

                prog_add_instr(prog, OP_PUSH, entry->local_addr);
                prog_add_instr(prog, OP_ASTORE, indexes_count);
            } else {
                compiler_compile(self, expr, prog);
                int status =
                    sym_try_assign(
                        sym, 
                        target->value, self->loc_counter
                    );

                sym_entry_t* entry = sym_try_get_value(
                    sym, 
                    target->value
                );

                assert(entry);

                if (entry->is_const) {
                    err_push(&self->err, target->line,
                             "cannot assign value to constant '%s'",
                             target->value);
                    return;
                }
                if (!status)
                {
                    err_push(&self->err, target->line,
                             "cannot assign value to '%s'",
                             target->value);
                    return;
                }

                prog_add_instr(prog, OP_LOCAL_STORE, self->loc_counter);

                self->loc_counter++;
            }

        } break;

        case NODE_CONSTDECL:
        case NODE_VARDECL: {
            node_t const* ident = node->children.data[0];
            node_t* expr = node->children.data[1];

            compiler_compile(self, expr, prog);
            int status =
                sym_declare(
                    sym, ident->value,
                    self->loc_counter,
                    node->kind == NODE_CONSTDECL
                );
            assert(status);
            prog_add_instr(prog, OP_LOCAL_STORE, self->loc_counter);
            self->loc_counter++;
        } break;

        case NODE_IDENT: {
            char const* name = node->value;
            sym_entry_t* entry = sym_try_get_value(sym, name);
            assert(entry);
            prog_add_instr(prog, OP_LOCAL_LOAD, entry->local_addr);
        } break;

        case NODE_ARRAY: {
            for (size_t i=0; i<node->children.size; i++)
            {
                size_t k = node->children.size - 1 - i;
                compiler_compile(self, node->children.data[k], prog);
            }
            prog_add_instr(prog, OP_MK_ARRAY, node->children.size);
        } break;

        case NODE_LT: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);
            prog_add_instr(prog, OP_LT, CCM_NO_PARAM);
        } break;

        case NODE_LE: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);
            prog_add_instr(prog, OP_GT, CCM_NO_PARAM);
            prog_add_instr(prog, OP_NOT, CCM_NO_PARAM);
        } break;

        case NODE_GT: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);
            prog_add_instr(prog, OP_GT, CCM_NO_PARAM);
        } break;

        case NODE_GE: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);
            prog_add_instr(prog, OP_LT, CCM_NO_PARAM);
            prog_add_instr(prog, OP_NOT, CCM_NO_PARAM);
        } break;

        case NODE_EQ: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);
            prog_add_instr(prog, OP_EQ, CCM_NO_PARAM);
        } break;

        case NODE_NE: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);
            prog_add_instr(prog, OP_EQ, CCM_NO_PARAM);
            prog_add_instr(prog, OP_NOT, CCM_NO_PARAM);
        } break;

        case NODE_INDEX: {
            for (size_t i=0; i<node->children.size; i++)
            {
                size_t k = node->children.size - 1 - i;
                compiler_compile(self, node->children.data[k], prog);
            }

            prog_add_instr(prog, OP_INDEX, node->children.size - 1);
        } break;

        case NODE_IN: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);

            prog_add_instr(prog, OP_IN, CCM_NO_PARAM);
        } break;

        case NODE_AND: {
            compiler_compile_and(self, node, prog);
        } break;
        case NODE_OR: {
            compiler_compile_or(self, node, prog);
        } break;
        case NODE_NOT: {
            compiler_compile(self, node->children.data[0], prog);
            prog_add_instr(prog, OP_NOT, CCM_NO_PARAM);
        } break;

        case NODE_ASSERT_EQ: {
            compiler_compile(self, node->children.data[0], prog);
            prog_add_instr(prog, OP_ASSERT_EQ, CCM_NO_PARAM);
        } break;

        case NODE_ASSERT_NE: {
            compiler_compile(self, node->children.data[0], prog);
            prog_add_instr(prog, OP_ASSERT_NE, CCM_NO_PARAM);
        } break;

        case NODE_MODULE: {
            for (size_t i=0; i<node->children.size; i++)
            {
                compiler_compile(self,
                                 node->children.data[i],
                                 prog);
            }
        } break;

        case NODE_STR: {
            size_t id = prog_add_new_constant(
                prog,
                ccm_to_str(&self->module->ccm,
                           node->value,
                           node->line)
            );

            prog_add_instr(prog, OP_PUSH, id);
        } break;

        case NODE_NUM: {
            size_t id = prog_add_new_constant(
                prog,
                ccm_to_num(&self->module->ccm,
                           atof(node->value),
                           node->line)
            );

            prog_add_instr(prog, OP_PUSH, id);
        } break;

        case NODE_TUPLE: {
            for (size_t i=0; i<node->children.size; i++)
            {
                size_t k = node->children.size - 1 - i;

                compiler_compile(self,
                                 node->children.data[k],
                                 prog);
            }

            prog_add_instr(prog, OP_MK_TUPLE, node->children.size);
        } break;

        case NODE_BOOL: {
            int val = strcmp(node->value, "true") == 0;
            CCM value = ccm_to_boolean(ccm, val, node->line);
            size_t id = prog_add_new_constant(prog, value);
            prog_add_instr(prog, OP_PUSH, id);
        } break;

        case NODE_SUB: {
            compiler_compile(self, node->children.data[0], prog);

            if (node->children.size == 2)
            {
                compiler_compile(self, node->children.data[1], prog);
                prog_add_instr(prog, OP_SUB, CCM_NO_PARAM);
            }
            else
            {
                prog_add_instr(prog, OP_USUB, CCM_NO_PARAM);
            }
        } break;

        case NODE_ADD:
        case NODE_MUL:
        case NODE_DIV:
        case NODE_MOD:
        case NODE_POW: {
            compiler_compile(self, node->children.data[0], prog);
            compiler_compile(self, node->children.data[1], prog);
            Opcode op;

            switch (node->kind)
            {
                case NODE_ADD: op = OP_ADD; break;
                case NODE_MUL: op = OP_MUL; break;
                case NODE_DIV: op = OP_DIV; break;
                case NODE_MOD: op = OP_MOD; break;
                case NODE_POW: op = OP_POW; break;
                default: abort();
            }

            prog_add_instr(prog, op, CCM_NO_PARAM);
        } break;

        default: {
            fprintf(stderr, "cannot compile node %s\n",
                    NodeKindStr[node->kind]);
            abort();
        }
    }
}

void compiler_compile_or(compiler_t* self,
                          node_t* node,
                          prog_t* prog)
{
    ccm_t* ccm = &self->module->ccm;
    vec_t to_true;
    vec_init(&to_true);

    for (size_t i=0; i<node->children.size; i++)
    {
        compiler_compile(self, node->children.data[i], prog);
        prog_add_instr(prog, OP_NOT, 0);
        size_t addr = prog_add_instr(prog, OP_BRF, 0);
        vec_push(&to_true, (void*) addr);
    }

    // FALSE
    prog_add_instr(prog, OP_PUSH, prog_add_new_constant(
        prog,
        ccm_to_boolean(ccm, 0, node->line)
    ));

    size_t to_end = prog_add_instr(prog, OP_BR, 0);

    // TRUE
    size_t true_point = prog->instrs.size;
    prog_add_instr(prog, OP_PUSH, prog_add_new_constant(
        prog,
        ccm_to_boolean(ccm, 1, node->line)
    ));

    for (size_t i=0; i<to_true.size; i++)
    {
        size_t k = ((size_t) to_true.data[i]);
        ((instr_t*) prog->instrs.data[k])->param = true_point;
    }

    size_t end_point = prog->instrs.size;
    ((instr_t*) prog->instrs.data[to_end])->param = end_point;

    vec_free(&to_true);
}

void compiler_compile_and(compiler_t* self,
                          node_t* node,
                          prog_t* prog)
{
    ccm_t* ccm = &self->module->ccm;
    vec_t to_false;
    vec_init(&to_false);

    for (size_t i=0; i<node->children.size; i++)
    {
        compiler_compile(self, node->children.data[i], prog);
        size_t addr = prog_add_instr(prog, OP_BRF, 0);
        vec_push(&to_false, (void*) addr);
    }

    // TRUE
    prog_add_instr(prog, OP_PUSH, prog_add_new_constant(
        prog,
        ccm_to_boolean(ccm, 1, node->line)
    ));

    size_t to_end = prog_add_instr(prog, OP_BR, 0);

    // FALSE
    size_t false_point = prog_add_instr(
        prog,
        OP_PUSH,
        prog_add_new_constant(
            prog,
            ccm_to_boolean(ccm, 0, node->line)
        )
    );

    for (size_t i=0; i<to_false.size; i++)
    {
        size_t k = ((size_t) to_false.data[i]);
        ((instr_t*) prog->instrs.data[k])->param = false_point;
    }

    size_t end_point = prog->instrs.size;
    ((instr_t*) prog->instrs.data[to_end])->param = end_point;

    vec_free(&to_false);
}