#include <amd64/scc_amd64.h>
#include <amd64/scc_amd64_abi.h>
#include <reg_alloc.h>
#include <scc_ir2mcode.h>

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

static void parse_location(scc_ir2mcode_ctx_t *ctx, scc_reg_loc_t *loc,
                           scc_ir_node_ref_t node_ref) {
    Assert(ctx != null && loc != null);
    scc_ir_node_t *node = scc_ir_module_get_node(GET_MODULE(ctx), node_ref);
    if (node == null) {
        LOG_FATAL("invalid node ref");
        UNREACHABLE();
        return;
    }
    usize idx = 0;
    switch (node->tag) {
    case SCC_IR_NODE_CONST_INT:
        scc_ir_type_t *type =
            scc_ir_module_get_type(GET_MODULE(ctx), node->type);
        Assert(type != 0);
        Assert(type->tag == SCC_IR_TYPE_u32 || type->tag == SCC_IR_TYPE_i32);
        *loc = (scc_reg_loc_t){
            .kind = SCC_REG_KIND_IMM,
            .idx = (usize)node->data.const_int.int32,
        };
        return;
    case SCC_IR_NODE_CONST_UINT:
    case SCC_IR_NODE_CONST_FLOAT:
        TODO();
        break;
    case SCC_IR_NODE_FUNC_ARG_REF: {
        scc_ir_type_t *type =
            scc_ir_module_get_type(GET_MODULE(ctx), node->type);
        Assert(type != 0);
        scc_reg_loc_t arg_loc;
        // arg_loc.kind = SCC_REG_KIND_FUNC_ARG;
        // arg_loc.idx = node->data.arg_ref.idx;
        arg_loc.kind = SCC_REG_KIND_STACK;
        arg_loc.idx = 8 * node->data.arg_ref.idx;
        *loc = arg_loc;
        return;
    }
    default:
        idx = (usize)scc_hashtable_get(ctx->noderef2regloc,
                                       (void *)(usize)node_ref);
        break;
    }

    Assert(idx > 0 && idx <= scc_vec_size(ctx->reg_alloc.reg_loc_vec));
    *loc = scc_vec_at(ctx->reg_alloc.reg_loc_vec, idx - 1);
}

static void load_value_to_reg(scc_mcode_t *mcode, scc_reg_loc_t *loc, int reg) {
    switch (loc->kind) {
    case SCC_REG_KIND_GPR:
        if (loc->idx != reg) {
            scc_mcode_amd64_mov_r64_r64(mcode, reg, loc->idx);
        } else {
            TODO();
        }
        break;
    case SCC_REG_KIND_STACK:
        // FIXME -8 for rdp
        scc_mcode_amd64_mov_r64_m64_disp32(mcode, reg, SCC_AMD64_RBP,
                                           -loc->idx - 8);
        break;
    case SCC_REG_KIND_IMM:
        scc_mcode_amd64_mov_r64_imm64(mcode, reg, loc->idx); // 或 imm32
        break;
    default:
        LOG_FATAL("unsupported location");
    }
}

static void store_value_from_reg(scc_mcode_t *mcode, scc_reg_loc_t *loc,
                                 int reg) {
    switch (loc->kind) {
    case SCC_REG_KIND_GPR:
        if (loc->idx != reg) {
            scc_mcode_amd64_mov_r64_r64(mcode, loc->idx, reg);
        }
        break;
    case SCC_REG_KIND_STACK:
        // FIXME -8 for rdp
        scc_mcode_amd64_mov_m64_disp32_r64(mcode, SCC_AMD64_RBP, -loc->idx - 8,
                                           reg);
        break;
    case SCC_REG_KIND_IMM:
        LOG_FATAL("cannot store to immediate");
        break;
    default:
        LOG_FATAL("unsupported location");
        break;
    }
}

// 临时存储待修补条目
typedef struct {
    usize pos;
    usize target_bb_ref;
} patch_t;
typedef SCC_VEC(patch_t) patch_vec_t;

static void parse_node(scc_ir2mcode_ctx_t *ctx, scc_ir_bblock_ref_t node_ref,
                       patch_vec_t *patches) {
    scc_ir_node_t *node = scc_ir_module_get_node(GET_MODULE(ctx), node_ref);
    if (node == null) {
        LOG_ERROR("invalid node ref");
        return;
    }

    switch (node->tag) {
    case SCC_IR_NODE_CONV: ///< 类型转换
        LOG_FATAL("Unsupported node type: %d", node->tag);
        break;
    ///< 函数参数引用
    case SCC_IR_NODE_FUNC_ARG_REF:
        ///< ABI
        break;
    case SCC_IR_NODE_BLOCK_ARG_REF: ///< 基本块参数引用
    case SCC_IR_NODE_ALLOC:         ///< 分配内存(stack)
    case SCC_IR_NODE_GLOBAL_ALLOC:  ///< 全局分配(bss)
        break;
    case SCC_IR_NODE_LOAD: ///< 加载数据
    {
        // node->data.load.target
        scc_reg_loc_t from;
        scc_reg_loc_t to;
        parse_location(ctx, &from, node->data.load.target);
        parse_location(ctx, &to, node_ref);
        load_value_to_reg(&ctx->sect_mcode, &from, SCC_AMD64_RAX);
        store_value_from_reg(&ctx->sect_mcode, &to, SCC_AMD64_RAX);
        break;
    }
    case SCC_IR_NODE_STORE: ///< 存储数据
    {
        scc_reg_loc_t from;
        scc_reg_loc_t to;
        parse_location(ctx, &from, node->data.store.value);
        parse_location(ctx, &to, node->data.store.target);
        load_value_to_reg(&ctx->sect_mcode, &from, SCC_AMD64_RAX);
        store_value_from_reg(&ctx->sect_mcode, &to, SCC_AMD64_RAX);
        break;
    }
    case SCC_IR_NODE_GET_PTR:      ///< 获取指针
    case SCC_IR_NODE_GET_ELEM_PTR: ///< 获取元素指针(used by array)
        TODO();
    ///< 二元运算
    case SCC_IR_NODE_OP: {
        scc_reg_loc_t loc_lhs;
        parse_location(ctx, &loc_lhs, node->data.op.lhs);
        scc_reg_loc_t loc_rhs;
        parse_location(ctx, &loc_rhs, node->data.op.rhs);
        scc_reg_loc_t loc_res;
        parse_location(ctx, &loc_res, node_ref);

        // 将左操作数加载到 RAX（临时结果寄存器）
        load_value_to_reg(&ctx->sect_mcode, &loc_lhs, SCC_AMD64_RAX);
        // 将右操作数加载到 RCX
        load_value_to_reg(&ctx->sect_mcode, &loc_rhs, SCC_AMD64_RCX);
        switch (node->data.op.op) {
        case SCC_IR_OP_ADD:
            scc_mcode_amd64_add_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            break;
        case SCC_IR_OP_SUB:
            scc_mcode_amd64_sub_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            break;
        case SCC_IR_OP_MUL:
            scc_mcode_amd64_mul_r64(&ctx->sect_mcode, SCC_AMD64_RCX);
            break;
            // [SCC_IR_OP_EMPTY] = "empty", [SCC_IR_OP_NEQ] = "!=",
            // [SCC_IR_OP_EQ] = "==",       [SCC_IR_OP_GT] = ">",
            // [SCC_IR_OP_LT] = "<",        [SCC_IR_OP_GE] = ">=",
            // [SCC_IR_OP_LE] = "<=",       [SCC_IR_OP_ADD] = "+",
            // [SCC_IR_OP_SUB] = "-",       [SCC_IR_OP_MUL] = "*",
            // [SCC_IR_OP_DIV] = "/",       [SCC_IR_OP_MOD] = "%",
            // [SCC_IR_OP_AND] = "&",       [SCC_IR_OP_OR] = "|",
            // [SCC_IR_OP_XOR] = "^",       [SCC_IR_OP_NOT] = "~",
            // [SCC_IR_OP_SHL] = "<<",      [SCC_IR_OP_SHR] = ">>",
            // [SCC_IR_OP_SAR] = ">>a", // Arithmetic shift right
        case SCC_IR_OP_NEQ:
            scc_mcode_amd64_cmp_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            scc_mcode_amd64_setcc_r8(&ctx->sect_mcode, SCC_AMD64_COND_NE,
                                     SCC_AMD64_RAX);
            scc_mcode_amd64_movzx_r64_r8(&ctx->sect_mcode, SCC_AMD64_RAX,
                                         SCC_AMD64_RAX);
            break;
        case SCC_IR_OP_EQ:
            scc_mcode_amd64_cmp_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            scc_mcode_amd64_setcc_r8(&ctx->sect_mcode, SCC_AMD64_COND_E,
                                     SCC_AMD64_RAX);
            scc_mcode_amd64_movzx_r64_r8(&ctx->sect_mcode, SCC_AMD64_RAX,
                                         SCC_AMD64_RAX);
            break;
        case SCC_IR_OP_GT:
            scc_mcode_amd64_cmp_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            scc_mcode_amd64_setcc_r8(&ctx->sect_mcode, SCC_AMD64_COND_G,
                                     SCC_AMD64_RAX);
            scc_mcode_amd64_movzx_r64_r8(&ctx->sect_mcode, SCC_AMD64_RAX,
                                         SCC_AMD64_RAX);
            break;
        case SCC_IR_OP_LT:
            scc_mcode_amd64_cmp_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            scc_mcode_amd64_setcc_r8(&ctx->sect_mcode, SCC_AMD64_COND_L,
                                     SCC_AMD64_RAX);
            scc_mcode_amd64_movzx_r64_r8(&ctx->sect_mcode, SCC_AMD64_RAX,
                                         SCC_AMD64_RAX);
            break;
        case SCC_IR_OP_GE:
            scc_mcode_amd64_cmp_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            scc_mcode_amd64_setcc_r8(&ctx->sect_mcode, SCC_AMD64_COND_GE,
                                     SCC_AMD64_RAX);
            scc_mcode_amd64_movzx_r64_r8(&ctx->sect_mcode, SCC_AMD64_RAX,
                                         SCC_AMD64_RAX);
            break;
        case SCC_IR_OP_LE:
            scc_mcode_amd64_cmp_r64_r64(&ctx->sect_mcode, SCC_AMD64_RAX,
                                        SCC_AMD64_RCX);
            scc_mcode_amd64_setcc_r8(&ctx->sect_mcode, SCC_AMD64_COND_LE,
                                     SCC_AMD64_RAX);
            scc_mcode_amd64_movzx_r64_r8(&ctx->sect_mcode, SCC_AMD64_RAX,
                                         SCC_AMD64_RAX);
            break;
        default:
            LOG_FATAL("unknown op: %d", node->data.op.op);
            break;
        }
        // 将 RAX 中的结果存储到 res 位置
        store_value_from_reg(&ctx->sect_mcode, &loc_res, SCC_AMD64_RAX);
        break;
    }
    ///< 有条件分支
    case SCC_IR_NODE_BRANCH: {
        scc_reg_loc_t loc;
        parse_location(ctx, &loc, node->data.branch.cond);
        // (void)loc;
        load_value_to_reg(&ctx->sect_mcode, &loc, SCC_AMD64_RAX);
        scc_mcode_amd64_cmp_r64_imm32(&ctx->sect_mcode, SCC_AMD64_RAX, 0);

        scc_mcode_amd64_jcc_rel32(&ctx->sect_mcode, SCC_AMD64_COND_NE, 0);
        patch_t patch_true = {.pos = scc_vec_size(ctx->sect_mcode.mcode),
                              .target_bb_ref =
                                  (usize)node->data.branch.true_bblock};
        scc_vec_push(*patches, patch_true);
        scc_mcode_amd64_jmp_rel32(&ctx->sect_mcode, 0);
        patch_t patch_false = {.pos = scc_vec_size(ctx->sect_mcode.mcode),
                               .target_bb_ref =
                                   (usize)node->data.branch.false_bblock};
        scc_vec_push(*patches, patch_false);
        break;
    }
    ///< 无条件跳转
    case SCC_IR_NODE_JUMP: {
        scc_mcode_amd64_jmp_rel32(&ctx->sect_mcode, 0);
        usize pos = scc_vec_size(ctx->sect_mcode.mcode);
        patch_t patch = {.pos = pos,
                         .target_bb_ref = (usize)node->data.jump.target_bblock};
        scc_vec_push(*patches, patch);
        break;
    }
    ///< 调用函数
    case SCC_IR_NODE_CALL: {
        scc_reg_loc_t loc;
        /*
        ABI
        RAX	不稳定的	返回值寄存器
        RCX	不稳定的	第一个整型自变量
        RDX	不稳定的	第二个整型自变量
        R8	不稳定的	第三个整型自变量
        R9	不稳定的	第四个整型自变量
        */
        scc_vec_foreach(node->data.call.args, i) {
            parse_location(ctx, &loc, scc_vec_at(node->data.call.args, i));
            if (i == 0) {
                load_value_to_reg(&ctx->sect_mcode, &loc, SCC_AMD64_RCX);
            } else if (i == 1) {
                load_value_to_reg(&ctx->sect_mcode, &loc, SCC_AMD64_RDX);
            } else if (i == 2) {
                load_value_to_reg(&ctx->sect_mcode, &loc, SCC_AMD64_R8);
            } else if (i == 3) {
                load_value_to_reg(&ctx->sect_mcode, &loc, SCC_AMD64_R9);
            } else {
                LOG_FATAL("not support more than 4 args");
            }
            // scc_mcode_amd64_push_r64();
        }

        scc_ir_func_t *func =
            scc_ir_module_get_func(GET_MODULE(ctx), node->data.call.callee);
        if (!func) {
            LOG_ERROR("invalid function reference");
            return;
        }

        scc_mcode_amd64_call_rel32(&ctx->sect_mcode, 0);
        usize sym_idx = sccf_builder_get_symbol_idx(ctx->builder, func->name);
        Assert(sym_idx != 0);
        sccf_builder_add_reloc(
            ctx->builder, (sccf_reloc_t){
                              .reloc_type = SCCF_RELOC_TYPE_REL,
                              .offset = scc_vec_size(ctx->sect_mcode.mcode) - 4,
                              .addend = 4,
                              .sect_type = SCCF_SECT_CODE,
                              .sym_idx = sym_idx,
                          });

        // 处理返回值
        scc_ir_type_t *func_type =
            scc_ir_module_get_type(GET_MODULE(ctx), func->type);
        Assert(func_type);
        scc_ir_type_t *ret_type = scc_ir_module_get_type(
            GET_MODULE(ctx), func_type->data.function.ret_type);
        if (ret_type && ret_type->tag != SCC_IR_TYPE_void) {
            parse_location(ctx, &loc, node_ref);
            store_value_from_reg(&ctx->sect_mcode, &loc, SCC_AMD64_RAX);
        }
        break;
    }
    ///< 函数返回
    case SCC_IR_NODE_RET: {
        if (node->data.ret.ret_val) {
            scc_reg_loc_t loc;
            parse_location(ctx, &loc, node->data.ret.ret_val);
            load_value_to_reg(&ctx->sect_mcode, &loc, SCC_AMD64_RAX);
        }
        scc_mcode_amd64_add_rsp_imm32(&ctx->sect_mcode, ctx->stack_size);
        scc_mcode_amd64_pop_r64(&ctx->sect_mcode, SCC_AMD64_RBP);
        scc_mcode_amd64_ret(&ctx->sect_mcode);
        break;
    }
    default:
        LOG_FATAL("unknown node type: %d", node->tag);
        UNREACHABLE();
        break;
    }
}

static void parse_bblock(scc_ir2mcode_ctx_t *ctx, scc_ir_bblock_t *bblock,
                         patch_vec_t *patches) {

    // 打印基本块中的每条指令
    for (usize i = 0; i < scc_vec_size(bblock->instrs); i++) {
        scc_ir_node_ref_t node_ref = scc_vec_at(bblock->instrs, i);
        parse_node(ctx, node_ref, patches);
    }
}

static u32 hash_func(const void *key) { return (u32)(usize)key; }
static int equal_func(const void *key1, const void *key2) {
    return (usize)key1 - (usize)key2;
}

static void parse_function(scc_ir2mcode_ctx_t *ctx, scc_ir_func_t *func) {
    ctx->noderef2regloc = scc_reg_alloc(&ctx->reg_alloc, func);
    // 对齐到 16 字节
    // FIXME
    ctx->stack_size += 8; ///< for rbp
    ctx->stack_size = (ctx->reg_alloc.alloc_stack_size + 15) & ~15;

    usize bblock_cnt = scc_vec_size(func->bblocks);
    usize *bblock_offsets = scc_calloc(bblock_cnt, sizeof(usize));
    // 建立 bblock_ref -> id 的映射
    scc_hashtable_t ref2id;
    scc_hashtable_init(&ref2id, hash_func, equal_func);
    for (usize i = 0; i < bblock_cnt; i++) {
        scc_ir_bblock_ref_t ref = scc_vec_at(func->bblocks, i);
        scc_hashtable_set(&ref2id, (void *)(usize)ref, (void *)i);
    }

    patch_vec_t patches;
    scc_vec_init(patches);

    scc_mcode_amd64_push_r64(&ctx->sect_mcode, SCC_AMD64_RBP);
    scc_mcode_amd64_sub_rsp_imm32(&ctx->sect_mcode, ctx->stack_size);
    scc_mcode_amd64_lea_r64_m64_disp32(&ctx->sect_mcode, SCC_AMD64_RBP,
                                       SCC_AMD64_RSP, ctx->stack_size);
    scc_reg_loc_t loc;
    scc_vec_foreach(func->params, i) {
        // scc_ir_node_t *param =
        //     scc_ir_module_get_node(GET_MODULE(ctx), );
        scc_ir_node_ref_t node_ref = scc_vec_at(func->params, i);
        parse_location(ctx, &loc, node_ref);
        if (i == 0) {
            store_value_from_reg(&ctx->sect_mcode, &loc, SCC_AMD64_RCX);
        } else if (i == 1) {
            store_value_from_reg(&ctx->sect_mcode, &loc, SCC_AMD64_RDX);
        } else if (i == 2) {
            store_value_from_reg(&ctx->sect_mcode, &loc, SCC_AMD64_R8);
        } else if (i == 3) {
            store_value_from_reg(&ctx->sect_mcode, &loc, SCC_AMD64_R9);
        } else {
            LOG_FATAL("not support more than 4 args");
        }
        // scc_mcode_amd64_push_r64();
    }

    for (usize i = 0; i < scc_vec_size(func->bblocks); i++) {
        scc_ir_bblock_ref_t bblock_ref = scc_vec_at(func->bblocks, i);
        scc_ir_bblock_t *bblock =
            scc_ir_module_get_bblock(GET_MODULE(ctx), bblock_ref);

        if (bblock == null) {
            LOG_FATAL("<invalid block>\n");
            return;
        }

        bblock_offsets[i] = scc_vec_size(ctx->sect_mcode.mcode);
        parse_bblock(ctx, bblock, &patches);
    }

    // 回填所有跳转偏移
    u8 *buf = scc_vec_unsafe_get_data(ctx->sect_mcode.mcode);
    scc_vec_foreach(patches, idx) {
        patch_t *p = &scc_vec_at(patches, idx);
        usize target_id =
            (usize)scc_hashtable_get(&ref2id, (void *)(usize)p->target_bb_ref);
        usize target_off = bblock_offsets[target_id];
        usize next_off = p->pos;
        i32 rel = (i32)(target_off - next_off);
        // FIXME 写入到指令的偏移字段（小端）
        *(u32 *)(&buf[p->pos - 4]) = rel;
    }
    scc_free(bblock_offsets);
    scc_vec_free(patches);
    scc_hashtable_drop(&ref2id);
}

void scc_ir2amd64(scc_ir2mcode_ctx_t *ctx) {
    scc_reg_alloc_init(&ctx->reg_alloc, scc_reg_alloc_with_stack,
                       GET_MODULE(ctx));

    scc_vec_foreach(ctx->cprog->func_decls, i) {
        scc_ir_node_ref_t func_ref = scc_vec_at(ctx->cprog->func_decls, i);
        scc_ir_func_t *func = scc_ir_module_get_func(GET_MODULE(ctx), func_ref);
        if (!func) {
            LOG_ERROR("invalid function reference");
            return;
        }

        sccf_sym_t sym = {0};
        if (scc_vec_size(func->bblocks)) {
            sym = (sccf_sym_t){
                .sccf_sect_offset = 0,
                .sccf_sect_type = SCCF_SECT_CODE,
                .sccf_sym_bind = SCCF_SYM_BIND_GLOBAL,
                .sccf_sym_size = 0,
                .sccf_sym_type = SCCF_SYM_TYPE_FUNC,
                .sccf_sym_vis = SCCF_SYM_VIS_DEFAULT,
            };
        } else {
            sym = (sccf_sym_t){
                .sccf_sect_offset = 0,
                .sccf_sect_type = SCCF_SECT_NONE,
                .sccf_sym_bind = SCCF_SYM_BIND_GLOBAL,
                .sccf_sym_size = 0,
                .sccf_sym_type = SCCF_SYM_TYPE_EXTERN,
                .sccf_sym_vis = SCCF_SYM_VIS_DEFAULT,
            };
        }
        usize sym_idx = sccf_builder_add_symbol(ctx->builder, func->name, &sym);
        Assert(sym_idx != 0);
    }

    scc_vec_foreach(ctx->cprog->func_defs, i) {
        scc_ir_node_ref_t func_ref = scc_vec_at(ctx->cprog->func_defs, i);
        scc_ir_func_t *func = scc_ir_module_get_func(GET_MODULE(ctx), func_ref);
        if (!func) {
            LOG_ERROR("invalid function reference");
            return;
        }
        sccf_sym_t *sym =
            sccf_builder_get_symbol_unsafe(ctx->builder, func->name);
        Assert(sym != null);
        sym->sccf_sect_offset = scc_vec_size(ctx->sect_mcode.mcode);
        parse_function(ctx, func);
    }

    u8 *buf = scc_vec_unsafe_get_data(ctx->sect_mcode.mcode);
    scc_vec_foreach(ctx->builder->relocs, i) {
        sccf_reloc_t *reloc = &scc_vec_at(ctx->builder->relocs, i);
        if (reloc->sym_idx == 0) {
            Panic("relocate to an invalid symbol");
        }
        sccf_sym_t *sym = &scc_vec_at(ctx->builder->symtab, reloc->sym_idx);
        if (sym->sccf_sym_type != SCCF_SYM_TYPE_EXTERN) {
            Assert(reloc->reloc_type == SCCF_RELOC_TYPE_REL);

            Assert(sym->sccf_sect_type == SCCF_SECT_CODE);
            Assert(sym->sccf_sym_type == SCCF_SYM_TYPE_FUNC);
            i64 target_off = sym->sccf_sect_offset;
            i64 next_off = reloc->offset + reloc->addend;
            i32 rel = (i32)(target_off - next_off);
            // FIXME 写入到指令的偏移字段（小端）
            *(i32 *)(&buf[reloc->offset]) = rel;

            reloc->reloc_type = SCCF_RELOC_TYPE_EMPTY;
        }
    }

    sccf_sect_data_t text_section;
    scc_vec_unsafe_from_buffer(text_section,
                               scc_vec_unsafe_get_data(ctx->sect_mcode.mcode),
                               scc_vec_size(ctx->sect_mcode.mcode));
    sccf_builder_add_text_section(ctx->builder, &text_section);
    sccf_sect_data_t data_section;
    scc_vec_init(data_section);
    sccf_builder_add_data_section(ctx->builder, &data_section);
    // FIXME
    ctx->builder->entry_symbol_name = "main";
}