#include "ir.h"

static size_t reg_gen_id = 0;
static size_t lab_gen_id = 0;

Operand
emit_arith(ProgramBASM *program, Node *node,  Operator op) {
    LineInfo line = (LineInfo){
        .line = node->line,
        .col = node->col,
    };
    Operand reg_a = emit_basm(program, node->builtin.args[0]);
    Operand reg_b;
    for (size_t i = 1; i < array_size(node->builtin.args); ++i) {
        Node *arg = node->builtin.args[i];
        reg_b = emit_basm(program, arg);
        Operand reg_dst = NEW_REG();
        EMIT_2(program, line, op, reg_dst, reg_a, reg_b);
        reg_a = reg_dst;
    }
    return reg_a;
}

Operand
emit_numcomp(ProgramBASM *program, Node *node, Operator op) {
    LineInfo line = (LineInfo){
        .line = node->line,
        .col = node->col,
    };
    Operand label_false = NEW_LAB();
    Operand label_end = NEW_LAB();
    Operand reg_a = emit_basm(program, node->builtin.args[0]);
    Operand reg_b;
    for (size_t i = 1; i < array_size(node->builtin.args); ++i) {
        Node *arg = node->builtin.args[i];
        reg_b = emit_basm(program, arg);
        EMIT_2(program, line, op, label_false, reg_a, reg_b);
        reg_a = reg_b;
    }
    Operand reg_out = NEW_REG();
    EMIT_1(program, line, OP_LD8, reg_out, NEW_S64(1));
    EMIT_0(program, line, OP_JMP, label_end);
    EMIT_0(program, line, OP_LABEL, label_false);
    EMIT_1(program, line, OP_LD8, reg_out, NEW_S64(0));
    EMIT_0(program, line, OP_LABEL, label_end);
    return reg_out;
}

Operand
emit_builtin(ProgramBASM *program, Node *node) {
    switch (node->builtin.type) {
        case TOKEN_ADD: { return emit_arith(program, node, OP_ADD); } break;
        case TOKEN_SUB: { return emit_arith(program, node, OP_SUB); } break;
        case TOKEN_MUL: { return emit_arith(program, node, OP_MUL); } break;
        case TOKEN_DIV: { return emit_arith(program, node, OP_DIV); } break;
        case TOKEN_MOD: { return emit_arith(program, node, OP_MOD); } break;
        case TOKEN_EQ: { return emit_numcomp(program, node, OP_JMP_NEQ); } break;
        case TOKEN_LT: { return emit_numcomp(program, node, OP_JMP_GE); } break;
        case TOKEN_GT: { return emit_numcomp(program, node, OP_JMP_LE); } break;
        case TOKEN_LE: { return emit_numcomp(program, node, OP_JMP_GT); } break;
        case TOKEN_GE: { return emit_numcomp(program, node, OP_JMP_LT); } break;
        default: {
            push_error(ERR_TYPE_BASM, ERR_UNIMPLEMENTED, node->line, node->col);
            return (Operand){0};
        } break;
    }
}

Operand
emit_number(ProgramBASM *program, Node *node) {
    // TODO: ldX depending on type of number.
    LineInfo line = (LineInfo){.line = node->line, .col = node->col};
    Operand reg_dst = NEW_REG();
    Operand num = NEW_S64(node->number.integral);
    EMIT_1(program, line, OP_LD64, reg_dst, num);
    return reg_dst;
}

Operand
emit_bool(ProgramBASM *program, Node *node) {
    LineInfo line = (LineInfo){.line = node->line, .col = node->col};
    Operand reg_dst = NEW_REG();
    Operand val;
    if (node->boolean) {
        val = NEW_S64(1);
    } else {
        val = NEW_S64(0);
    }
    EMIT_1(program, line, OP_LD8, reg_dst, val);
    return reg_dst;
}

Operand
emit_if(ProgramBASM *program, Node *node) {
    LineInfo line = (LineInfo){.line = node->line, .col = node->col};
    Operand label_false = NEW_LAB();
    Operand label_end = NEW_LAB();

    // If type of expression is `void`, there is nothing to return. Don't
    // generate a new register or copy output.
    Operand reg_out = (Operand){0};
    if (node->expr_type != &default_types[TYPE_VOID]) {
        reg_out = NEW_REG();
    }

    Operand cond = emit_basm(program, node->ifexpr.cond);
    EMIT_1(program, line, OP_JMP_FALSE, label_false, cond);

    Operand ret_true = emit_basm(program, node->ifexpr.expr_true);
    if (node->ifexpr.expr_false != NULL) {
        if (node->expr_type != &default_types[TYPE_VOID]) {
            EMIT_1(program, line, OP_CP64, reg_out, ret_true); // TODO: depends on the type of output
        }
        EMIT_0(program, line, OP_JMP, label_end);
        EMIT_0(program, line, OP_LABEL, label_false);
        Operand ret_false = emit_basm(program, node->ifexpr.expr_false);
        if (node->expr_type != &default_types[TYPE_VOID]) {
            EMIT_1(program, line, OP_CP64, reg_out, ret_false); // TODO: depends on the type of output
        }
        EMIT_0(program, line, OP_LABEL, label_end);
    } else {
        EMIT_0(program, line, OP_LABEL, label_false);
    }

    return reg_out;
}

// TODO: emit_and
// TODO: emit_or
// TODO: emit_not
// TODO: emit_block
// TODO: emit_procedure
// TODO: emit_local_vars
// TODO: emit_param_vars
// TODO: emit_proc_call
// TODO: emit_global_vars

Operand
emit_basm(ProgramBASM *program, Node *node) {
    switch (node->type) {
        case NODE_BOOL: { return emit_bool(program, node); } break;
        case NODE_NUMBER: { return emit_number(program, node); } break;
        case NODE_BUILTIN: { return emit_builtin(program, node); } break;
        case NODE_IF: { return emit_if(program, node); } break;
        default: {
            push_error(ERR_TYPE_BASM, ERR_UNIMPLEMENTED, node->line, node->col);
            return (Operand){0};
        } break;
    }
}

ProgramBASM *
generate_basm(ParseTree *parse_tree) {
    ProgramBASM *program = malloc(sizeof(ProgramBASM));
    array_init(program->inst, 0);
    array_init(program->lines, 0);
    for (size_t i = 0; i < array_size(parse_tree->roots); ++i) {
        Node *root = parse_tree->roots[i];
        emit_basm(program, root);
    }
    return program;
}