#include <ir_builder.h>
#include <ir_builtin.h>

void scc_ir_builder_init(scc_ir_builder_t *builder) {
    builder->current_bblock = SCC_IR_REF_NULL;
    builder->current_func = SCC_IR_REF_NULL;

    scc_ir_cprog_init(&builder->cprog);
    scc_ir_ctx_init(&builder->ctx);
}

void scc_ir_builder_drop(scc_ir_builder_t *builder) {
    scc_ir_cprog_drop(&builder->cprog);
    scc_ir_ctx_drop(&builder->ctx);
}

scc_ir_func_ref_t scc_ir_builder_begin_func(scc_ir_builder_t *builder,
                                            const char *name,
                                            scc_ir_type_ref_t type,
                                            const char **param_names) {
    scc_ir_func_t func = {0};

    // 初始化参数和基本块向量
    func.name = name;
    func.type = type;
    scc_vec_init(func.params);
    scc_vec_init(func.bblocks);

    // 创建函数并设置为当前函数
    scc_ir_func_ref_t func_ref = scc_ir_ctx_new_func(&builder->ctx, &func);
    builder->current_func = func_ref;

    // 如果提供了参数名称，则创建参数节点
    if (param_names == null) {
        return func_ref;
    }

    scc_ir_func_t *func_ptr = scc_ir_ctx_get_func(&builder->ctx, func_ref);
    scc_ir_type_t *func_type = scc_ir_ctx_get_type(&builder->ctx, type);

    if (func_type == null || func_type->tag != SCC_IR_TYPE_FUNC) {
        LOG_ERROR("Invalid function type");
        return func_ref;
    }

    scc_vec_foreach(func_type->data.function.params, i) {
        scc_ir_type_ref_t param_type =
            scc_vec_at(func_type->data.function.params, i);

        scc_ir_node_t param_node = {0};
        param_node.tag = SCC_IR_NODE_NULL; // 参数节点标记
        param_node.type = param_type;
        param_node.name = param_names[i];
        scc_vec_init(param_node.used_by);

        scc_ir_node_ref_t param_ref =
            scc_ir_ctx_new_node(&builder->ctx, &param_node);
        scc_vec_push(func_ptr->params, param_ref);
    }

    return func_ref;
}

void scc_ir_builder_end_func(scc_ir_builder_t *builder) {
    scc_ir_func_t *func_ptr =
        scc_ir_ctx_get_func(&builder->ctx, builder->current_func);
    if (func_ptr == null) {
        LOG_FATAL("Invalid function reference");
        return;
    }
    if (scc_vec_size(func_ptr->bblocks) == 0) {
        scc_vec_push(builder->cprog.func_decls, builder->current_func);
    } else {
        scc_vec_push(builder->cprog.func_defs, builder->current_func);
    }
    builder->current_func = 0;
}

scc_ir_func_ref_t scc_ir_builder_current_func(scc_ir_builder_t *builder) {
    return builder->current_func;
}

scc_ir_bblock_ref_t scc_ir_builder_begin_bblock(scc_ir_builder_t *builder,
                                                const char *label) {
    scc_ir_bblock_t bblock = {0};

    if (label) {
        bblock.label = label;
    } else {
        // TODO 自动生成标签
        char *auto_label = scc_malloc(sizeof(char) * 64);
        if (auto_label)
            LOG_FATAL("allocate memory failed");
        static int auto_counter = 0;
        scc_snprintf(auto_label, sizeof(auto_label), ".BB%d", auto_counter++);
        bblock.label = auto_label;
    }

    scc_vec_init(bblock.instrs);

    scc_ir_bblock_ref_t bblock_ref =
        scc_ir_ctx_new_bblock(&builder->ctx, &bblock);
    builder->current_bblock = bblock_ref;

    // 将基本块添加到当前函数
    scc_ir_func_t *current_func =
        scc_ir_ctx_get_func(&builder->ctx, builder->current_func);
    if (current_func) {
        scc_vec_push(current_func->bblocks, bblock_ref);
    }

    return bblock_ref;
}

void scc_ir_builder_end_bblock(scc_ir_builder_t *builder) {
    builder->current_bblock = 0;
}

void scc_ir_builder_set_current_bblock(scc_ir_builder_t *builder,
                                       scc_ir_bblock_ref_t bblock) {
    builder->current_bblock = bblock;
}

static void scc_ir_builder_add_instr(scc_ir_builder_t *builder,
                                     scc_ir_node_ref_t node) {
    scc_ir_bblock_t *current_bblock =
        scc_ir_ctx_get_bblock(&builder->ctx, builder->current_bblock);
    if (current_bblock) {
        scc_vec_push(current_bblock->instrs, node);
    } else {
        LOG_ERROR("Current basic block is not set");
    }
}

scc_ir_node_ref_t scc_ir_builder_alloca(scc_ir_builder_t *builder,
                                        scc_ir_type_ref_t type,
                                        const char *name) {
    scc_ir_node_t alloc_node = {0};
    alloc_node.tag = SCC_IR_NODE_ALLOC;
    alloc_node.type = scc_ir_ctx_new_type(
        &builder->ctx,
        &(scc_ir_type_t){.tag = SCC_IR_TYPE_PTR, .data.pointer.base = type});
    alloc_node.name = name;

    scc_ir_node_ref_t node_ref =
        scc_ir_ctx_new_node(&builder->ctx, &alloc_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_load(scc_ir_builder_t *builder,
                                      scc_ir_node_ref_t ptr) {
    scc_ir_node_t load_node = {0};
    load_node.tag = SCC_IR_NODE_LOAD;
    load_node.data.load.target = ptr;

    // 设置类型为指针指向的类型
    scc_ir_node_t *ptr_node = scc_ir_ctx_get_node(&builder->ctx, ptr);
    if (ptr_node) {
        scc_ir_type_t *ptr_type =
            scc_ir_ctx_get_type(&builder->ctx, ptr_node->type);
        if (ptr_type && ptr_type->tag == SCC_IR_TYPE_PTR) {
            load_node.type = ptr_type->data.pointer.base;
        }
    }

    scc_ir_node_ref_t node_ref = scc_ir_ctx_new_node(&builder->ctx, &load_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_store(scc_ir_builder_t *builder,
                                       scc_ir_node_ref_t ptr,
                                       scc_ir_node_ref_t value) {
    scc_ir_node_t store_node = {0};
    store_node.tag = SCC_IR_NODE_STORE;
    store_node.data.store.target = ptr;
    store_node.data.store.value = value;

    scc_ir_node_ref_t node_ref =
        scc_ir_ctx_new_node(&builder->ctx, &store_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_get_ptr(scc_ir_builder_t *builder,
                                         scc_ir_node_ref_t ptr,
                                         scc_ir_node_ref_t index) {
    scc_ir_node_t get_ptr_node = {0};
    get_ptr_node.tag = SCC_IR_NODE_GET_PTR;
    get_ptr_node.data.get_ptr.src_addr = ptr;
    get_ptr_node.data.get_ptr.index = index;

    // 类型应与源地址相同（都是指针）
    scc_ir_node_t *src_node = scc_ir_ctx_get_node(&builder->ctx, ptr);
    if (src_node) {
        get_ptr_node.type = src_node->type;
    }

    scc_ir_node_ref_t node_ref =
        scc_ir_ctx_new_node(&builder->ctx, &get_ptr_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_binop(scc_ir_builder_t *builder,
                                       scc_ir_op_type_t op,
                                       scc_ir_node_ref_t lhs,
                                       scc_ir_node_ref_t rhs) {
    scc_ir_node_t binop_node = {0};
    binop_node.tag = SCC_IR_NODE_OP;
    binop_node.data.op.op = op;
    binop_node.data.op.lhs = lhs;
    binop_node.data.op.rhs = rhs;

    // 类型通常与操作数相同（对于算术运算）
    scc_ir_node_t *lhs_node = scc_ir_ctx_get_node(&builder->ctx, lhs);
    if (lhs_node) {
        binop_node.type = lhs_node->type;
    }

    scc_ir_node_ref_t node_ref =
        scc_ir_ctx_new_node(&builder->ctx, &binop_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_cmp(scc_ir_builder_t *builder,
                                     scc_ir_op_type_t op, scc_ir_node_ref_t lhs,
                                     scc_ir_node_ref_t rhs) {
    scc_ir_node_t cmp_node = {0};
    cmp_node.tag = SCC_IR_NODE_OP;
    cmp_node.data.op.op = op;
    cmp_node.data.op.lhs = lhs;
    cmp_node.data.op.rhs = rhs;

    // 比较操作的结果通常是布尔值
    cmp_node.type = scc_ir_ctx_get_builtin_i32(&builder->ctx);

    scc_ir_node_ref_t node_ref = scc_ir_ctx_new_node(&builder->ctx, &cmp_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_jump(scc_ir_builder_t *builder,
                                      scc_ir_bblock_ref_t target) {
    scc_ir_node_t jump_node = {0};
    jump_node.tag = SCC_IR_NODE_JUMP;
    jump_node.data.jump.target_bblock = target;

    scc_ir_node_ref_t node_ref = scc_ir_ctx_new_node(&builder->ctx, &jump_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_branch(scc_ir_builder_t *builder,
                                        scc_ir_node_ref_t cond,
                                        scc_ir_bblock_ref_t true_target,
                                        scc_ir_bblock_ref_t false_target) {
    scc_ir_node_t branch_node = {0};
    branch_node.tag = SCC_IR_NODE_BRANCH;
    branch_node.data.branch.cond = cond;
    branch_node.data.branch.true_bblock = true_target;
    branch_node.data.branch.false_bblock = false_target;

    scc_ir_node_ref_t node_ref =
        scc_ir_ctx_new_node(&builder->ctx, &branch_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_call(scc_ir_builder_t *builder,
                                      scc_ir_func_ref_t callee,
                                      const scc_ir_node_ref_t *args,
                                      usize arg_count) {
    scc_ir_node_t call_node = {0};
    call_node.tag = SCC_IR_NODE_CALL;
    call_node.data.call.callee = callee;

    scc_vec_init(call_node.data.call.args);
    for (usize i = 0; i < arg_count; i++) {
        scc_vec_push(call_node.data.call.args, args[i]);
    }

    // 设置返回类型为被调用函数的返回类型
    scc_ir_func_t *callee_func = scc_ir_ctx_get_func(&builder->ctx, callee);
    if (callee_func) {
        scc_ir_type_t *func_type =
            scc_ir_ctx_get_type(&builder->ctx, callee_func->type);
        if (func_type && func_type->tag == SCC_IR_TYPE_FUNC) {
            call_node.type = func_type->data.function.ret_type;
        }
    }

    scc_ir_node_ref_t node_ref = scc_ir_ctx_new_node(&builder->ctx, &call_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_ret(scc_ir_builder_t *builder,
                                     scc_ir_node_ref_t value) {
    scc_ir_node_t ret_node = {0};
    ret_node.tag = SCC_IR_NODE_RET;
    ret_node.data.ret.ret_val = value;

    scc_ir_node_ref_t node_ref = scc_ir_ctx_new_node(&builder->ctx, &ret_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);

    return node_ref;
}

scc_ir_node_ref_t scc_ir_builder_ret_void(scc_ir_builder_t *builder) {
    scc_ir_node_t ret_node = {0};
    ret_node.tag = SCC_IR_NODE_RET;
    ret_node.data.ret.ret_val = 0; // 无返回值

    scc_ir_node_ref_t node_ref = scc_ir_ctx_new_node(&builder->ctx, &ret_node);

    // 添加到当前基本块
    scc_ir_builder_add_instr(builder, node_ref);
    return node_ref;
}