#include <arch/scc_x86_mir.h>
#include <arch/scc_x86_reg_alloc.h>
#include <target/scc_win64.h>

// #include <core_pass/scc_frame_layout.h>
// #include <core_pass/scc_prolog_epilog.h>
// #include <core_pass/scc_reg_alloc.h>

#include <scc_mir_module.h>
#include <scc_mir_pass.h>

void scc_frame_layout(scc_frame_layout_t *ctx, scc_mir_module_t *module) {
    Assert(ctx && ctx->impl_fn && module);
    scc_vec_foreach(module->cfg_module.funcs, i) {
        if (i == 0)
            continue;
        ctx->impl_fn(ctx, module, &scc_vec_at(module->cfg_module.funcs, i));
    }
}

void scc_prolog_epilog(scc_prolog_epilog_t *ctx, scc_mir_module_t *module) {
    Assert(ctx && ctx->epilog && ctx->prolog && module);
    scc_vec_foreach(module->cfg_module.funcs, i) {
        if (i == 0)
            continue;
        scc_mir_func_t *func = &scc_vec_at(module->cfg_module.funcs, i);
        scc_vec_foreach(func->bblocks, i) {
            scc_cfg_bblock_id_t bb_id = scc_vec_at(func->bblocks, i);
            scc_cfg_bblock_t *bb =
                scc_cfg_module_unsafe_get_bblock(&module->cfg_module, bb_id);
            Assert(bb != nullptr);
            scc_mir_instr_vec_t *old_instrs = SCC_MIR_BBLOCK_VALUES_PTR(bb);
            scc_mir_instr_vec_t new_instrs;
            scc_vec_init(new_instrs);

            if (i == 0) {
                ctx->prolog(&new_instrs, func);
            }

            scc_vec_foreach(*old_instrs, i) {
                scc_mir_instr_t *ins =
                    scc_vec_sized_at_ptr(*old_instrs, module->instr_size, i);
                if (ins->opcode < 0) {
                    continue;
                }
                if (ctx->need_epilog(ins)) {
                    ctx->epilog(&new_instrs, func);
                }
                scc_vec_sized_push(new_instrs, module->instr_size, ins,
                                   module->instr_size);
            }
            scc_vec_free(*old_instrs);
            *old_instrs = new_instrs;
        }
    }
}

void scc_mir_pass(scc_mir_module_t *mir_module, scc_mir_pass_stage_t stage) {
    scc_reg_alloc_ctx_t reg_alloc_ctx = {.func = nullptr, .ops = {0}};
    scc_reg_alloc_fill_arch_x86(&reg_alloc_ctx.ops);
    scc_win_pc_x64_reg_alloc_fill(&reg_alloc_ctx.ops);

    scc_reg_alloc(&reg_alloc_ctx, mir_module);
    if (stage == SCC_MIR_STAGE_REGALLOC) {
        return;
    }

    void scc_x86_peephole_optimize(scc_mir_module_t * module);
    // scc_x86_peephole_optimize(mir_module);

    scc_frame_layout_t frame_layout_ctx = {0};
    scc_win_pc_x64_frame_alloc_init(&frame_layout_ctx);
    scc_frame_layout(&frame_layout_ctx, mir_module);
    if (stage == SCC_MIR_STAGE_FRAME_LAYOUT) {
        return;
    }

    scc_prolog_epilog_t prolog_epilog_ctx = {0};
    scc_win_pc_x64_prolog_epilog_init(&prolog_epilog_ctx);
    scc_prolog_epilog(&prolog_epilog_ctx, mir_module);
    if (stage == SCC_MIR_STAGE_PROLOGUE_EPILOGUE) {
        return;
    }
}