#include "mir/MIR.h"

#include <ostream>
#include <stdexcept>

#include "utils/Log.h"

namespace mir {
namespace {

const FrameSlot& GetFrameSlot(const MachineFunction& function,
                              const Operand& operand) {
  if (operand.GetKind() != Operand::Kind::FrameIndex) {
    throw std::runtime_error(FormatError("mir", "期望 FrameIndex 操作数"));
  }
  return function.GetFrameSlot(operand.GetFrameIndex());
}

void PrintMovImm(std::ostream& os, PhysReg reg, int imm) {
  const char* reg_name = PhysRegName(reg);
  if (imm >= -32768 && imm <= 65535) {
    os << "  mov " << reg_name << ", #" << imm << "\n";
  } else {
    uint32_t uimm = static_cast<uint32_t>(imm);
    os << "  mov " << reg_name << ", #" << (uimm & 0xFFFF) << "\n";
    os << "  movk " << reg_name << ", #" << ((uimm >> 16) & 0xFFFF) << ", lsl #16\n";
  }
}

void PrintStackAccess(std::ostream& os, const char* mnemonic, PhysReg reg,
                      int offset) {
  if (offset >= -256 && offset <= 255) {
    os << "  " << mnemonic << " " << PhysRegName(reg) << ", [x29, #" << offset
       << "]\n";
  } else {
    // Offset out of range for ldur/stur
    if (offset < 0) {
      PrintMovImm(os, PhysReg::X16, -offset);
      os << "  sub x16, x29, x16\n";
    } else {
      PrintMovImm(os, PhysReg::X16, offset);
      os << "  add x16, x29, x16\n";
    }
    
    if (mnemonic[0] == 'l') { // load
      os << "  ldr " << PhysRegName(reg) << ", [x16]\n";
    } else { // store
      os << "  str " << PhysRegName(reg) << ", [x16]\n";
    }
  }
}

const char* CondCodeName(CondCode cc) {
  switch (cc) {
    case CondCode::EQ: return "eq";
    case CondCode::NE: return "ne";
    case CondCode::LT: return "lt";
    case CondCode::LE: return "le";
    case CondCode::GT: return "gt";
    case CondCode::GE: return "ge";
  }
  return "??";
}

}  // namespace

void PrintAsm(const MachineModule& module, std::ostream& os) {
  // Print global variables
  if (!module.GetGlobals().empty()) {
    os << ".data\n";
    for (const auto& gv : module.GetGlobals()) {
      os << ".global " << gv.name << "\n";
      os << ".align 4\n";
      os << gv.name << ":\n";
      if (gv.size > 4 || gv.init_value == 0) {
        os << "  .zero " << gv.size << "\n";
      } else {
        os << "  .word " << gv.init_value << "\n";
      }
    }
    os << "\n";
  }

  os << ".text\n";
  for (const auto& function : module.GetFunctions()) {
    os << ".global " << function->GetName() << "\n";
    os << ".type " << function->GetName() << ", %function\n";
    os << function->GetName() << ":\n";

    for (const auto& block : function->GetBlocks()) {
      os << ".L" << function->GetName() << "_" << block->GetName() << ":\n";
      
      for (const auto& inst : block->GetInstructions()) {
        const auto& ops = inst.GetOperands();
        switch (inst.GetOpcode()) {
          case Opcode::Prologue:
            os << "  stp x29, x30, [sp, #-16]!\n";
            os << "  mov x29, sp\n";
            if (function->GetFrameSize() > 0) {
              if (function->GetFrameSize() <= 4095) {
                os << "  sub sp, sp, #" << function->GetFrameSize() << "\n";
              } else {
                PrintMovImm(os, PhysReg::X11, function->GetFrameSize());
                os << "  sub sp, sp, x11\n";
              }
            }
            break;
          case Opcode::Epilogue:
            if (function->GetFrameSize() > 0) {
              if (function->GetFrameSize() <= 4095) {
                os << "  add sp, sp, #" << function->GetFrameSize() << "\n";
              } else {
                PrintMovImm(os, PhysReg::X11, function->GetFrameSize());
                os << "  add sp, sp, x11\n";
              }
            }
            os << "  ldp x29, x30, [sp], #16\n";
            break;
          case Opcode::MovImm:
            if (ops.at(1).GetKind() == Operand::Kind::Global) {
              os << "  adrp " << PhysRegName(ops.at(0).GetReg()) << ", " << ops.at(1).GetGlobal() << "\n";
              os << "  add " << PhysRegName(ops.at(0).GetReg()) << ", " << PhysRegName(ops.at(0).GetReg())
                 << ", :lo12:" << ops.at(1).GetGlobal() << "\n";
            } else {
              PrintMovImm(os, ops.at(0).GetReg(), ops.at(1).GetImm());
            }
            break;
          case Opcode::MovRR: {
            const char* dst = PhysRegName(ops.at(0).GetReg());
            const char* src = PhysRegName(ops.at(1).GetReg());
            if (dst[0] == 's' && src[0] == 'w') {
              os << "  fmov " << dst << ", " << src << "\n";
            } else if (dst[0] == 'w' && src[0] == 's') {
              os << "  fmov " << dst << ", " << src << "\n";
            } else if (dst[0] == 's' && src[0] == 's') {
              os << "  fmov " << dst << ", " << src << "\n";
            } else {
              os << "  mov " << dst << ", " << src << "\n";
            }
            break;
          }
          case Opcode::LoadStack: {
            const auto& slot = GetFrameSlot(*function, ops.at(1));
            PrintStackAccess(os, "ldur", ops.at(0).GetReg(), slot.offset);
            break;
          }
          case Opcode::StoreStack: {
            const auto& slot = GetFrameSlot(*function, ops.at(1));
            PrintStackAccess(os, "stur", ops.at(0).GetReg(), slot.offset);
            break;
          }
          case Opcode::AddrStack: {
            const auto& slot = GetFrameSlot(*function, ops.at(1));
            int offset = slot.offset;
            if (offset >= 0) {
              if (offset <= 4095) {
                os << "  add " << PhysRegName(ops.at(0).GetReg()) << ", x29, #" << offset << "\n";
              } else {
                PrintMovImm(os, PhysReg::X16, offset);
                os << "  add " << PhysRegName(ops.at(0).GetReg()) << ", x29, x16\n";
              }
            } else {
              int abs_offset = -offset;
              if (abs_offset <= 4095) {
                os << "  sub " << PhysRegName(ops.at(0).GetReg()) << ", x29, #" << abs_offset << "\n";
              } else {
                PrintMovImm(os, PhysReg::X16, abs_offset);
                os << "  sub " << PhysRegName(ops.at(0).GetReg()) << ", x29, x16\n";
              }
            }
            break;
          }
          case Opcode::LoadGlobal:
            os << "  adrp x16, " << ops.at(1).GetGlobal() << "\n";
            os << "  add x16, x16, :lo12:" << ops.at(1).GetGlobal() << "\n";
            os << "  ldr " << PhysRegName(ops.at(0).GetReg()) << ", [x16]\n";
            break;
          case Opcode::StoreGlobal:
            os << "  adrp x16, " << ops.at(1).GetGlobal() << "\n";
            os << "  add x16, x16, :lo12:" << ops.at(1).GetGlobal() << "\n";
            os << "  str " << PhysRegName(ops.at(0).GetReg()) << ", [x16]\n";
            break;
          case Opcode::AddRR:
            os << "  add " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::AddRRI:
            os << "  add " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", #" << ops.at(2).GetImm() << "\n";
            break;
          case Opcode::AddRRR_LSL: {
            const char* reg2_name = PhysRegName(ops.at(2).GetReg());
            std::string reg2_str = reg2_name;
            std::string extension = "lsl";
            if (reg2_name[0] == 'w') {
              extension = "sxtw";
            }
            os << "  add " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << reg2_str << ", " << extension << " #" << ops.at(3).GetImm() << "\n";
            break;
          }
          case Opcode::SubRR:
            os << "  sub " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::MulRR:
            os << "  mul " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::SDivRR:
            os << "  sdiv " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::MSubRRR:
            os << "  msub " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << ", "
               << PhysRegName(ops.at(3).GetReg()) << "\n";
            break;
          case Opcode::Sxtw:
            os << "  sxtw " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
            break;
          case Opcode::NegR:
            os << "  neg " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
            break;
          case Opcode::CmpRR:
            os << "  cmp " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
            break;
          case Opcode::CSet:
            os << "  cset " << PhysRegName(ops.at(0).GetReg()) << ", "
               << CondCodeName(ops.at(1).GetCond()) << "\n";
            break;
          case Opcode::FAdd:
            os << "  fadd " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::FSub:
            os << "  fsub " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::FMUL:
            os << "  fmul " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::FDiv:
            os << "  fdiv " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << ", "
               << PhysRegName(ops.at(2).GetReg()) << "\n";
            break;
          case Opcode::FNeg:
            os << "  fneg " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
            break;
          case Opcode::FCmp:
            os << "  fcmp " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
            break;
          case Opcode::FCvtSI2FP:
            os << "  scvtf " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
            break;
          case Opcode::FCvtFP2SI:
            os << "  fcvtzs " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
            break;
          case Opcode::LoadR:
            os << "  ldr " << PhysRegName(ops.at(0).GetReg()) << ", ["
               << PhysRegName(ops.at(1).GetReg()) << "]\n";
            break;
          case Opcode::StoreR:
            os << "  str " << PhysRegName(ops.at(0).GetReg()) << ", ["
               << PhysRegName(ops.at(1).GetReg()) << "]\n";
            break;
          case Opcode::Call:
            os << "  bl " << ops.at(0).GetLabel() << "\n";
            break;
          case Opcode::B:
            os << "  b .L" << function->GetName() << "_" << ops.at(0).GetLabel() << "\n";
            break;
          case Opcode::BCond:
            os << "  cmp " << PhysRegName(ops.at(1).GetReg()) << ", #0\n";
            os << "  b." << CondCodeName(ops.at(0).GetCond()) << " .L" << function->GetName() << "_" << ops.at(2).GetLabel() << "\n";
            break;
          case Opcode::Ret:
            os << "  ret\n";
            break;
        }
      }
    }
    os << ".size " << function->GetName() << ", .-" << function->GetName() << "\n\n";
  }
}

}  // namespace mir