#include "mir/MIR.h"

#include <stdexcept>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include <cstring>

#include "ir/IR.h"
#include "utils/Log.h"
#include <iostream>

namespace mir {
namespace {

using ValueSlotMap = std::unordered_map<const ir::Value*, int>;

uint32_t GetTypeSize(const ir::Type* type) {
  if (type->IsInt32() || type->IsFloat()) {
    return 4;
  }
  if (type->IsPtrInt32() || type->IsPtrFloat()) {
    return 8; // 64-bit pointers
  }
  if (type->IsArray()) {
    auto* arr_ty = const_cast<ir::Type*>(type)->GetAsArrayType().get();
    return arr_ty->GetNumElements() * GetTypeSize(arr_ty->GetElementType().get());
  }
  return 4;
}

std::unordered_set<const ir::Value*> IdentifyPointerValues(const ir::Function& function) {
  std::unordered_set<const ir::Value*> pointers;

  // 1. Arguments that are pointers
  for (const auto& arg : function.GetArguments()) {
    if (arg->GetType()->IsPtrInt32() || arg->GetType()->IsPtrFloat()) {
      pointers.insert(arg.get());
    }
  }

  // 2. Alloca instructions that store a pointer argument
  for (const auto& bbPtr : function.GetBlocks()) {
    for (const auto& instPtr : bbPtr->GetInstructions()) {
      const auto* inst = instPtr.get();
      if (inst->GetOpcode() == ir::Opcode::Alloca) {
        bool stores_ptr = false;
        auto* parent_bb = inst->GetParent();
        if (parent_bb) {
          auto* parent_func = parent_bb->GetParent();
          if (parent_func) {
            for (const auto& other_bb : parent_func->GetBlocks()) {
              for (const auto& other_inst : other_bb->GetInstructions()) {
                if (other_inst->GetOpcode() == ir::Opcode::Store) {
                  auto* store = static_cast<const ir::StoreInst*>(other_inst.get());
                  if (store->GetPtr() == inst) {
                    auto* val = store->GetValue();
                    if (val->GetType()->IsPtrInt32() || val->GetType()->IsPtrFloat() || pointers.find(val) != pointers.end()) {
                      stores_ptr = true;
                      break;
                    }
                  }
                }
              }
              if (stores_ptr) break;
            }
          }
        }
        if (stores_ptr) {
          pointers.insert(inst);
        }
      }
    }
  }

  // 3. GEP instructions
  for (const auto& bbPtr : function.GetBlocks()) {
    for (const auto& instPtr : bbPtr->GetInstructions()) {
      const auto* inst = instPtr.get();
      if (inst->GetOpcode() == ir::Opcode::GEP) {
        pointers.insert(inst);
      }
    }
  }

  // 4. Load instructions that load from those pointer-storing allocas
  for (const auto& bbPtr : function.GetBlocks()) {
    for (const auto& instPtr : bbPtr->GetInstructions()) {
      const auto* inst = instPtr.get();
      if (inst->GetOpcode() == ir::Opcode::Load) {
        auto* load = static_cast<const ir::LoadInst*>(inst);
        if (pointers.find(load->GetPtr()) != pointers.end()) {
          if (auto* alloca = dynamic_cast<const ir::Instruction*>(load->GetPtr())) {
            if (alloca->GetOpcode() == ir::Opcode::Alloca) {
              pointers.insert(inst);
            }
          }
        }
      }
    }
  }

  return pointers;
}

uint32_t GetAllocaSize(const ir::Instruction& inst, const std::unordered_set<const ir::Value*>& pointers) {
  auto type = inst.GetType();
  if (pointers.find(&inst) != pointers.end()) {
    return 8; // Stores a 64-bit pointer
  }
  return GetTypeSize(type.get());
}

bool IsPowerOfTwo(uint32_t value) {
  return value != 0 && (value & (value - 1)) == 0;
}

int Log2(uint32_t value) {
  int shift = 0;
  while (value > 1) {
    value >>= 1;
    shift++;
  }
  return shift;
}

std::vector<uint32_t> GetGepStrides(const ir::GetElementPtrInst& gep) {
  std::vector<uint32_t> strides;
  auto curr_type = gep.GetPtr()->GetType();
  if (curr_type->IsPtrInt32() || curr_type->IsPtrFloat()) {
    strides.push_back(4);
  } else if (curr_type->IsArray()) {
    strides.push_back(GetTypeSize(curr_type.get()));
    for (size_t i = 2; i < gep.GetNumOperands(); ++i) {
      curr_type = curr_type->GetAsArrayType()->GetElementType();
      strides.push_back(GetTypeSize(curr_type.get()));
    }
  }
  return strides;
}

void EmitAddressToReg(const ir::Value* value, PhysReg target,
                      const ValueSlotMap& slots, MachineBasicBlock& block) {
  if (auto* alloca = dynamic_cast<const ir::Instruction*>(value)) {
    if (alloca->GetOpcode() == ir::Opcode::Alloca) {
      auto it = slots.find(value);
      if (it == slots.end()) {
        throw std::runtime_error(FormatError("mir", "找不到局部变量的栈槽: " + value->GetName()));
      }
      block.Append(Opcode::AddRegImm, {Operand::Reg(target), Operand::Reg(PhysReg::X29), Operand::FrameIndex(it->second)});
      return;
    }
  }

  if (value->IsGlobalValue()) {
    block.Append(Opcode::Adrp, {Operand::Reg(target), Operand::Global(value->GetName())});
    block.Append(Opcode::AddRegImm, {Operand::Reg(target), Operand::Reg(target), Operand::Global(value->GetName())});
    return;
  }

  // Otherwise, the address itself is stored in a stack slot
  auto it = slots.find(value);
  if (it == slots.end()) {
    throw std::runtime_error(FormatError("mir", "找不到指针的值槽: " + value->GetName()));
  }
  block.Append(Opcode::LoadStack, {Operand::Reg(target), Operand::FrameIndex(it->second)});
}

void EmitValueToReg(const ir::Value* value, PhysReg target,
                    const ValueSlotMap& slots, MachineBasicBlock& block) {
  if (auto* constant = dynamic_cast<const ir::ConstantInt*>(value)) {
    block.Append(Opcode::MovImm, {Operand::Reg(target), Operand::Imm(constant->GetValue())});
    return;
  }

  if (auto* constant = dynamic_cast<const ir::ConstantFloat*>(value)) {
    float fval = constant->GetValue();
    int bits;
    std::memcpy(&bits, &fval, sizeof(float));
    block.Append(Opcode::MovImm, {Operand::Reg(target), Operand::Imm(bits)});
    return;
  }

  if (value->IsGlobalValue()) {
    EmitAddressToReg(value, target, slots, block);
    return;
  }

  auto it = slots.find(value);
  if (it == slots.end()) {
    throw std::runtime_error(FormatError("mir", "找不到值对应的栈槽: " + value->GetName()));
  }

  block.Append(Opcode::LoadStack, {Operand::Reg(target), Operand::FrameIndex(it->second)});
}

void LowerInstruction(const ir::Instruction& inst, MachineFunction& function,
                      ValueSlotMap& slots, MachineBasicBlock& block,
                      const std::unordered_set<const ir::Value*>& pointers) {
  switch (inst.GetOpcode()) {
    case ir::Opcode::Alloca: {
      slots.emplace(&inst, function.CreateFrameIndex(GetAllocaSize(inst, pointers)));
      return;
    }
    case ir::Opcode::Store: {
      auto& store = static_cast<const ir::StoreInst&>(inst);
      
      if (auto* alloca = dynamic_cast<const ir::Instruction*>(store.GetPtr())) {
        if (alloca->GetOpcode() == ir::Opcode::Alloca) {
          auto it = slots.find(alloca);
          if (it != slots.end()) {
            bool is_ptr = store.GetValue()->GetType()->IsPtrInt32() || 
                          store.GetValue()->GetType()->IsPtrFloat() || 
                          pointers.find(store.GetValue()) != pointers.end() ||
                          store.GetValue()->IsGlobalValue();
            PhysReg val_reg = store.GetValue()->GetType()->IsFloat() ? PhysReg::S8 :
                              is_ptr ? PhysReg::X8 : PhysReg::W8;
            EmitValueToReg(store.GetValue(), val_reg, slots, block);
            block.Append(Opcode::StoreStack, {Operand::Reg(val_reg), Operand::FrameIndex(it->second)});
            return;
          }
        }
      }

      // Dynamic store
      bool is_ptr = store.GetValue()->GetType()->IsPtrInt32() || 
                    store.GetValue()->GetType()->IsPtrFloat() || 
                    pointers.find(store.GetValue()) != pointers.end() ||
                    store.GetValue()->IsGlobalValue();
      PhysReg val_reg = store.GetValue()->GetType()->IsFloat() ? PhysReg::S8 :
                        is_ptr ? PhysReg::X8 : PhysReg::W8;
      EmitValueToReg(store.GetValue(), val_reg, slots, block);
      EmitAddressToReg(store.GetPtr(), PhysReg::X9, slots, block);
      block.Append(Opcode::StrRegReg, {Operand::Reg(val_reg), Operand::Reg(PhysReg::X9)});
      return;
    }
    case ir::Opcode::Load: {
      auto& load = static_cast<const ir::LoadInst&>(inst);
      bool is_ptr = load.GetType()->IsPtrInt32() || 
                    load.GetType()->IsPtrFloat() || 
                    pointers.find(&inst) != pointers.end();
      int dst_slot = function.CreateFrameIndex(is_ptr ? 8 : GetTypeSize(load.GetType().get()));
      slots.emplace(&inst, dst_slot);

      if (auto* alloca = dynamic_cast<const ir::Instruction*>(load.GetPtr())) {
        if (alloca->GetOpcode() == ir::Opcode::Alloca) {
          auto it = slots.find(alloca);
          if (it != slots.end()) {
            PhysReg val_reg = load.GetType()->IsFloat() ? PhysReg::S8 :
                              is_ptr ? PhysReg::X8 : PhysReg::W8;
            block.Append(Opcode::LoadStack, {Operand::Reg(val_reg), Operand::FrameIndex(it->second)});
            block.Append(Opcode::StoreStack, {Operand::Reg(val_reg), Operand::FrameIndex(dst_slot)});
            return;
          }
        }
      }

      // Dynamic load
      PhysReg val_reg = load.GetType()->IsFloat() ? PhysReg::S8 :
                        is_ptr ? PhysReg::X8 : PhysReg::W8;
      EmitAddressToReg(load.GetPtr(), PhysReg::X9, slots, block);
      block.Append(Opcode::LdrRegReg, {Operand::Reg(val_reg), Operand::Reg(PhysReg::X9)});
      block.Append(Opcode::StoreStack, {Operand::Reg(val_reg), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::Add:
    case ir::Opcode::Sub:
    case ir::Opcode::Mul:
    case ir::Opcode::Div:
    case ir::Opcode::Mod: {
      auto& bin = static_cast<const ir::BinaryInst&>(inst);
      int dst_slot = function.CreateFrameIndex(4);
      slots.emplace(&inst, dst_slot);

      if (inst.GetOpcode() == ir::Opcode::Add) {
        const ir::Value* variable = nullptr;
        int constant = 0;
        if (auto* lhs_const = dynamic_cast<const ir::ConstantInt*>(bin.GetLhs())) {
          variable = bin.GetRhs();
          constant = lhs_const->GetValue();
        } else if (auto* rhs_const = dynamic_cast<const ir::ConstantInt*>(bin.GetRhs())) {
          variable = bin.GetLhs();
          constant = rhs_const->GetValue();
        }
        if (variable && constant >= -4095 && constant <= 4095) {
          EmitValueToReg(variable, PhysReg::W8, slots, block);
          block.Append(Opcode::AddRegImm, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8), Operand::Imm(constant)});
          block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
          return;
        }
      } else if (inst.GetOpcode() == ir::Opcode::Sub) {
        if (auto* rhs_const = dynamic_cast<const ir::ConstantInt*>(bin.GetRhs())) {
          int constant = rhs_const->GetValue();
          if (constant >= -4095 && constant <= 4095) {
            EmitValueToReg(bin.GetLhs(), PhysReg::W8, slots, block);
            block.Append(Opcode::AddRegImm, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8), Operand::Imm(-constant)});
            block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
            return;
          }
        }
      }

      if (inst.GetOpcode() == ir::Opcode::Mul) {
        const ir::Value* variable = nullptr;
        int constant = 0;
        if (auto* lhs_const = dynamic_cast<const ir::ConstantInt*>(bin.GetLhs())) {
          variable = bin.GetRhs();
          constant = lhs_const->GetValue();
        } else if (auto* rhs_const = dynamic_cast<const ir::ConstantInt*>(bin.GetRhs())) {
          variable = bin.GetLhs();
          constant = rhs_const->GetValue();
        }

        if (variable && constant > 1) {
          EmitValueToReg(variable, PhysReg::W8, slots, block);
          if (IsPowerOfTwo(static_cast<uint32_t>(constant))) {
            block.Append(Opcode::LslImm, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8), Operand::Imm(Log2(static_cast<uint32_t>(constant)))});
            block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
            return;
          }
          if (IsPowerOfTwo(static_cast<uint32_t>(constant - 1))) {
            block.Append(Opcode::LslImm, {Operand::Reg(PhysReg::W9), Operand::Reg(PhysReg::W8), Operand::Imm(Log2(static_cast<uint32_t>(constant - 1)))});
            block.Append(Opcode::AddRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9), Operand::Reg(PhysReg::W8)});
            block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
            return;
          }
        }
      }

      EmitValueToReg(bin.GetLhs(), PhysReg::W8, slots, block);
      EmitValueToReg(bin.GetRhs(), PhysReg::W9, slots, block);

      if (inst.GetOpcode() == ir::Opcode::Add) {
        block.Append(Opcode::AddRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9)});
      } else if (inst.GetOpcode() == ir::Opcode::Sub) {
        block.Append(Opcode::SubRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9)});
      } else if (inst.GetOpcode() == ir::Opcode::Mul) {
        block.Append(Opcode::MulRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9)});
      } else if (inst.GetOpcode() == ir::Opcode::Div) {
        block.Append(Opcode::SDivRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9)});
      } else if (inst.GetOpcode() == ir::Opcode::Mod) {
        block.Append(Opcode::SDivRR, {Operand::Reg(PhysReg::W10), Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9)});
        block.Append(Opcode::MSubRRRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W10), Operand::Reg(PhysReg::W9), Operand::Reg(PhysReg::W8)});
      }

      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::FAdd:
    case ir::Opcode::FSub:
    case ir::Opcode::FMul:
    case ir::Opcode::FDiv: {
      auto& bin = static_cast<const ir::BinaryInst&>(inst);
      int dst_slot = function.CreateFrameIndex(4);
      slots.emplace(&inst, dst_slot);

      EmitValueToReg(bin.GetLhs(), PhysReg::S8, slots, block);
      EmitValueToReg(bin.GetRhs(), PhysReg::S9, slots, block);

      if (inst.GetOpcode() == ir::Opcode::FAdd) {
        block.Append(Opcode::FAddRRR, {Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S9)});
      } else if (inst.GetOpcode() == ir::Opcode::FSub) {
        block.Append(Opcode::FSubRRR, {Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S9)});
      } else if (inst.GetOpcode() == ir::Opcode::FMul) {
        block.Append(Opcode::FMulRRR, {Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S9)});
      } else if (inst.GetOpcode() == ir::Opcode::FDiv) {
        block.Append(Opcode::FDivRRR, {Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S9)});
      }

      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::S8), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::ICmpEQ:
    case ir::Opcode::ICmpNE:
    case ir::Opcode::ICmpLT:
    case ir::Opcode::ICmpGT:
    case ir::Opcode::ICmpLE:
    case ir::Opcode::ICmpGE: {
      auto& cmp = static_cast<const ir::BinaryInst&>(inst);
      int dst_slot = function.CreateFrameIndex(4);
      slots.emplace(&inst, dst_slot);

      EmitValueToReg(cmp.GetLhs(), PhysReg::W8, slots, block);
      EmitValueToReg(cmp.GetRhs(), PhysReg::W9, slots, block);
      block.Append(Opcode::CmpRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9)});

      std::string cond;
      switch (inst.GetOpcode()) {
        case ir::Opcode::ICmpEQ: cond = "eq"; break;
        case ir::Opcode::ICmpNE: cond = "ne"; break;
        case ir::Opcode::ICmpLT: cond = "lt"; break;
        case ir::Opcode::ICmpGT: cond = "gt"; break;
        case ir::Opcode::ICmpLE: cond = "le"; break;
        case ir::Opcode::ICmpGE: cond = "ge"; break;
        default: break;
      }

      block.Append(Opcode::Cset, {Operand::Reg(PhysReg::W8), Operand::Cond(cond)});
      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::FCmpEQ:
    case ir::Opcode::FCmpNE:
    case ir::Opcode::FCmpLT:
    case ir::Opcode::FCmpGT:
    case ir::Opcode::FCmpLE:
    case ir::Opcode::FCmpGE: {
      auto& cmp = static_cast<const ir::BinaryInst&>(inst);
      int dst_slot = function.CreateFrameIndex(4);
      slots.emplace(&inst, dst_slot);

      EmitValueToReg(cmp.GetLhs(), PhysReg::S8, slots, block);
      EmitValueToReg(cmp.GetRhs(), PhysReg::S9, slots, block);
      block.Append(Opcode::FCmpRR, {Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::S9)});

      std::string cond;
      switch (inst.GetOpcode()) {
        case ir::Opcode::FCmpEQ: cond = "eq"; break;
        case ir::Opcode::FCmpNE: cond = "ne"; break;
        case ir::Opcode::FCmpLT: cond = "mi"; break;
        case ir::Opcode::FCmpGT: cond = "gt"; break;
        case ir::Opcode::FCmpLE: cond = "ls"; break;
        case ir::Opcode::FCmpGE: cond = "ge"; break;
        default: break;
      }

      block.Append(Opcode::Cset, {Operand::Reg(PhysReg::W8), Operand::Cond(cond)});
      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::ZExt: {
      auto& cast = static_cast<const ir::CastInst&>(inst);
      int dst_slot = function.CreateFrameIndex(4);
      slots.emplace(&inst, dst_slot);

      EmitValueToReg(cast.GetValue(), PhysReg::W8, slots, block);
      block.Append(Opcode::ZExt, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W8)});
      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::SIToFP: {
      auto& cast = static_cast<const ir::CastInst&>(inst);
      int dst_slot = function.CreateFrameIndex(4);
      slots.emplace(&inst, dst_slot);

      EmitValueToReg(cast.GetValue(), PhysReg::W8, slots, block);
      block.Append(Opcode::SIToFP, {Operand::Reg(PhysReg::S8), Operand::Reg(PhysReg::W8)});
      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::S8), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::FPToSI: {
      auto& cast = static_cast<const ir::CastInst&>(inst);
      int dst_slot = function.CreateFrameIndex(4);
      slots.emplace(&inst, dst_slot);

      EmitValueToReg(cast.GetValue(), PhysReg::S8, slots, block);
      block.Append(Opcode::FPToSI, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::S8)});
      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
      return;
    }
    case ir::Opcode::Br: {
      auto& br = static_cast<const ir::BranchInst&>(inst);
      
      auto emit_phi_copies = [&](const ir::BasicBlock* succ) {
        if (!succ) return;
        for (const auto& succ_inst : succ->GetInstructions()) {
          if (succ_inst->GetOpcode() == ir::Opcode::Phi) {
            auto* phi = static_cast<const ir::PhiInst*>(succ_inst.get());
            const ir::Value* incoming_val = nullptr;
            for (size_t i = 0; i < phi->GetNumIncoming(); ++i) {
              if (phi->GetIncomingBlock(i) == inst.GetParent()) {
                incoming_val = phi->GetIncomingValue(i);
                break;
              }
            }
            if (incoming_val) {
              auto slot_it = slots.find(phi);
              if (slot_it != slots.end()) {
                int phi_slot = slot_it->second;
                bool is_ptr = phi->GetType()->IsPtrInt32() || 
                              phi->GetType()->IsPtrFloat() || 
                              pointers.find(phi) != pointers.end() ||
                              (incoming_val && (pointers.find(incoming_val) != pointers.end() || incoming_val->IsGlobalValue()));
                PhysReg val_reg = phi->GetType()->IsFloat() ? PhysReg::S8 :
                                  is_ptr ? PhysReg::X8 : PhysReg::W8;
                EmitValueToReg(incoming_val, val_reg, slots, block);
                block.Append(Opcode::StoreStack, {Operand::Reg(val_reg), Operand::FrameIndex(phi_slot)});
              }
            }
          }
        }
      };

      if (br.IsConditional()) {
        emit_phi_copies(br.GetIfTrue());
        emit_phi_copies(br.GetIfFalse());
        EmitValueToReg(br.GetCondition(), PhysReg::W8, slots, block);
        block.Append(Opcode::MovImm, {Operand::Reg(PhysReg::W9), Operand::Imm(0)});
        block.Append(Opcode::CmpRR, {Operand::Reg(PhysReg::W8), Operand::Reg(PhysReg::W9)});
        block.Append(Opcode::BCond, {Operand::Cond("ne"), Operand::Label(br.GetIfTrue()->GetName())});
        block.Append(Opcode::B, {Operand::Label(br.GetIfFalse()->GetName())});
      } else {
        emit_phi_copies(br.GetDest());
        block.Append(Opcode::B, {Operand::Label(br.GetDest()->GetName())});
      }
      return;
    }
    case ir::Opcode::Phi: {
      return;
    }
    case ir::Opcode::Ret: {
      auto& ret = static_cast<const ir::ReturnInst&>(inst);
      if (ret.GetValue()) {
        bool is_ptr = ret.GetValue()->GetType()->IsPtrInt32() || 
                      ret.GetValue()->GetType()->IsPtrFloat() || 
                      pointers.find(ret.GetValue()) != pointers.end() ||
                      ret.GetValue()->IsGlobalValue();
        PhysReg ret_reg = ret.GetValue()->GetType()->IsFloat() ? PhysReg::S0 :
                          is_ptr ? PhysReg::X0 : PhysReg::W0;
        EmitValueToReg(ret.GetValue(), ret_reg, slots, block);
      }
      block.Append(Opcode::Ret);
      return;
    }
    case ir::Opcode::Call: {
      auto& call = static_cast<const ir::CallInst&>(inst);
      bool is_ret_ptr = call.GetType()->IsPtrInt32() || 
                        call.GetType()->IsPtrFloat() || 
                        pointers.find(&inst) != pointers.end();
      int dst_slot = -1;
      if (!call.GetType()->IsVoid()) {
        dst_slot = function.CreateFrameIndex(is_ret_ptr ? 8 : GetTypeSize(call.GetType().get()));
        slots.emplace(&inst, dst_slot);
      }

      int int_idx = 0;
      int float_idx = 0;
      for (size_t i = 1; i < call.GetNumOperands(); ++i) {
        auto* arg = call.GetOperand(i);
        if (arg->GetType()->IsFloat()) {
          PhysReg reg = static_cast<PhysReg>(static_cast<int>(PhysReg::S0) + float_idx);
          EmitValueToReg(arg, reg, slots, block);
          float_idx++;
        } else {
          bool is_arg_ptr = arg->GetType()->IsPtrInt32() || 
                            arg->GetType()->IsPtrFloat() || 
                            pointers.find(arg) != pointers.end() ||
                            arg->IsGlobalValue();
          PhysReg reg = is_arg_ptr
                            ? static_cast<PhysReg>(static_cast<int>(PhysReg::X0) + int_idx)
                            : static_cast<PhysReg>(static_cast<int>(PhysReg::W0) + int_idx);
          EmitValueToReg(arg, reg, slots, block);
          int_idx++;
        }
      }

      block.Append(Opcode::Call, {Operand::Global(call.GetFunction()->GetName())});

      if (dst_slot != -1) {
        if (call.GetType()->IsFloat()) {
          block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::S0), Operand::FrameIndex(dst_slot)});
        } else {
          PhysReg ret_reg = is_ret_ptr ? PhysReg::X0 : PhysReg::W0;
          block.Append(Opcode::StoreStack, {Operand::Reg(ret_reg), Operand::FrameIndex(dst_slot)});
        }
      }
      return;
    }
    case ir::Opcode::GEP: {
      auto& gep = static_cast<const ir::GetElementPtrInst&>(inst);
      int dst_slot = function.CreateFrameIndex(8);
      slots.emplace(&inst, dst_slot);

      // Load base pointer address into X8
      if (gep.GetPtr()->IsGlobalValue()) {
        EmitAddressToReg(gep.GetPtr(), PhysReg::X8, slots, block);
      } else if (auto* alloca = dynamic_cast<const ir::AllocaInst*>(gep.GetPtr())) {
        if (alloca->GetType()->IsArray()) {
          EmitAddressToReg(gep.GetPtr(), PhysReg::X8, slots, block);
        } else {
          EmitValueToReg(gep.GetPtr(), PhysReg::X8, slots, block);
        }
      } else {
        EmitValueToReg(gep.GetPtr(), PhysReg::X8, slots, block);
      }

      auto strides = GetGepStrides(gep);
      for (size_t i = 1; i < gep.GetNumOperands(); ++i) {
        auto* idx = gep.GetOperand(i);
        uint32_t stride = strides.at(i - 1);
        
        if (auto* ci = dynamic_cast<const ir::ConstantInt*>(idx)) {
          int64_t offset = static_cast<int64_t>(ci->GetValue()) * stride;
          if (offset != 0) {
            block.Append(Opcode::AddRegImm, {Operand::Reg(PhysReg::X8), Operand::Reg(PhysReg::X8), Operand::Imm(offset)});
          }
          continue;
        }

        EmitValueToReg(idx, PhysReg::W9, slots, block);
        bool shifted = false;
        if (stride > 1 && IsPowerOfTwo(stride)) {
          block.Append(Opcode::ZExt, {Operand::Reg(PhysReg::X9), Operand::Reg(PhysReg::W9)});
          block.Append(Opcode::LslImm, {Operand::Reg(PhysReg::X9), Operand::Reg(PhysReg::X9), Operand::Imm(Log2(stride))});
          shifted = true;
        } else if (stride > 1) {
          block.Append(Opcode::MovImm, {Operand::Reg(PhysReg::W10), Operand::Imm(stride)});
          block.Append(Opcode::MulRR, {Operand::Reg(PhysReg::W9), Operand::Reg(PhysReg::W9), Operand::Reg(PhysReg::W10)});
        }

        // Extend W9 to X9 and add to base address X8
        if (!shifted) {
          block.Append(Opcode::ZExt, {Operand::Reg(PhysReg::X9), Operand::Reg(PhysReg::W9)});
        }
        block.Append(Opcode::AddRR, {Operand::Reg(PhysReg::X8), Operand::Reg(PhysReg::X8), Operand::Reg(PhysReg::X9)});
      }

      // Store address into GEP's stack slot
      block.Append(Opcode::StoreStack, {Operand::Reg(PhysReg::X8), Operand::FrameIndex(dst_slot)});
      return;
    }
  }

  throw std::runtime_error(FormatError("mir", "暂不支持该 IR 指令: " + std::to_string(static_cast<int>(inst.GetOpcode()))));
}

}  // namespace

std::vector<std::unique_ptr<MachineFunction>> LowerToMIR(const ir::Module& module) {
  DefaultContext();
  std::vector<std::unique_ptr<MachineFunction>> mfuncs;

  for (const auto& funcPtr : module.GetFunctions()) {
    const auto& func = *funcPtr;
    if (func.GetBlocks().empty()) continue; // skip declarations

    auto machine_func = std::make_unique<MachineFunction>(func.GetName());
    ValueSlotMap slots;
    auto pointers = IdentifyPointerValues(func);

    // First, create all basic blocks in MachineFunction
    std::unordered_map<const ir::BasicBlock*, MachineBasicBlock*> bb_map;
    machine_func->GetBlocks().reserve(func.GetBlocks().size());
    for (const auto& bbPtr : func.GetBlocks()) {
      auto& mbb = machine_func->CreateBlock(bbPtr->GetName());
      bb_map[bbPtr.get()] = &mbb;
    }

    // Pre-allocate stack slots for all Phi instructions in the function
    for (const auto& bbPtr : func.GetBlocks()) {
      for (const auto& inst : bbPtr->GetInstructions()) {
        if (inst->GetOpcode() == ir::Opcode::Phi) {
          bool is_phi_ptr = inst->GetType()->IsPtrInt32() || 
                            inst->GetType()->IsPtrFloat() || 
                            pointers.find(inst.get()) != pointers.end();
          int slot = machine_func->CreateFrameIndex(is_phi_ptr ? 8 : GetTypeSize(inst->GetType().get()));
          slots.emplace(inst.get(), slot);
        }
      }
    }

    auto& entry_block = *bb_map.at(func.GetEntry());

    // Lower function arguments at the start of the entry block
    const auto& args = func.GetArguments();
    int int_idx = 0;
    int float_idx = 0;
    for (const auto& arg : args) {
      bool is_arg_ptr = arg->GetType()->IsPtrInt32() || 
                        arg->GetType()->IsPtrFloat() || 
                        pointers.find(arg.get()) != pointers.end();
      int slot = machine_func->CreateFrameIndex(is_arg_ptr ? 8 : GetTypeSize(arg->GetType().get()));
      slots.emplace(arg.get(), slot);

      if (arg->GetType()->IsFloat()) {
        PhysReg reg = static_cast<PhysReg>(static_cast<int>(PhysReg::S0) + float_idx);
        entry_block.Append(Opcode::StoreStack, {Operand::Reg(reg), Operand::FrameIndex(slot)});
        float_idx++;
      } else {
        PhysReg reg = is_arg_ptr
                          ? static_cast<PhysReg>(static_cast<int>(PhysReg::X0) + int_idx)
                          : static_cast<PhysReg>(static_cast<int>(PhysReg::W0) + int_idx);
        entry_block.Append(Opcode::StoreStack, {Operand::Reg(reg), Operand::FrameIndex(slot)});
        int_idx++;
      }
    }

    // Now, lower all instructions block by block
    for (const auto& bbPtr : func.GetBlocks()) {
      auto& mbb = *bb_map.at(bbPtr.get());
      for (const auto& inst : bbPtr->GetInstructions()) {
        LowerInstruction(*inst, *machine_func, slots, mbb, pointers);
      }
    }

    mfuncs.push_back(std::move(machine_func));
  }

  return mfuncs;
}

}  // namespace mir
