#include <scc_mir_module.h>
#include <x86/scc_x86_iform.h>

#include <arch/scc_x86_isel.h>
#include <core_pass/scc_prolog_epilog.h>
#include <target/scc_win64.h>

static const int WIN64_STACK_ALIGN = 16;

static void frame_alloc_impl(scc_frame_layout_t *ctx,
                             scc_mir_module_t *mir_module,
                             scc_mir_func_t *mir_func) {
    ctx->offset = 8;
    scc_mir_func_meta_t *func_meta = SCC_MIR_FUNC_META(mir_func);
    scc_vec_foreach(mir_func->bblocks, i) {
        scc_cfg_bblock_id_t id = scc_vec_at(mir_func->bblocks, i);
        scc_cfg_bblock_t *bb =
            scc_cfg_module_unsafe_get_bblock(&mir_module->cfg_module, id);
        scc_mir_x86_instr_vec_t *instrs = SCC_MIR_X86_BBLOCK_INSTRS(bb);
        scc_vec_foreach(*instrs, j) {
            scc_mir_x86_instr_t *ins = &scc_vec_at(*instrs, j);
            for (int k = 0; k < ins->x86_instr.num_operands; k++) {
                scc_x86_operand_value_t *op = &ins->x86_instr.operands[k];
                if (scc_x86_op_is_slot(op)) {
                    int slot_id = scc_x86_op_slot_id(op);
                    scc_mir_stack_slot_t *slot =
                        scc_mir_unsafe_slot(mir_func, slot_id);
                    if (slot->offset == 0) {
                        ctx->offset += slot->size;
                        slot->offset = ctx->offset;
                    }
                    *op = scc_x86_op_mem((scc_x86_mem_t){
                        .seg = SCC_X86_REG_INVALID,
                        .base = SCC_X86_REG_RSP,
                        .disp = {.displacement = -slot->offset},
                        .index = SCC_X86_REG_INVALID,
                        .scale = 1,
                    });
                }
            }
        }
    }

    int total_size = ctx->offset + 32; // shadow space
    ctx->offset =
        (total_size + WIN64_STACK_ALIGN - 1) & ~(WIN64_STACK_ALIGN - 1);
    func_meta->frame_size = ctx->offset;
}

void scc_win_pc_x64_frame_alloc_init(scc_frame_layout_t *ctx) {
    ctx->impl_fn = frame_alloc_impl;
}

/*
windows x86 prolog epilog
https://learn.microsoft.com/zh-cn/cpp/build/prolog-and-epilog?view=msvc-180
*/
static void prologue(scc_mir_instr_vec_t *userdata,
                     const scc_mir_func_t *func) {
    scc_mir_x86_instr_vec_t *instrs = (void *)userdata;
    scc_mir_x86_instr_t instr;
    /*
    mov    [RSP + 8], RCX
    push   R15
    push   R14
    push   R13
    sub    RSP, fixed-allocation-size
    lea    R13, 128[RSP]
    ...
    */
    scc_mir_func_meta_t *meta = SCC_MIR_FUNC_META(func);
    int frame_size = meta->frame_size;

    // va_arg
    if (meta->need_va_args) {
        /// FILL to shadow space
    }

    // sub rsp, frame_size
    if (frame_size > 0) {
        scc_mir_x86_instr_2(&instr, SCC_X86_IFORM_SUB_GPRV_IMMZ,
                            scc_x86_op_preg(SCC_X86_REG_RSP),
                            scc_x86_op_imm(frame_size), scc_pos_create());
        scc_vec_push(*instrs, instr);
    }
}

static void epilogue(scc_mir_instr_vec_t *userdata,
                     const scc_mir_func_t *func) {
    scc_mir_x86_instr_vec_t *instrs = (void *)userdata;
    scc_mir_x86_instr_t instr;
    /*
    add      RSP, fixed-allocation-size
    pop      R13
    pop      R14
    pop      R15
    ret
    */
    scc_mir_func_meta_t *meta = SCC_MIR_FUNC_META(func);
    int frame_size = meta->frame_size;

    // add rsp, frame_size
    if (frame_size > 0) {
        scc_mir_x86_instr_2(&instr, SCC_X86_IFORM_ADD_GPRV_IMMZ,
                            scc_x86_op_preg(SCC_X86_REG_RSP),
                            scc_x86_op_imm(frame_size), scc_pos_create());
        scc_vec_push(*instrs, instr);
    }
}

static int need_epilog(const scc_mir_instr_t *instr) {
    scc_mir_x86_instr_t *ins = (scc_mir_x86_instr_t *)instr;
    switch (ins->x86_instr.opcode) {
    case SCC_X86_IFORM_RET_FAR:
    case SCC_X86_IFORM_RET_NEAR:
    case SCC_X86_IFORM_RET_FAR_IMMW:
    case SCC_X86_IFORM_RET_NEAR_IMMW:
        return true;
    default:
        break;
    }
    return false;
}

void scc_win_pc_x64_prolog_epilog_init(scc_prolog_epilog_t *ctx) {
    ctx->prolog = prologue;
    ctx->epilog = epilogue;
    ctx->need_epilog = need_epilog;
}

static void lower_call(void *userdata, const scc_lir_instr_t *instr) {
    scc_x86_64_isel_t *isel = userdata;

    for (int i = instr->metadata.call.arg_count - 1; i >= 0; i -= 1) {
        scc_lir_val_t *args = &instr->metadata.call.args[i];
        switch (i) {
        case 0:
            scc_x86_emit_move(isel, scc_x86_op_preg(SCC_X86_REG_RCX),
                              scc_x86_lir_val_to_mir_op(isel, args), 8);
            break;
        case 1:
            scc_x86_emit_move(isel, scc_x86_op_preg(SCC_X86_REG_RDX),
                              scc_x86_lir_val_to_mir_op(isel, args), 8);
            break;
        case 2:
            scc_x86_emit_move(isel, scc_x86_op_preg(SCC_X86_REG_R8),
                              scc_x86_lir_val_to_mir_op(isel, args), 8);
            break;
        case 3:
            scc_x86_emit_move(isel, scc_x86_op_preg(SCC_X86_REG_R9),
                              scc_x86_lir_val_to_mir_op(isel, args), 8);
            break;
        default:
            scc_x86_operand_value_t op = scc_x86_lir_val_to_mir_op(isel, args);
            add_instr_1(isel,
                        op.kind == SCC_X86_OPR_REG ? SCC_X86_IFORM_PUSH_GPRV_50
                                                   : SCC_X86_IFORM_PUSH_IMMZ,
                        op, 8);
            break;
        }
    }

    emit_direct_call(isel, instr->metadata.call.callee);

    scc_x86_operand_value_t ret_reg =
        scc_x86_lir_val_to_mir_op(isel, &instr->to);
    if (ret_reg.kind != SCC_X86_OPR_NONE) {
        scc_x86_emit_move(isel, ret_reg, scc_x86_op_preg(SCC_X86_REG_RAX), 8);
    }
}

static void lower_param(void *userdata, const scc_lir_val_t *val,
                        void *out_op) {
    scc_x86_operand_value_t *out = out_op;
    Assert(val->kind == SCC_LIR_INSTR_KIND_ARG);
    switch (val->data.arg) {
    case 0:
        *out = scc_x86_op_preg(SCC_X86_REG_RCX);
        break;
    case 1:
        *out = scc_x86_op_preg(SCC_X86_REG_RDX);
        break;
    case 2:
        *out = scc_x86_op_preg(SCC_X86_REG_R8);
        break;
    case 3:
        *out = scc_x86_op_preg(SCC_X86_REG_R9);
        break;
    default:
        *out = scc_x86_op_slot(-val->data.arg);
        break;
    }
}

static void lower_ret(void *userdata, const scc_lir_instr_t *instr) {
    scc_x86_64_isel_t *isel = userdata;
    scc_lir_val_t ret_val = instr->metadata.ret_val;
    if (ret_val.kind != SCC_LIR_INSTR_KIND_NONE) {
        scc_x86_emit_move(isel, scc_x86_op_preg(SCC_X86_REG_RAX),
                          scc_x86_lir_val_to_mir_op(isel, &ret_val), 8);
    }
    emit_ret(isel);
}

void scc_win_pc_x64_abi_lowering(scc_abi_lowering_t *abi_lowering) {
    abi_lowering->lower_call = lower_call;
    abi_lowering->lower_ret = lower_ret;
    abi_lowering->lower_param = lower_param;

    abi_lowering->lower_va_start = nullptr;
    abi_lowering->lower_va_arg = nullptr;
    abi_lowering->lower_va_copy = nullptr;
    abi_lowering->lower_va_end = nullptr;
}

void scc_win_pc_x64_reg_alloc_fill(scc_reg_alloc_op_t *ops) {
    (void)ops;
    // 寄存器池已在 scc_reg_alloc_fill_arch_x86 中注册
}