#include "mir/MIR.h"

#include <stdexcept>
#include <unordered_map>

#include "ir/IR.h"
#include "utils/Log.h"

namespace mir {
namespace {

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

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->value())});
    return;
  }

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

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

void LowerInstruction(const ir::Instruction& inst, MachineFunction& function,
                      ValueSlotMap& slots) {
  auto& block = function.entry();

  switch (inst.opcode()) {
    case ir::Opcode::Alloca: {
      slots.emplace(&inst, function.CreateFrameIndex());
      return;
    }
    case ir::Opcode::Store: {
      auto& store = static_cast<const ir::StoreInst&>(inst);
      auto dst = slots.find(store.ptr());
      if (dst == slots.end()) {
        throw std::runtime_error(
            FormatError("mir", "暂不支持对非栈变量地址进行写入"));
      }
      EmitValueToReg(store.value(), PhysReg::W8, slots, block);
      block.Append(Opcode::StoreStack,
                   {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst->second)});
      return;
    }
    case ir::Opcode::Load: {
      auto& load = static_cast<const ir::LoadInst&>(inst);
      auto src = slots.find(load.ptr());
      if (src == slots.end()) {
        throw std::runtime_error(
            FormatError("mir", "暂不支持对非栈变量地址进行读取"));
      }
      int dst_slot = function.CreateFrameIndex();
      block.Append(Opcode::LoadStack,
                   {Operand::Reg(PhysReg::W8), Operand::FrameIndex(src->second)});
      block.Append(Opcode::StoreStack,
                   {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
      slots.emplace(&inst, dst_slot);
      return;
    }
    case ir::Opcode::Add: {
      auto& bin = static_cast<const ir::BinaryInst&>(inst);
      int dst_slot = function.CreateFrameIndex();
      EmitValueToReg(bin.lhs(), PhysReg::W8, slots, block);
      EmitValueToReg(bin.rhs(), PhysReg::W9, slots, block);
      block.Append(Opcode::AddRR, {Operand::Reg(PhysReg::W8),
                                   Operand::Reg(PhysReg::W8),
                                   Operand::Reg(PhysReg::W9)});
      block.Append(Opcode::StoreStack,
                   {Operand::Reg(PhysReg::W8), Operand::FrameIndex(dst_slot)});
      slots.emplace(&inst, dst_slot);
      return;
    }
    case ir::Opcode::Ret: {
      auto& ret = static_cast<const ir::ReturnInst&>(inst);
      EmitValueToReg(ret.value(), PhysReg::W0, slots, block);
      block.Append(Opcode::Ret);
      return;
    }
    case ir::Opcode::Sub:
    case ir::Opcode::Mul:
      throw std::runtime_error(FormatError("mir", "暂不支持该二元运算"));
  }

  throw std::runtime_error(FormatError("mir", "暂不支持该 IR 指令"));
}

}  // namespace

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

  if (module.functions().size() != 1) {
    throw std::runtime_error(FormatError("mir", "暂不支持多个函数"));
  }

  const auto& func = *module.functions().front();
  if (func.name() != "main") {
    throw std::runtime_error(FormatError("mir", "暂不支持非 main 函数"));
  }

  auto machine_func = std::make_unique<MachineFunction>(func.name());
  ValueSlotMap slots;
  const auto* entry = func.entry();
  if (!entry) {
    throw std::runtime_error(FormatError("mir", "IR 函数缺少入口基本块"));
  }

  for (const auto& inst : entry->instructions()) {
    LowerInstruction(*inst, *machine_func, slots);
  }

  return machine_func;
}

}  // namespace mir