#include "CalleeSavedHandler.h"
#include <set>
#include <algorithm>

namespace sysy {

char CalleeSavedHandler::ID = 0;

bool CalleeSavedHandler::runOnFunction(Function *F, AnalysisManager& AM) {
    // This pass works on MachineFunction level, not IR level
    return false;
}

void CalleeSavedHandler::runOnMachineFunction(MachineFunction* mfunc) {
    StackFrameInfo& frame_info = mfunc->getFrameInfo();
    std::set<PhysicalReg> used_callee_saved;

    // 1. 扫描所有指令，找出被使用的s寄存器
    for (auto& mbb : mfunc->getBlocks()) {
        for (auto& instr : mbb->getInstructions()) {
            for (auto& op : instr->getOperands()) {
                
                // 辅助Lambda，用于检查和插入寄存器
                auto check_and_insert_reg = [&](RegOperand* reg_op) {
                    if (!reg_op->isVirtual()) {
                        PhysicalReg preg = reg_op->getPReg();
                        // --- 关键检查点 ---
                        // 必须严格判断是否在 s0-s11 的范围内。
                        // a0, t0 等寄存器绝对不应被视为被调用者保存寄存器。
                        if (preg >= PhysicalReg::S0 && preg <= PhysicalReg::S11) {
                            used_callee_saved.insert(preg);
                        }
                    }
                };

                if (op->getKind() == MachineOperand::KIND_REG) {
                    check_and_insert_reg(static_cast<RegOperand*>(op.get()));
                } else if (op->getKind() == MachineOperand::KIND_MEM) {
                    check_and_insert_reg(static_cast<MemOperand*>(op.get())->getBase());
                }
            }
        }
    }
    
    // 如果没有使用s寄存器（除了可能作为帧指针的s0），则无需操作
    if (used_callee_saved.empty() || (used_callee_saved.size() == 1 && used_callee_saved.count(PhysicalReg::S0))) {
        return;
    }

    // 将结果存入StackFrameInfo，供后续使用
    frame_info.used_callee_saved_regs = used_callee_saved;

    // 2. 在函数序言插入保存指令
    MachineBasicBlock* entry_block = mfunc->getBlocks().front().get();
    auto& entry_instrs = entry_block->getInstructions();
    auto prologue_end = entry_instrs.begin();
    
    // 找到序言结束的位置（通常是addi s0, sp, size之后）
    for (auto it = entry_instrs.begin(); it != entry_instrs.end(); ++it) {
        if ((*it)->getOpcode() == RVOpcodes::ADDI && 
            (*it)->getOperands()[0]->getKind() == MachineOperand::KIND_REG &&
            static_cast<RegOperand*>((*it)->getOperands()[0].get())->getPReg() == PhysicalReg::S0) 
        {
            prologue_end = std::next(it);
            break;
        }
    }
    
    // 为了栈帧布局确定性，对寄存器进行排序
    std::vector<PhysicalReg> sorted_regs(used_callee_saved.begin(), used_callee_saved.end());
    std::sort(sorted_regs.begin(), sorted_regs.end());
    
    int current_offset = -16; // ra和s0已经占用了-8和-16的位置
    for (PhysicalReg reg : sorted_regs) {
        if (reg == PhysicalReg::S0) continue; // s0已经在序言中处理
        current_offset -= 8;
        auto sd = std::make_unique<MachineInstr>(RVOpcodes::SD);
        sd->addOperand(std::make_unique<RegOperand>(reg));
        sd->addOperand(std::make_unique<MemOperand>(
            std::make_unique<RegOperand>(PhysicalReg::S0), // 假设s0是帧指针
            std::make_unique<ImmOperand>(current_offset)
        ));
        entry_instrs.insert(prologue_end, std::move(sd));
    }

    // 3. 在函数结尾（ret之前）插入恢复指令
    for (auto& mbb : mfunc->getBlocks()) {
        for (auto it = mbb->getInstructions().begin(); it != mbb->getInstructions().end(); ++it) {
            if ((*it)->getOpcode() == RVOpcodes::RET) {
                // 以相反的顺序恢复
                current_offset = -16;
                for (PhysicalReg reg : sorted_regs) {
                    if (reg == PhysicalReg::S0) continue;
                    current_offset -= 8;
                    auto ld = std::make_unique<MachineInstr>(RVOpcodes::LD);
                    ld->addOperand(std::make_unique<RegOperand>(reg));
                    ld->addOperand(std::make_unique<MemOperand>(
                        std::make_unique<RegOperand>(PhysicalReg::S0),
                        std::make_unique<ImmOperand>(current_offset)
                    ));
                    mbb->getInstructions().insert(it, std::move(ld));
                }
                break; // 处理完一个基本块的ret即可
            }
        }
    }
}

} // namespace sysy