#include "mir/MIR.h"
#include <algorithm>
#include <stdexcept>
#include <vector>

#include "utils/Log.h"

namespace mir {
namespace {

int AlignTo(int value, int align) {
  return ((value + align - 1) / align) * align;
}

}  // namespace

void RunFrameLowering(MachineFunction& function) {
  auto& saved = const_cast<std::vector<PhysReg>&>(function.GetCalleeSavedRegs());
  auto has_reg = [&](PhysReg r) {
      return std::find(saved.begin(), saved.end(), r) != saved.end();
  };
  if (!has_reg(PhysReg::RA)) saved.push_back(PhysReg::RA);
  if (!has_reg(PhysReg::S0)) saved.push_back(PhysReg::S0);
  int cursor = 0;
  std::vector<FrameSlot> newSlots;

  // 1. 为每个现有栈槽分配偏移（包括局部变量、溢出槽）
  for (const auto& slot : function.GetFrameSlots()) {
    int align = slot.size;
    cursor = AlignTo(cursor, align);
    FrameSlot newSlot = slot;
    newSlot.offset = cursor;
    newSlots.push_back(newSlot);
    cursor += slot.size;
  }

  // 2. 为 callee-saved 寄存器分配空间（先计算需要保存的寄存器）
  const auto& savedRegs = function.GetCalleeSavedRegs();
  int savedSize = 0;
  for (PhysReg reg : savedRegs) {
    bool isFloat = (reg >= PhysReg::FT0 && reg <= PhysReg::FT11) ||
                   (reg >= PhysReg::FA0 && reg <= PhysReg::FA7) ||
                   (reg == PhysReg::FS0 || reg == PhysReg::FS1);
    savedSize += isFloat ? 4 : 8;
  }
  
  // 3. 总帧大小 = 局部变量区 + callee-saved 区 + 对齐
  int calleeSavedArea = AlignTo(savedSize, 8);
  cursor = AlignTo(cursor, 16);       // 局部变量区对齐
  int localVarsSize = cursor;
  function.SetLocalVarsSize(localVarsSize);
  
  int totalFrame = localVarsSize + calleeSavedArea;
  function.SetFrameSize(totalFrame);

  // 4. 记录每个 callee-saved 寄存器的偏移（相对于 sp）
  int offset = localVarsSize;
  for (PhysReg reg : savedRegs) {
    bool isFloat = (reg >= PhysReg::FT0 && reg <= PhysReg::FT11) ||
                   (reg >= PhysReg::FA0 && reg <= PhysReg::FA7) ||
                   (reg == PhysReg::FS0 || reg == PhysReg::FS1);
    int size = isFloat ? 4 : 8;
    function.SetCalleeSavedOffset(reg, offset);
    offset += size;
  }
  const_cast<std::vector<FrameSlot>&>(function.GetFrameSlots()).clear();
  for (auto& slot : newSlots) {
    const_cast<std::vector<FrameSlot>&>(function.GetFrameSlots()).push_back(slot);
  }
}

}  // namespace mir