#include "mir/MIR.h"
#include "ir/IR.h"

#include <ostream>
#include <stdexcept>
#include <cstdint>
#include <vector>
#include <cstring>

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

bool IsFloatReg(PhysReg reg) {
  return reg >= PhysReg::S0 && reg <= PhysReg::S15;
}

void PrintStackAccess(std::ostream& os, const char* mnemonic, PhysReg reg,
                      int offset, int frame_size) {
  bool is_float = IsFloatReg(reg);
  const char* ldr_cmd = is_float ? "ldr" : "ldr";
  const char* str_cmd = is_float ? "str" : "str";
  const char* base_mnemonic = (std::strcmp(mnemonic, "ldur") == 0) ? ldr_cmd : str_cmd;

  if (offset >= -256 && offset <= 255) {
    if (is_float) {
      os << "  " << base_mnemonic << " " << PhysRegName(reg) << ", [x29, #" << offset << "]\n";
    } else {
      os << "  " << mnemonic << " " << PhysRegName(reg) << ", [x29, #" << offset << "]\n";
    }
  } else {
    int sp_offset = frame_size + offset;
    int access_size = 4;
    if ((reg >= PhysReg::X0 && reg <= PhysReg::X28) ||
        reg == PhysReg::X29 || reg == PhysReg::X30 ||
        reg == PhysReg::SP) {
      access_size = 8;
    }
    int max_offset = access_size == 8 ? 32760 : 16380;
    if (sp_offset >= 0 && sp_offset <= max_offset &&
        sp_offset % access_size == 0) {
      os << "  " << base_mnemonic << " " << PhysRegName(reg)
         << ", [sp, #" << sp_offset << "]\n";
    } else {
      os << "  ldr x10, =" << offset << "\n";
      os << "  " << base_mnemonic << " " << PhysRegName(reg) << ", [x29, x10]\n";
    }
  }
}

std::string GetBlockLabel(const std::string& func_name, const std::string& block_name) {
  if (block_name == "entry" || block_name.empty()) {
    return func_name;
  }
  return ".L_" + func_name + "_" + block_name;
}

}  // namespace

void PrintAsm(const MachineFunction& function, std::ostream& os) {
  os << ".text\n";
  os << ".global " << function.GetName() << "\n";
  os << ".type " << function.GetName() << ", %function\n";

  struct FloatConstant {
    std::string label;
    float value;
  };
  std::vector<FloatConstant> float_constants;

  for (size_t b = 0; b < function.GetBlocks().size(); ++b) {
    const auto& block = function.GetBlocks()[b];
    
    // Print the block label
    if (b == 0) {
      os << function.GetName() << ":\n";
    } else {
      os << GetBlockLabel(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) {
            int size = function.GetFrameSize();
            if (size <= 4095) {
              os << "  sub sp, sp, #" << size << "\n";
            } else {
              os << "  ldr x9, =" << size << "\n";
              os << "  sub sp, sp, x9\n";
            }
          }
          break;
        case Opcode::Epilogue:
          if (function.GetFrameSize() > 0) {
            int size = function.GetFrameSize();
            if (size <= 4095) {
              os << "  add sp, sp, #" << size << "\n";
            } else {
              os << "  ldr x9, =" << size << "\n";
              os << "  add sp, sp, x9\n";
            }
          }
          os << "  ldp x29, x30, [sp], #16\n";
          break;
        case Opcode::MovImm: {
          PhysReg dst = ops.at(0).GetReg();
          if (IsFloatReg(dst)) {
            // Load float constant
            int bits = ops.at(1).GetImm();
            float val;
            std::memcpy(&val, &bits, sizeof(float));
            std::string flabel = ".LC_" + function.GetName() + "_" + std::to_string(float_constants.size());
            float_constants.push_back({flabel, val});
            
            os << "  adrp x8, " << flabel << "\n";
            os << "  ldr " << PhysRegName(dst) << ", [x8, :lo12:" << flabel << "]\n";
          } else {
            int imm = ops.at(1).GetImm();
            if (imm >= 0 && imm <= 65535) {
              os << "  mov " << PhysRegName(dst) << ", #" << imm << "\n";
            } else {
              os << "  ldr " << PhysRegName(dst) << ", =" << imm << "\n";
            }
          }
          break;
        }
        case Opcode::LoadStack: {
          const auto& slot = GetFrameSlot(function, ops.at(1));
          PrintStackAccess(os, "ldur", ops.at(0).GetReg(), slot.offset,
                           function.GetFrameSize());
          break;
        }
        case Opcode::StoreStack: {
          const auto& slot = GetFrameSlot(function, ops.at(1));
          PrintStackAccess(os, "stur", ops.at(0).GetReg(), slot.offset,
                           function.GetFrameSize());
          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::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::MSubRRRR:
          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::FAddRRR:
          os << "  fadd " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << ", "
             << PhysRegName(ops.at(2).GetReg()) << "\n";
          break;
        case Opcode::FSubRRR:
          os << "  fsub " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << ", "
             << PhysRegName(ops.at(2).GetReg()) << "\n";
          break;
        case Opcode::FMulRRR:
          os << "  fmul " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << ", "
             << PhysRegName(ops.at(2).GetReg()) << "\n";
          break;
        case Opcode::FDivRRR:
          os << "  fdiv " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << ", "
             << PhysRegName(ops.at(2).GetReg()) << "\n";
          break;
        case Opcode::CmpRR:
          os << "  cmp " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << "\n";
          break;
        case Opcode::FCmpRR:
          os << "  fcmp " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << "\n";
          break;
        case Opcode::Cset:
          os << "  cset " << PhysRegName(ops.at(0).GetReg()) << ", "
             << ops.at(1).GetCondCode() << "\n";
          break;
        case Opcode::B:
          os << "  b " << GetBlockLabel(function.GetName(), ops.at(0).GetLabelName()) << "\n";
          break;
        case Opcode::BCond:
          os << "  b." << ops.at(0).GetCondCode() << " "
             << GetBlockLabel(function.GetName(), ops.at(1).GetLabelName()) << "\n";
          break;
        case Opcode::Call:
          os << "  bl " << ops.at(0).GetGlobalName() << "\n";
          break;
        case Opcode::Ret:
          os << "  ret\n";
          break;
        case Opcode::MovReg:
          if (IsFloatReg(ops.at(0).GetReg()) || IsFloatReg(ops.at(1).GetReg())) {
            os << "  fmov " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
          } else {
            os << "  mov " << PhysRegName(ops.at(0).GetReg()) << ", "
               << PhysRegName(ops.at(1).GetReg()) << "\n";
          }
          break;
        case Opcode::Adrp:
          os << "  adrp " << PhysRegName(ops.at(0).GetReg()) << ", "
             << ops.at(1).GetGlobalName() << "\n";
          break;
        case Opcode::AddRegImm: {
          PhysReg dst = ops.at(0).GetReg();
          PhysReg src = ops.at(1).GetReg();
          if (ops.at(2).GetKind() == Operand::Kind::FrameIndex) {
            const auto& slot = function.GetFrameSlot(ops.at(2).GetFrameIndex());
            int offset = slot.offset;
            if (offset >= -4095 && offset <= 4095) {
              if (offset >= 0) {
                os << "  add " << PhysRegName(dst) << ", " << PhysRegName(src) << ", #" << offset << "\n";
              } else {
                os << "  sub " << PhysRegName(dst) << ", " << PhysRegName(src) << ", #" << (-offset) << "\n";
              }
            } else {
              os << "  ldr x9, =" << offset << "\n";
              os << "  add " << PhysRegName(dst) << ", " << PhysRegName(src) << ", x9\n";
            }
          } else if (ops.at(2).GetKind() == Operand::Kind::Global) {
            os << "  add " << PhysRegName(dst) << ", " << PhysRegName(src) << ", :lo12:" << ops.at(2).GetGlobalName() << "\n";
          } else {
            int imm = ops.at(2).GetImm();
            if (imm >= -4095 && imm <= 4095) {
              if (imm >= 0) {
                os << "  add " << PhysRegName(dst) << ", " << PhysRegName(src) << ", #" << imm << "\n";
              } else {
                os << "  sub " << PhysRegName(dst) << ", " << PhysRegName(src) << ", #" << (-imm) << "\n";
              }
            } else {
              os << "  ldr x9, =" << imm << "\n";
              os << "  add " << PhysRegName(dst) << ", " << PhysRegName(src) << ", x9\n";
            }
          }
          break;
        }
        case Opcode::LslImm:
          os << "  lsl " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << ", #"
             << ops.at(2).GetImm() << "\n";
          break;
        case Opcode::LdrRegReg: {
          PhysReg reg = ops.at(0).GetReg();
          const char* ldr_cmd = IsFloatReg(reg) ? "ldr" : "ldr";
          os << "  " << ldr_cmd << " " << PhysRegName(reg) << ", ["
             << PhysRegName(ops.at(1).GetReg()) << "]\n";
          break;
        }
        case Opcode::StrRegReg: {
          PhysReg reg = ops.at(0).GetReg();
          const char* str_cmd = IsFloatReg(reg) ? "str" : "str";
          os << "  " << str_cmd << " " << PhysRegName(reg) << ", ["
             << PhysRegName(ops.at(1).GetReg()) << "]\n";
          break;
        }
        case Opcode::SIToFP:
          os << "  scvtf " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << "\n";
          break;
        case Opcode::FPToSI:
          os << "  fcvtzs " << PhysRegName(ops.at(0).GetReg()) << ", "
             << PhysRegName(ops.at(1).GetReg()) << "\n";
          break;
        case Opcode::ZExt:
          if (ops.at(0).GetReg() >= PhysReg::X0 && ops.at(0).GetReg() <= PhysReg::X28) {
            os << "  sxtw " << PhysRegName(ops.at(0).GetReg()) << ", " << PhysRegName(ops.at(1).GetReg()) << "\n";
          } else {
            os << "  and " << PhysRegName(ops.at(0).GetReg()) << ", " << PhysRegName(ops.at(1).GetReg()) << ", #1\n";
          }
          break;
      }
    }
  }

  os << ".size " << function.GetName() << ", .-" << function.GetName() << "\n";

  // Print read-only data segment if there are float constants
  if (!float_constants.empty()) {
    os << ".section .rodata\n";
    os << ".align 2\n";
    for (const auto& fc : float_constants) {
      os << fc.label << ":\n";
      uint32_t bits;
      std::memcpy(&bits, &fc.value, sizeof(float));
      os << "  .word " << bits << " // float " << fc.value << "\n";
    }
  }
}

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

void PrintGlobals(const ir::Module& module, std::ostream& os) {
  for (const auto& gv : module.GetGlobalValues()) {
    os << ".global " << gv->GetName() << "\n";
    
    std::shared_ptr<ir::Type> actual_ty = gv->GetType();
    if (actual_ty->IsPtrInt32()) actual_ty = ir::Type::GetInt32Type();
    else if (actual_ty->IsPtrFloat()) actual_ty = ir::Type::GetFloatType();
    
    uint32_t actual_size = GetTypeSize(actual_ty.get());

    if (gv->GetInitializer()) {
      os << ".data\n";
      os << ".align 2\n";
      os << ".size " << gv->GetName() << ", " << actual_size << "\n";
      os << gv->GetName() << ":\n";
      
      if (actual_ty->IsFloat()) {
        float val = 0.0f;
        if (auto* cf = dynamic_cast<const ir::ConstantFloat*>(gv->GetInitializer())) {
          val = cf->GetValue();
        } else if (auto* ci = dynamic_cast<const ir::ConstantInt*>(gv->GetInitializer())) {
          val = static_cast<float>(ci->GetValue());
        }
        uint32_t bits;
        std::memcpy(&bits, &val, sizeof(float));
        os << "  .word " << bits << " // float " << val << "\n";
      } else {
        int val = 0;
        if (auto* ci = dynamic_cast<const ir::ConstantInt*>(gv->GetInitializer())) {
          val = ci->GetValue();
        } else if (auto* cf = dynamic_cast<const ir::ConstantFloat*>(gv->GetInitializer())) {
          val = static_cast<int>(cf->GetValue());
        }
        os << "  .word " << val << "\n";
      }
    } else {
      os << ".bss\n";
      os << ".align 4\n";
      os << ".size " << gv->GetName() << ", " << actual_size << "\n";
      os << gv->GetName() << ":\n";
      os << "  .zero " << actual_size << "\n";
    }
    os << "\n";
  }
}

}  // namespace mir
