#include "ir_def.h"
#include <ir_builder.h>
#include <ir_prog.h>

#define GET_MODULE(builder) (&(builder->cprog->module))

void scc_ir_builder_init(scc_ir_builder_t *builder, scc_ir_cprog_t *cprog) {
    builder->current_bblock = SCC_IR_REF_nullptr;
    builder->current_func = SCC_IR_REF_nullptr;
    builder->cprog = cprog;

    scc_ir_ctx_init(&builder->ctx, GET_MODULE(builder));
#ifndef SCC_NO_DEBUG
    builder->borrow_depth = 0;
    builder->dbg_file = nullptr;
    builder->dbg_line = 0;
#endif
}

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

scc_ir_func_ref_t scc_ir_builder_func(scc_ir_builder_t *builder,
                                      scc_ir_type_ref_t type_ref,
                                      const char *name) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_func_ref_t func_ref =
        scc_ir_ctx_declare_func(&builder->ctx, type_ref, name);
    scc_vec_push(builder->cprog->func_decls, func_ref);
    return func_ref;
}

scc_ir_type_ref_t scc_ir_builder_type(scc_ir_builder_t *builder,
                                      const scc_ir_type_t *type_desc) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    return scc_ir_ctx_get_type(&builder->ctx, type_desc);
}

scc_ir_value_ref_t scc_ir_builder_const_string(scc_ir_builder_t *builder,
                                               const char *str, usize len) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_type_ref_t u8_type = scc_ir_builder_type_u8(builder);
    scc_ir_type_t array_type = {
        .tag = SCC_IR_TYPE_ARRAY,
        .data.array.base = u8_type,
        .data.array.len = len - 1, // 包含 nullptr 结尾
    };
    scc_ir_type_ref_t array_type_ref =
        scc_ir_ctx_get_type(&builder->ctx, &array_type);

    // 5. 创建聚合节点
    scc_ir_value_t const_array_value = {
        .tag = SCC_IR_VALUE_TAG_CONST_ARRAY,
        .type = array_type_ref,
        .data.const_array.base_type = u8_type,
    };
    char *buff = scc_malloc(len - 1);
    Assert(buff);
    // FIXME content to real string
    for (usize i = 1; i < len - 1; i++) {
        buff[i - 1] = str[i];
    }
    buff[len - 2] = '\0';
    scc_vec_unsafe_from_buffer(const_array_value.data.const_array.fields,
                               (u8 *)buff, len - 1);
    scc_ir_value_ref_t const_array_ref =
        scc_ir_module_add_value(builder->ctx.module, &const_array_value);
    Assert(const_array_ref != SCC_IR_REF_nullptr);

    // 3. 创建全局变量节点，类型为指针，初始值指向常量数组
    scc_ir_value_ref_t global_value_ref =
        scc_ir_builder_global_alloca(builder, array_type_ref, const_array_ref);
    // scc_hashtable_insert(builder);

    scc_ir_value_ref_t pointer_to_global_value = scc_ir_module_add_value(
        builder->ctx.module,
        &(scc_ir_value_t){
            .tag = SCC_IR_VALUE_TAG_GET_ELEM_PTR,
            .data.get_elem_ptr.src_addr = global_value_ref,
            .data.get_elem_ptr.index = SCC_IR_VALUE_TAG_NULLPTR,
        });
    scc_ir_builder_add_instr(builder, pointer_to_global_value);
    return pointer_to_global_value;
}

void scc_ir_builder_begin_func(scc_ir_builder_t *builder,
                               scc_ir_func_ref_t func_ref,
                               const char **param_names) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    builder->current_func = func_ref;

    // 借用 func_ptr 和 func_type 以获取参数类型列表
    scc_ir_func_t *func_ptr = nullptr;
    scc_ir_type_t *func_type = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, func_ptr,
        scc_ir_module_get_func(GET_MODULE(builder), func_ref));
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, func_type,
        scc_ir_module_get_type(GET_MODULE(builder), func_ptr->type));

    if (func_type == nullptr || func_type->tag != SCC_IR_TYPE_FUNC) {
        LOG_ERROR("Invalid function type");
        SCC_IR_BUILDER_END_BORROW(builder); // func_type
        SCC_IR_BUILDER_END_BORROW(builder); // func_ptr
        return;
    }
    if (func_ptr == nullptr) {
        LOG_ERROR("Invalid function reference");
        SCC_IR_BUILDER_END_BORROW(builder); // func_type
        SCC_IR_BUILDER_END_BORROW(builder); // func_ptr
        return;
    }
    if (scc_vec_size(func_ptr->bblocks) != 0 ||
        scc_vec_size(func_ptr->params) != 0) {
        LOG_FATAL("Multiple function definitions");
        SCC_IR_BUILDER_END_BORROW(builder); // func_type
        SCC_IR_BUILDER_END_BORROW(builder); // func_ptr
        return;
    }

    // 释放借用，因为下面要调用 add_value（可能 realloc）
    SCC_IR_BUILDER_END_BORROW(builder); // func_type
    SCC_IR_BUILDER_END_BORROW(builder); // func_ptr
}

void scc_ir_builder_end_func(scc_ir_builder_t *builder) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_func_t *func_ptr = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, func_ptr,
        scc_ir_module_get_func(GET_MODULE(builder), builder->current_func));
    if (func_ptr == nullptr) {
        LOG_FATAL("Invalid function reference");
        SCC_IR_BUILDER_END_BORROW(builder);
        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);
    }
    SCC_IR_BUILDER_END_BORROW(builder);
    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_bblock(scc_ir_builder_t *builder,
                                          const char *label) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_bblock_t bblock = {0};
    if (label) {
        bblock.label = label;
    }
    scc_vec_init(bblock.instrs);
    scc_ir_bblock_ref_t bblock_ref =
        scc_ir_module_add_bblock(GET_MODULE(builder), &bblock);

    // 将基本块添加到当前函数
    scc_ir_func_t *current_func = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, current_func,
        scc_ir_module_get_func(GET_MODULE(builder), builder->current_func));
    if (current_func) {
        scc_vec_push(current_func->bblocks, bblock_ref);
    }
    SCC_IR_BUILDER_END_BORROW(builder);
    return bblock_ref;
}

scc_ir_bblock_ref_t scc_ir_builder_begin_bblock(scc_ir_builder_t *builder,
                                                const char *label) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    builder->current_bblock = scc_ir_builder_bblock(builder, label);
    return builder->current_bblock;
}

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

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

scc_ir_func_ref_t scc_ir_builder_current_bblock(scc_ir_builder_t *builder) {
    return builder->current_bblock;
}

void scc_ir_builder_add_instr(scc_ir_builder_t *builder,
                              scc_ir_value_ref_t instr) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_bblock_t *current_bblock = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, current_bblock,
        scc_ir_module_get_bblock(GET_MODULE(builder), builder->current_bblock));
    if (current_bblock) {
        scc_vec_push(current_bblock->instrs, instr);
    } else {
        LOG_ERROR("Current basic block is not set");
    }
    SCC_IR_BUILDER_END_BORROW(builder);
}

scc_ir_value_ref_t scc_ir_builder_global_alloca(scc_ir_builder_t *builder,
                                                scc_ir_type_ref_t type,
                                                scc_ir_value_ref_t value) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    char *name = scc_malloc(32);
    scc_ir_value_ref_t global_value_ref = scc_ir_module_add_value(
        builder->ctx.module, &(scc_ir_value_t){
                                 .name = name,
                                 .tag = SCC_IR_VALUE_TAG_GLOBAL_ALLOC,
                                 .type = type,
                                 .data.global_alloc.value = value,
                             });
    scc_snprintf(name, 32, "$G%u", global_value_ref);
    scc_vec_push(builder->cprog->global_vals, global_value_ref);
    return global_value_ref;
}

scc_ir_value_ref_t scc_ir_builder_alloca(scc_ir_builder_t *builder,
                                         scc_ir_type_ref_t type,
                                         const char *name) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t alloc_node = {0};
    alloc_node.tag = SCC_IR_VALUE_TAG_ALLOC;
    alloc_node.type = scc_ir_module_add_type(
        GET_MODULE(builder),
        &(scc_ir_type_t){.tag = SCC_IR_TYPE_PTR, .data.pointer.base = type});
    alloc_node.name = name;

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &alloc_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_func_arg_ref(scc_ir_builder_t *builder,
                                               scc_ir_type_ref_t type,
                                               const char *name,
                                               usize arg_idx) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t value = {0};
    value.tag = SCC_IR_VALUE_TAG_FUNC_ARG_REF;
    value.type = type;
    value.name = name;
    value.data.arg_ref.idx = arg_idx;

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &value);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_load(scc_ir_builder_t *builder,
                                       scc_ir_value_ref_t target) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t load_node = {0};
    load_node.tag = SCC_IR_VALUE_TAG_LOAD;
    load_node.data.load.target = target;

    // 借用 ptr_node 和 ptr_type 获取类型信息
    scc_ir_value_t *ptr_node = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, ptr_node,
        scc_ir_module_get_value(GET_MODULE(builder), target));
    if (ptr_node) {
        scc_ir_type_t *ptr_type = nullptr;
        SCC_IR_BUILDER_BEGIN_BORROW(
            builder, ptr_type,
            scc_ir_module_get_type(GET_MODULE(builder), ptr_node->type));
        if (ptr_type && ptr_type->tag == SCC_IR_TYPE_PTR) {
            load_node.type = ptr_type->data.pointer.base;
        }
        SCC_IR_BUILDER_END_BORROW(builder); // ptr_type
    }
    SCC_IR_BUILDER_END_BORROW(builder); // ptr_node

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &load_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_store(scc_ir_builder_t *builder,
                                        scc_ir_value_ref_t target,
                                        scc_ir_value_ref_t value) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    Assert(target != SCC_IR_REF_nullptr && value != SCC_IR_REF_nullptr);
    scc_ir_value_t store_node = {0};
    store_node.tag = SCC_IR_VALUE_TAG_STORE;
    store_node.data.store.target = target;
    store_node.data.store.value = value;

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &store_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_get_elem_ptr(scc_ir_builder_t *builder,
                                               scc_ir_value_ref_t target,
                                               scc_ir_value_ref_t index) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t get_ptr_node = {0};
    get_ptr_node.tag = SCC_IR_VALUE_TAG_GET_ELEM_PTR;
    get_ptr_node.data.get_elem_ptr.src_addr = target;
    get_ptr_node.data.get_elem_ptr.index = index;

    // 借用类型信息
    scc_ir_type_t *type_ref = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, type_ref,
        scc_ir_module_get_type_by_value(GET_MODULE(builder), target));
    Assert(type_ref != nullptr);
    scc_ir_type_t type = *type_ref;     // 拷贝一份，避免后续借用
    SCC_IR_BUILDER_END_BORROW(builder); // type_ref

    if (type.tag == SCC_IR_TYPE_PTR) {
        scc_ir_type_t *base_type = nullptr;
        SCC_IR_BUILDER_BEGIN_BORROW(
            builder, base_type,
            scc_ir_module_get_type(GET_MODULE(builder),
                                   type_ref->data.pointer.base));
        if (base_type->tag == SCC_IR_TYPE_ARRAY) {
            scc_ir_type_t type = (scc_ir_type_t){
                .tag = SCC_IR_TYPE_PTR,
                .data.pointer.base = base_type->data.array.base,
            };
            SCC_IR_BUILDER_END_BORROW(builder); // base_type
            get_ptr_node.type = scc_ir_builder_type(builder, &type);
        } else {
            SCC_IR_BUILDER_END_BORROW(builder); // base_type
            get_ptr_node.type = scc_ir_builder_type(builder, &type);
        }
    } else {
        get_ptr_node.type = scc_ir_builder_type(builder, &type);
    }

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &get_ptr_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_binop(scc_ir_builder_t *builder,
                                        scc_ir_op_type_t op,
                                        scc_ir_value_ref_t lhs,
                                        scc_ir_value_ref_t rhs) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t binop_node = {0};
    binop_node.tag = SCC_IR_VALUE_TAG_OP;
    binop_node.data.op.op = op;
    binop_node.data.op.lhs = lhs;
    binop_node.data.op.rhs = rhs;

    // 借用 lhs_node 获取类型
    scc_ir_value_t *lhs_node = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, lhs_node, scc_ir_module_get_value(GET_MODULE(builder), lhs));
    if (lhs_node) {
        binop_node.type = lhs_node->type;
    }
    SCC_IR_BUILDER_END_BORROW(builder); // lhs_node

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &binop_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_cmp(scc_ir_builder_t *builder,
                                      scc_ir_op_type_t op,
                                      scc_ir_value_ref_t lhs,
                                      scc_ir_value_ref_t rhs) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t cmp_node = {0};
    cmp_node.tag = SCC_IR_VALUE_TAG_OP;
    cmp_node.data.op.op = op;
    cmp_node.data.op.lhs = lhs;
    cmp_node.data.op.rhs = rhs;
    cmp_node.type = 0; // FIXME

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &cmp_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_jump(scc_ir_builder_t *builder,
                                       scc_ir_bblock_ref_t target) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t jump_node = {0};
    jump_node.tag = SCC_IR_VALUE_TAG_JUMP;
    jump_node.data.jump.target_bblock = target;

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &jump_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_branch(scc_ir_builder_t *builder,
                                         scc_ir_value_ref_t cond,
                                         scc_ir_bblock_ref_t true_target,
                                         scc_ir_bblock_ref_t false_target) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t branch_node = {0};
    branch_node.tag = SCC_IR_VALUE_TAG_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_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &branch_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_call(scc_ir_builder_t *builder,
                                       scc_ir_func_ref_t callee,
                                       const scc_ir_value_ref_t *args,
                                       usize arg_count) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t call_node = {0};
    call_node.tag = SCC_IR_VALUE_TAG_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]);
    }

    // 借用 callee_func 和 func_type 获取返回类型
    scc_ir_func_t *callee_func = nullptr;
    SCC_IR_BUILDER_BEGIN_BORROW(
        builder, callee_func,
        scc_ir_module_get_func(GET_MODULE(builder), callee));
    if (callee_func) {
        scc_ir_type_t *func_type = nullptr;
        SCC_IR_BUILDER_BEGIN_BORROW(
            builder, func_type,
            scc_ir_module_get_type(GET_MODULE(builder), callee_func->type));
        if (func_type && func_type->tag == SCC_IR_TYPE_FUNC) {
            call_node.type = func_type->data.function.ret_type;
        }
        SCC_IR_BUILDER_END_BORROW(builder); // func_type
    }
    SCC_IR_BUILDER_END_BORROW(builder); // callee_func

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &call_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_ret(scc_ir_builder_t *builder,
                                      scc_ir_value_ref_t value) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t ret_node = {0};
    ret_node.tag = SCC_IR_VALUE_TAG_RET;
    ret_node.data.ret.ret_val = value;

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &ret_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}

scc_ir_value_ref_t scc_ir_builder_ret_void(scc_ir_builder_t *builder) {
    SCC_IR_BUILDER_CHECK_NO_BORROW(builder);
    scc_ir_value_t ret_node = {0};
    ret_node.tag = SCC_IR_VALUE_TAG_RET;
    ret_node.data.ret.ret_val = 0;

    scc_ir_value_ref_t value_ref =
        scc_ir_module_add_value(GET_MODULE(builder), &ret_node);
    scc_ir_builder_add_instr(builder, value_ref);
    return value_ref;
}