#pragma once

#include <initializer_list>
#include <iosfwd>
#include <memory>
#include <string>
#include <vector>

namespace ir
{
  class Module;
}

namespace mir
{

  class MIRContext
  {
  public:
    MIRContext() = default;
  };

  MIRContext &DefaultContext();

  enum class PhysReg
  {
    W0,
    W1,
    W2,
    W3,
    W4,
    W5,
    W6,
    W7,
    W8,
    W9,
    W10,
    W11,
    W12,
    W13,
    W14,
    W15,
    W16,
    W17,
    W18,
    W19,
    W20,
    W21,
    W22,
    W23,
    W24,
    W25,
    W26,
    W27,
    W28,
    W29,
    W30,
    X0,
    X1,
    X2,
    X3,
    X4,
    X5,
    X6,
    X7,
    X8,
    X9,
    X10,
    X11,
    X12,
    X13,
    X14,
    X15,
    X16,
    X17,
    X18,
    X19,
    X20,
    X21,
    X22,
    X23,
    X24,
    X25,
    X26,
    X27,
    X28,
    X29,
    X30,
    S0,
    S1,
    S2,
    S3,
    S4,
    S5,
    S6,
    S7,
    S8,
    S9,
    S10,
    S11,
    S12,
    S13,
    S14,
    S15,
    S16,
    S17,
    S18,
    S19,
    S20,
    S21,
    S22,
    S23,
    S24,
    S25,
    S26,
    S27,
    S28,
    S29,
    S30,
    S31,
    XZR,
    SP,
    WZR
  };

  const char *PhysRegName(PhysReg reg);

  enum class VRegClass
  {
    Int,
    Float,
    Ptr
  };

  enum class Opcode
  {
    Prologue,
    Epilogue,
    MovImm,
    LoadStack,
    StoreStack,
    LoadStackAddr,
    LoadGlobal,
    StoreGlobal,
    LoadGlobalAddr,
    LoadMem,
    StoreMem,
    AddRR,
    SubRR,
    MulRR,
    DivRR,
    ModRR,
    AndRR,
    OrRR,
    XorRR,
    ShlRR,
    ShrRR,
    AsrRR,
    Asr64RR,
    Uxtw,
    Sxtw,
    CmpRR,
    CmpImm,
    FCmpRR,
    CSet,
    Csel,
    Smull,
    Msub,
    NegRR,
    FAddRR,
    FSubRR,
    FMulRR,
    FDivRR,
    Scvtf,
    FCvtzs,
    FMovWS,
    Br,
    CondBr,
    Call,
    Ret,
    LoadAddr,
    MovReg,
  };

  enum class CondCode
  {
    EQ,
    NE,
    LT,
    LE,
    GT,
    GE
  };

  class Operand
  {
  public:
    enum class Kind
    {
      Reg,
      VReg,
      Imm,
      FrameIndex,
      Label,
      Symbol
    };

    static Operand Reg(PhysReg reg);
    static Operand VReg(int id, VRegClass vreg_class);
    static Operand Imm(int value);
    static Operand FrameIndex(int index);
    static Operand Label(int label_id);
    static Operand Symbol(std::string symbol);

    Kind GetKind() const { return kind_; }
    PhysReg GetReg() const { return reg_; }
    int GetImm() const { return imm_; }
    int GetFrameIndex() const { return imm_; }
    int GetLabel() const { return imm_; }
    const std::string &GetSymbol() const { return symbol_; }
    int GetVRegId() const { return imm_; }
    VRegClass GetVRegClass() const { return vreg_class_; }

  private:
    Operand(Kind kind, PhysReg reg, int imm,
            VRegClass vreg_class = VRegClass::Int, std::string symbol = "");

    Kind kind_;
    PhysReg reg_;
    int imm_;
    std::string symbol_;
    VRegClass vreg_class_;
  };

  class MachineInstr
  {
  public:
    MachineInstr(Opcode opcode, std::vector<Operand> operands = {});

    Opcode GetOpcode() const { return opcode_; }
    const std::vector<Operand> &GetOperands() const { return operands_; }
    std::vector<Operand> &GetOperands() { return operands_; }

    bool IsRematerializable() const { return is_rematerializable_; }
    MachineInstr &SetRematerializable(bool val)
    {
      is_rematerializable_ = val;
      return *this;
    }

    int GetRematImm() const { return remat_imm_; }
    MachineInstr &SetRematImm(int val)
    {
      remat_imm_ = val;
      return *this;
    }

  private:
    Opcode opcode_;
    std::vector<Operand> operands_;
    bool is_rematerializable_ = false;
    int remat_imm_ = 0;
  };

  struct FrameSlot
  {
    int index = 0;
    int size = 4;
    int offset = 0;
    bool is_stack_arg = false;
    bool is_callee_stack_arg = false;
  };

  class MachineBasicBlock
  {
  public:
    explicit MachineBasicBlock(std::string name, int label_id = -1);

    const std::string &GetName() const { return name_; }
    int GetLabelId() const { return label_id_; }
    void SetLabelId(int label_id) { label_id_ = label_id; }

    std::vector<MachineInstr> &GetInstructions() { return instructions_; }
    const std::vector<MachineInstr> &GetInstructions() const { return instructions_; }

    MachineInstr &Append(Opcode opcode,
                         std::initializer_list<Operand> operands = {});

  private:
    std::string name_;
    int label_id_ = -1;
    std::vector<MachineInstr> instructions_;
  };

  class MachineFunction
  {
  public:
    explicit MachineFunction(std::string name);

    const std::string &GetName() const { return name_; }

    MachineBasicBlock &GetEntry() { return *entry_; }
    const MachineBasicBlock &GetEntry() const { return *entry_; }

    MachineBasicBlock *GetEntryPtr() { return entry_; }
    const MachineBasicBlock *GetEntryPtr() const { return entry_; }

    MachineBasicBlock &CreateBlock(std::string name);
    MachineBasicBlock *FindBlock(const std::string &name);
    const MachineBasicBlock *FindBlock(const std::string &name) const;

    std::vector<std::unique_ptr<MachineBasicBlock>> &GetBlocks()
    {
      return blocks_;
    }
    const std::vector<std::unique_ptr<MachineBasicBlock>> &GetBlocks() const
    {
      return blocks_;
    }

    int CreateLabel();

    int CreateFrameIndex(int size = 4);
    int CreateStackArgFrameIndex(int size = 4);
    int CreateCalleeStackArgFrameIndex(int size = 4);
    FrameSlot &GetFrameSlot(int index);
    const FrameSlot &GetFrameSlot(int index) const;
    const std::vector<FrameSlot> &GetFrameSlots() const { return frame_slots_; }
    std::vector<FrameSlot> &GetFrameSlots() { return frame_slots_; }

    int GetFrameSize() const { return frame_size_; }
    void SetFrameSize(int size) { frame_size_ = size; }

    int CreateVReg(VRegClass vreg_class);
    VRegClass GetVRegClass(int vreg_id) const;
    int GetNumVRegs() const { return static_cast<int>(vreg_classes_.size()); }

    void AddCalleeSavedReg(PhysReg reg);
    const std::vector<PhysReg> &GetCalleeSavedRegs() const { return callee_saved_regs_; }

  private:
    std::string name_;
    std::vector<std::unique_ptr<MachineBasicBlock>> blocks_;
    MachineBasicBlock *entry_ = nullptr;

    std::vector<FrameSlot> frame_slots_;
    int frame_size_ = 0;
    int next_label_id_ = 0;

    std::vector<VRegClass> vreg_classes_;
    std::vector<PhysReg> callee_saved_regs_;
  };

  struct MachineGlobal
  {
    enum class Kind
    {
      I32Scalar,
      I32Array
    };

    std::string name;
    Kind kind = Kind::I32Scalar;
    int init_value = 0;
    size_t array_size = 0;
    std::vector<int> init_values;
  };

  class MachineModule
  {
  public:
    MachineModule() = default;

    MachineFunction &CreateFunction(std::string name);
    MachineFunction *GetFunction(const std::string &name);
    const MachineFunction *GetFunction(const std::string &name) const;

    std::vector<std::unique_ptr<MachineFunction>> &GetFunctions()
    {
      return functions_;
    }
    const std::vector<std::unique_ptr<MachineFunction>> &GetFunctions() const
    {
      return functions_;
    }

    void AddGlobalI32(std::string name, int init_value)
    {
      MachineGlobal g;
      g.name = std::move(name);
      g.kind = MachineGlobal::Kind::I32Scalar;
      g.init_value = init_value;
      globals_.push_back(std::move(g));
    }

    void AddGlobalArrayI32(std::string name, size_t array_size,
                           std::vector<int> init_values = {})
    {
      MachineGlobal g;
      g.name = std::move(name);
      g.kind = MachineGlobal::Kind::I32Array;
      g.array_size = array_size;
      g.init_values = std::move(init_values);
      globals_.push_back(std::move(g));
    }

    std::vector<MachineGlobal> &GetGlobals() { return globals_; }
    const std::vector<MachineGlobal> &GetGlobals() const { return globals_; }

  private:
    std::vector<std::unique_ptr<MachineFunction>> functions_;
    std::vector<MachineGlobal> globals_;
  };

  std::unique_ptr<MachineModule> LowerModuleToMIR(const ir::Module &module);
  std::unique_ptr<MachineFunction> LowerToMIR(const ir::Module &module);

  void RunRegAlloc(MachineFunction &function);
  void RunRegAlloc(MachineModule &module);

  void RunFrameLowering(MachineFunction &function);
  void RunFrameLowering(MachineModule &module);

  void RunPeephole(MachineFunction &function);
  void RunPeephole(MachineModule &module);

  void PrintAsm(const MachineFunction &function, std::ostream &os);
  void PrintAsm(const MachineModule &module, std::ostream &os);

} // namespace mir