#include <scc_hir_dump.h>
#include <scc_hir_prog.h>
#include <scc_tree_dump.h>

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

void scc_hir_dump_init(scc_hir_dump_t *ctx, scc_tree_dump_t *td,
                       scc_hir_cprog_t *cprog) {
    ctx->dump_ctx = td;
    ctx->cprog = cprog;
}

static const char *get_node_type_str(scc_hir_value_tag_t tag) {
    static const char *node_types[] = {
        [SCC_HIR_VALUE_TAG_NULLPTR] = "NullPtr",
        [SCC_HIR_VALUE_TAG_BUILTIN] = "Builtin",
        [SCC_HIR_VALUE_TAG_INTEGER] = "ConstInt",
        [SCC_HIR_VALUE_TAG_DECIMAL] = "ConstFloat",
        [SCC_HIR_VALUE_TAG_ARRAY] = "ConstArray",
        [SCC_HIR_VALUE_TAG_AGGREGATE] = "Aggregate",
        [SCC_HIR_VALUE_TAG_CONV] = "Convert",
        [SCC_HIR_VALUE_TAG_FUNC_ARG_REF] = "FuncArgRef",
        [SCC_HIR_VALUE_TAG_BLOCK_ARG_REF] = "BlockArgRef",
        [SCC_HIR_VALUE_TAG_ALLOC] = "Alloc",
        [SCC_HIR_VALUE_TAG_GLOBAL_ALLOC] = "GlobalAlloc",
        [SCC_HIR_VALUE_TAG_LOAD] = "Load",
        [SCC_HIR_VALUE_TAG_STORE] = "Store",
        [SCC_HIR_VALUE_TAG_GET_ELEM_PTR] = "GetElemPtr",
        [SCC_HIR_VALUE_TAG_OP] = "Op",
        [SCC_HIR_VALUE_TAG_BRANCH] = "Branch",
        [SCC_HIR_VALUE_TAG_JUMP] = "Jump",
        [SCC_HIR_VALUE_TAG_CALL] = "Call",
        [SCC_HIR_VALUE_TAG_RET] = "Ret",
    };
    if (tag >= 0 && (usize)tag < sizeof(node_types) / sizeof(node_types[0]) &&
        node_types[tag])
        return node_types[tag];
    return "UnknownIRNode";
}

static const char *get_op_str(scc_hir_op_type_t op) {
    static const char *ops[] = {
        [SCC_HIR_OP_EMPTY] = "empty", [SCC_HIR_OP_NEQ] = "!=",
        [SCC_HIR_OP_EQ] = "==",       [SCC_HIR_OP_GT] = ">",
        [SCC_HIR_OP_LT] = "<",        [SCC_HIR_OP_GE] = ">=",
        [SCC_HIR_OP_LE] = "<=",       [SCC_HIR_OP_ADD] = "+",
        [SCC_HIR_OP_SUB] = "-",       [SCC_HIR_OP_MUL] = "*",
        [SCC_HIR_OP_DIV] = "/",       [SCC_HIR_OP_MOD] = "%",
        [SCC_HIR_OP_AND] = "&",       [SCC_HIR_OP_OR] = "|",
        [SCC_HIR_OP_XOR] = "^",       [SCC_HIR_OP_NOT] = "~",
        [SCC_HIR_OP_SHL] = "<<",      [SCC_HIR_OP_SHR] = ">>",
        [SCC_HIR_OP_SAR] = ">>a",     [SCC_HIR_OP_ULT] = "u<",
        [SCC_HIR_OP_ULE] = "u<=",     [SCC_HIR_OP_UGT] = "u>",
        [SCC_HIR_OP_UGE] = "u>=",
    };
    if (op >= 0 && (usize)op < sizeof(ops) / sizeof(ops[0]) && ops[op])
        return ops[op];
    return "<unknown_op>";
}

static const char *get_type_tag_str(scc_hir_type_tag_t tag) {
    static const char *type_tags[] = {
        [SCC_HIR_TYPE_void] = "void",     [SCC_HIR_TYPE_u8] = "u8",
        [SCC_HIR_TYPE_u16] = "u16",       [SCC_HIR_TYPE_u32] = "u32",
        [SCC_HIR_TYPE_u64] = "u64",       [SCC_HIR_TYPE_u128] = "u128",
        [SCC_HIR_TYPE_i8] = "i8",         [SCC_HIR_TYPE_i16] = "i16",
        [SCC_HIR_TYPE_i32] = "i32",       [SCC_HIR_TYPE_i64] = "i64",
        [SCC_HIR_TYPE_i128] = "i128",     [SCC_HIR_TYPE_f16] = "f16",
        [SCC_HIR_TYPE_f32] = "f32",       [SCC_HIR_TYPE_f64] = "f64",
        [SCC_HIR_TYPE_f128] = "f128",     [SCC_HIR_TYPE_PTR] = "ptr",
        [SCC_HIR_TYPE_ARRAY] = "array",   [SCC_HIR_TYPE_FUNC] = "func",
        [SCC_HIR_TYPE_STRUCT] = "struct", [SCC_HIR_TYPE_VECTOR] = "vector",
    };
    if (tag >= 0 && (usize)tag < sizeof(type_tags) / sizeof(type_tags[0]) &&
        type_tags[tag])
        return type_tags[tag];
    return "<unknown_type>";
}

static void dump_type_linear_with_visited(scc_hir_dump_t *ctx,
                                          scc_hir_type_ref_t type_ref,
                                          scc_hashtable_t *visited,
                                          cbool extened) {
    if (!ctx || !type_ref) {
        LOG_ERROR("invalid parameter");
        return;
    }
    // 检查循环
    if (scc_hashtable_get(visited, (void *)(usize)type_ref)) {
        extened = false;
    }

    scc_hir_type_t *type = scc_hir_module_get_type(GET_MODULE(ctx), type_ref);
    if (!type) {
        LOG_ERROR("invalid type ref");
        return;
    }
    switch (type->tag) {
    case SCC_HIR_TYPE_unknown:
    case SCC_HIR_TYPE_void:
    case SCC_HIR_TYPE_i8:
    case SCC_HIR_TYPE_i16:
    case SCC_HIR_TYPE_i32:
    case SCC_HIR_TYPE_i64:
    case SCC_HIR_TYPE_i128:
    case SCC_HIR_TYPE_u8:
    case SCC_HIR_TYPE_u16:
    case SCC_HIR_TYPE_u32:
    case SCC_HIR_TYPE_u64:
    case SCC_HIR_TYPE_u128:
    case SCC_HIR_TYPE_f16:
    case SCC_HIR_TYPE_f32:
    case SCC_HIR_TYPE_f64:
    case SCC_HIR_TYPE_f128:
        scc_tree_dump_value(ctx->dump_ctx, "%s", get_type_tag_str(type->tag));
        break;
    case SCC_HIR_TYPE_ARRAY:
        scc_tree_dump_append(ctx->dump_ctx, "[");
        dump_type_linear_with_visited(ctx, type->data.array.base, visited,
                                      extened);
        scc_tree_dump_append_fmt(ctx->dump_ctx, ", %zu]", type->data.array.len);
        break;
    case SCC_HIR_TYPE_PTR:
        scc_tree_dump_append(ctx->dump_ctx, "*");
        dump_type_linear_with_visited(ctx, type->data.pointer.base, visited,
                                      extened);
        break;
    case SCC_HIR_TYPE_FUNC:
        scc_tree_dump_append(ctx->dump_ctx, "(");
        for (usize i = 0; i < scc_vec_size(type->data.function.params); i++) {
            if (i > 0)
                scc_tree_dump_append(ctx->dump_ctx, ", ");
            dump_type_linear_with_visited(
                ctx, scc_vec_at(type->data.function.params, i), visited,
                extened);
        }
        if (type->data.function.ret_type) {
            scc_tree_dump_append(ctx->dump_ctx, ") -> ");
            dump_type_linear_with_visited(ctx, type->data.function.ret_type,
                                          visited, extened);
        } else {
            scc_tree_dump_append(ctx->dump_ctx, ")");
        }
        break;
    case SCC_HIR_TYPE_STRUCT:
    case SCC_HIR_TYPE_UNION:
        scc_hashtable_set(visited, (void *)(usize)type_ref, (void *)1);
        const char *tag_name =
            type->tag == SCC_HIR_TYPE_STRUCT ? "struct" : "union";
        if (type->name) {
            scc_tree_dump_append_fmt(ctx->dump_ctx, "%s %s {", tag_name,
                                     type->name);
        } else {
            scc_tree_dump_append_fmt(ctx->dump_ctx, "%s T%u {", tag_name,
                                     type_ref);
        }
        if (extened) {
            for (usize i = 0; i < scc_vec_size(type->data.aggregate.fields);
                 i++) {
                dump_type_linear_with_visited(
                    ctx, scc_vec_at(type->data.aggregate.fields, i), visited,
                    extened);
                scc_tree_dump_append(ctx->dump_ctx, ";");
            }
        } else {
            scc_tree_dump_append(ctx->dump_ctx, "...");
        }
        scc_tree_dump_append(ctx->dump_ctx, "}");
        break;
    default:
        LOG_ERROR("invalid type tag");
        break;
    }
}

void scc_hir_dump_type_linear(scc_hir_dump_t *ctx,
                              scc_hir_type_ref_t type_ref) {
    scc_hashtable_t visited;
    scc_hashtable_usize_init(&visited);
    dump_type_linear_with_visited(ctx, type_ref, &visited, false);
    scc_hashtable_drop(&visited);
}
static void scc_hir_dump_types_linear(scc_hir_dump_t *ctx) {
    scc_hir_module_t *module = GET_MODULE(ctx);
    // 遍历所有类型（跳过索引0，保留给无效引用）
    for (usize i = 1; i < scc_vec_size(module->types); i++) {
        scc_hir_type_t *type = &module->types.data[i];
        if (type->tag == SCC_HIR_TYPE_unknown)
            continue;

        scc_hashtable_t visited;
        scc_hashtable_usize_init(&visited);

        scc_tree_dump_begin_line(ctx->dump_ctx);
        scc_tree_dump_node(ctx->dump_ctx, "T%u", (unsigned)i);
        scc_tree_dump_append(ctx->dump_ctx, " = ");
        dump_type_linear_with_visited(ctx, (scc_hir_type_ref_t)i, &visited,
                                      true);

        scc_hashtable_drop(&visited);
    }
}

static void format_ref_or_value(scc_hir_dump_t *ctx,
                                scc_hir_value_ref_t value_ref) {
    scc_hir_value_t *value =
        scc_hir_module_get_value(GET_MODULE(ctx), value_ref);
    if (!value) {
        scc_tree_dump_append_fmt(ctx->dump_ctx, "%%%u", value_ref);
        return;
    }
    if (value->tag == SCC_HIR_VALUE_TAG_INTEGER) {
        // FIXME hack ap
        scc_tree_dump_append_fmt(ctx->dump_ctx, "%d",
                                 value->data.integer.data.digit);
        return;
    }
    if (value->name && value->name[0] != '\0') {
        scc_tree_dump_node(ctx->dump_ctx, "%%%u[%s]", value_ref, value->name);
    } else {
        scc_tree_dump_node(ctx->dump_ctx, "%%%u", value_ref);
    }
    if (value->type != SCC_HIR_REF_nullptr) {
        scc_tree_dump_append(ctx->dump_ctx, ":");
        scc_hir_dump_type_linear(ctx, value->type);
    }
}

void scc_hir_dump_value_linear(scc_hir_dump_t *ctx,
                               scc_hir_value_ref_t value_ref) {
    scc_hir_value_t *value =
        scc_hir_module_get_value(GET_MODULE(ctx), value_ref);
    if (!value) {
        scc_tree_dump_append(ctx->dump_ctx, "<invalid value>\n");
        return;
    }

    cbool needs_equals = (value->tag != SCC_HIR_VALUE_TAG_BRANCH &&
                          value->tag != SCC_HIR_VALUE_TAG_JUMP &&
                          value->tag != SCC_HIR_VALUE_TAG_RET &&
                          value->tag != SCC_HIR_VALUE_TAG_STORE);

    if (needs_equals) {
        format_ref_or_value(ctx, value_ref);
        scc_tree_dump_append(ctx->dump_ctx, " = ");
    }

    switch (value->tag) {
    case SCC_HIR_VALUE_TAG_BUILTIN: {
        scc_tree_dump_append(ctx->dump_ctx, "@scc::");
        switch (value->data.builtin.tag) {
        case SCC_HIR_BUILTIN_TAG_MEMCPY:
            scc_tree_dump_append(ctx->dump_ctx, "memcpy");
            scc_tree_dump_append(ctx->dump_ctx, "(");
            format_ref_or_value(ctx, value->data.builtin.func.memcpy.dest);
            scc_tree_dump_append(ctx->dump_ctx, ", ");
            format_ref_or_value(ctx, value->data.builtin.func.memcpy.src);
            scc_tree_dump_append(ctx->dump_ctx, ", ");
            format_ref_or_value(ctx, value->data.builtin.func.memcpy.size);
            scc_tree_dump_append(ctx->dump_ctx, ")");
            break;
        case SCC_HIR_BUILTIN_TAG_MEMSET:
            scc_tree_dump_append(ctx->dump_ctx, "memset");
            break;
        case SCC_HIR_BUILTIN_TAG_VA_ARG:
            scc_tree_dump_append(ctx->dump_ctx, "va_arg");
            break;
        case SCC_HIR_BUILTIN_TAG_VA_END:
            scc_tree_dump_append(ctx->dump_ctx, "va_end");
            break;
        case SCC_HIR_BUILTIN_TAG_VA_COPY:
            scc_tree_dump_append(ctx->dump_ctx, "va_copy");
            break;
        case SCC_HIR_BUILTIN_TAG_VA_START:
            scc_tree_dump_append(ctx->dump_ctx, "va_start");
            break;
        default:
            Panic("Unknown builtin tag");
            break;
        }
        break;
    }
    case SCC_HIR_VALUE_TAG_INTEGER:
        // 值已经在 format 中输出，这里不需要再做
        break;
    case SCC_HIR_VALUE_TAG_AGGREGATE:
        scc_tree_dump_append(ctx->dump_ctx, "{");
        if (scc_vec_size(value->data.aggregate.fields) > 0) {
            scc_tree_dump_push(ctx->dump_ctx, true);
            scc_vec_foreach(value->data.aggregate.fields, i) {
                scc_tree_dump_begin_line(ctx->dump_ctx);
                scc_hir_dump_value_linear(
                    ctx, scc_vec_at(value->data.aggregate.fields, i));
            }
            scc_tree_dump_pop(ctx->dump_ctx);
            scc_tree_dump_begin_line(ctx->dump_ctx);
        }
        scc_tree_dump_append(ctx->dump_ctx, "}");
        return;
    case SCC_HIR_VALUE_TAG_ALLOC:
        scc_tree_dump_append(ctx->dump_ctx, "alloc");
        break;
    case SCC_HIR_VALUE_TAG_LOAD:
        scc_tree_dump_append(ctx->dump_ctx, "load ");
        format_ref_or_value(ctx, value->data.load.target);
        break;
    case SCC_HIR_VALUE_TAG_STORE:
        scc_tree_dump_append(ctx->dump_ctx, "store ");
        format_ref_or_value(ctx, value->data.store.value);
        scc_tree_dump_append(ctx->dump_ctx, " -> ");
        format_ref_or_value(ctx, value->data.store.target);
        break;
    case SCC_HIR_VALUE_TAG_GET_ELEM_PTR:
        scc_tree_dump_append(ctx->dump_ctx, "getelemptr ");
        format_ref_or_value(ctx, value->data.get_elem_ptr.src_addr);
        scc_tree_dump_append(ctx->dump_ctx, ", ");
        format_ref_or_value(ctx, value->data.get_elem_ptr.index);
        break;
    case SCC_HIR_VALUE_TAG_OP:
        format_ref_or_value(ctx, value->data.op.lhs);
        scc_tree_dump_append_fmt(ctx->dump_ctx, " %s ",
                                 get_op_str(value->data.op.op));
        format_ref_or_value(ctx, value->data.op.rhs);
        break;
    case SCC_HIR_VALUE_TAG_BRANCH:
        if (value->data.branch.cond) {
            scc_tree_dump_append(ctx->dump_ctx, "br ");
            format_ref_or_value(ctx, value->data.branch.cond);
            scc_tree_dump_append_fmt(ctx->dump_ctx,
                                     ", label %%L%u, label %%L%u",
                                     value->data.branch.true_bblock,
                                     value->data.branch.false_bblock);
        } else {
            scc_tree_dump_append_fmt(ctx->dump_ctx, "br label %%L%u",
                                     value->data.branch.true_bblock);
        }
        break;
    case SCC_HIR_VALUE_TAG_JUMP:
        scc_tree_dump_append_fmt(ctx->dump_ctx, "jmp label %%L%u",
                                 value->data.jump.target_bblock);
        break;
    case SCC_HIR_VALUE_TAG_CALL: {
        scc_hir_func_t *func = scc_hir_module_get_func(
            GET_MODULE(ctx), value->data.call.callee.func_ref);
        scc_tree_dump_append_fmt(ctx->dump_ctx, "call @%s(",
                                 func ? (func->name ? func->name : "<unnamed>")
                                      : "<invalid>");
        for (usize i = 0; i < scc_vec_size(value->data.call.args); i++) {
            if (i > 0)
                scc_tree_dump_append(ctx->dump_ctx, ", ");
            format_ref_or_value(ctx, scc_vec_at(value->data.call.args, i));
        }
        scc_tree_dump_append(ctx->dump_ctx, ")");
        break;
    }
    case SCC_HIR_VALUE_TAG_RET:
        if (value->data.ret.ret_val != 0) {
            scc_tree_dump_append(ctx->dump_ctx, "ret ");
            format_ref_or_value(ctx, value->data.ret.ret_val);
        } else {
            scc_tree_dump_append(ctx->dump_ctx, "ret void");
        }
        break;
    case SCC_HIR_VALUE_TAG_FUNC_ARG_REF:
        scc_tree_dump_append_fmt(ctx->dump_ctx, "arg[%zu]",
                                 value->data.arg_ref.idx);
        break;
    case SCC_HIR_VALUE_TAG_GLOBAL_ALLOC:
        scc_tree_dump_append_fmt(ctx->dump_ctx, "global %s:", value->name);
        scc_tree_dump_begin_line(ctx->dump_ctx);
        scc_tree_dump_append(ctx->dump_ctx, "  ");
        if (value->data.global_alloc.value) {
            scc_hir_dump_value_linear(ctx, value->data.global_alloc.value);
        } else {
            scc_tree_dump_append(ctx->dump_ctx, "<zero initializer>");
        }
        break;
    case SCC_HIR_VALUE_TAG_ARRAY:
        scc_tree_dump_append(ctx->dump_ctx, "const_array ");
        scc_hir_dump_type_linear(ctx, value->data.const_array.base_type);
        scc_tree_dump_append(ctx->dump_ctx, " [");
        scc_vec_foreach(value->data.const_array.fields, i) {
            u8 ch = scc_vec_at(value->data.const_array.fields, i);
            scc_tree_dump_append_fmt(ctx->dump_ctx, " `%c`, ", ch ? ch : ' ');
        }
        scc_tree_dump_append(ctx->dump_ctx, " ]");
        break;
    case SCC_HIR_VALUE_TAG_CONV:
        static const char *conv_name[] = {
            [SCC_HIR_CONV_NONE] = "null",
            [SCC_HIR_CONV_SEXT] = "sext",
            [SCC_HIR_CONV_ZEXT] = "zext",
            [SCC_HIR_CONV_TRUNC] = "trunc",
        };
        scc_tree_dump_append_fmt(ctx->dump_ctx, "conv.%s(",
                                 conv_name[value->data.conv.conv_type]);
        scc_hir_dump_type_linear(ctx, value->data.conv.target_type);
        scc_tree_dump_append(ctx->dump_ctx, ") ");
        format_ref_or_value(ctx, value->data.conv.operand);
        break;
    default:
        scc_tree_dump_append_fmt(ctx->dump_ctx, "<%s value %u>",
                                 get_node_type_str(value->tag), value_ref);
        break;
    }
}

void scc_hir_dump_bblock_linear(scc_hir_dump_t *ctx,
                                scc_hir_bblock_ref_t bblock_ref) {
    scc_tree_dump_begin_line(ctx->dump_ctx);
    scc_hir_bblock_t *bblock =
        scc_hir_module_get_bblock(GET_MODULE(ctx), bblock_ref);
    if (!bblock) {
        scc_tree_dump_append(ctx->dump_ctx, "<invalid block>");
        return;
    }
    if (bblock->name && bblock->name[0] != '\0')
        scc_tree_dump_append_fmt(ctx->dump_ctx, "%%L%d %s:", bblock_ref,
                                 bblock->name);
    else
        scc_tree_dump_append_fmt(ctx->dump_ctx, "%%L%d <unnamed>:", bblock_ref);

    scc_hir_value_ref_vec_t *values = SCC_HIR_BBLOCK_VALUES(bblock);
    for (usize i = 0; i < scc_vec_size(*values); i++) {
        scc_tree_dump_begin_line(ctx->dump_ctx);
        scc_tree_dump_append(ctx->dump_ctx, "  ");
        scc_hir_dump_value_linear(ctx, scc_vec_at(*values, i));
    }
}

void scc_hir_dump_func_linear(scc_hir_dump_t *ctx, scc_hir_func_ref_t func_ref,
                              int is_decl) {
    scc_hir_func_t *func = scc_hir_module_get_func(GET_MODULE(ctx), func_ref);
    scc_tree_dump_begin_line(ctx->dump_ctx);
    if (!func) {
        scc_tree_dump_append(ctx->dump_ctx, "<invalid function>");
        return;
    }
    scc_tree_dump_node(ctx->dump_ctx, "func @%s",
                       (func->name && func->name[0]) ? func->name
                                                     : "<unnamed>");

    if (SCC_HIR_FUNC_META(func)->is_variadic) {
        scc_tree_dump_append(ctx->dump_ctx, "(...)");
    } else {
        scc_tree_dump_append(ctx->dump_ctx, "()");
    }

    scc_tree_dump_append(ctx->dump_ctx, ": ");
    scc_hir_dump_type_linear(ctx, SCC_HIR_FUNC_META(func)->type);

    if (is_decl) {
        scc_tree_dump_append(ctx->dump_ctx, ";");
        return;
    }

    scc_tree_dump_append(ctx->dump_ctx, " {");
    for (usize i = 0; i < scc_vec_size(func->bblocks); i++) {
        scc_hir_dump_bblock_linear(ctx, scc_vec_at(func->bblocks, i));
        scc_tree_dump_append(ctx->dump_ctx, "");
    }
    scc_tree_dump_begin_line(ctx->dump_ctx);
    scc_tree_dump_append(ctx->dump_ctx, "}");
}

void scc_hir_dump_cprog_linear(scc_hir_dump_t *ctx) {
    scc_hir_dump_types_linear(ctx);
    for (usize i = 0; i < scc_vec_size(ctx->cprog->global_vals); i++) {
        scc_tree_dump_begin_line(ctx->dump_ctx);
        scc_hir_dump_value_linear(ctx, scc_vec_at(ctx->cprog->global_vals, i));
    }
    for (usize i = 0; i < scc_vec_size(ctx->cprog->func_decls); i++) {
        scc_hir_func_ref_t func_decl = scc_vec_at(ctx->cprog->func_decls, i);
        scc_hir_func_t *func =
            scc_hir_module_get_func(GET_MODULE(ctx), func_decl);
        if (func && scc_vec_size(func->bblocks) == 0)
            scc_hir_dump_func_linear(ctx, func_decl, 1);
    }
    for (usize i = 0; i < scc_vec_size(ctx->cprog->func_defs); i++) {
        scc_hir_dump_func_linear(ctx, scc_vec_at(ctx->cprog->func_defs, i), 0);
    }
}