#pragma once
#include "common.h"
#include "llir.h"

using std::make_shared;

namespace CompSysY {

// a?, t?, ra are caller-saved
// s? are callee-saved
// x0,gp,tp are preserved in user-space
enum class RV64Reg {
  x0 = 0, // zero N/A
  ra,     // ra, caller, treated as callee-saved
  sp,     // sp, callee
  gp,     // gp, N/A
  tp,     // tp, N/A
  t0,     // t0, caller
  t1,
  t2, // t1-2, caller
  s0, // s0/fp, callee
  s1, // s1, callee
  a0,
  a1, // a0-1,caller
  a2,
  a3,
  a4,
  a5,
  a6,
  a7, // a2-7,caller
  s2,
  s3,
  s4,
  s5,
  s6,
  s7,
  s8,
  s9,
  s10,
  s11, // s2-11,callee
  t3,
  t4,
  t5,
  t6, // t3-6,caller
};
inline const char *enum_to_string(const RV64Reg &reg) {
  const static char *_str_tab[] = {
    "x0", "ra", "sp", "gp", "tp", "t0", "t1", "t2", "s0", "s1", "a0",  "a1",  "a2", "a3", "a4", "a5",
    "a6", "a7", "s2", "s3", "s4", "s5", "s6", "s7", "s8", "s9", "s10", "s11", "t3", "t4", "t5", "t6",
  };
  assert((unsigned)reg < 32);
  return _str_tab[(unsigned)reg];
}

inline bool is_callee_saved(const RV64Reg &reg) {
  return reg == RV64Reg::sp || reg == RV64Reg::ra || reg == RV64Reg::s0 || reg == RV64Reg::s1 ||
         (RV64Reg::s2 <= reg && reg <= RV64Reg::s11);
}

// riscv calling convention see: https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-cc.adoc

const unsigned XLEN      = 8;
const unsigned XLEN_MASK = XLEN - 1;
inline unsigned xlen_rnd_up(unsigned src) {
  return (src & XLEN_MASK) ? (src & (~XLEN_MASK)) + XLEN : src;
}
inline bool is_in_imm_range(int imm) {
  return -2048 <= imm && imm <= 2047;
}

inline std::string RV64_RegName(RV64Reg reg) {
  std::string regname = "x" + std::to_string((int)reg);
  return regname;
}

inline RV64Reg RV64_RegOffset(RV64Reg reg, int offset) {
  auto xi = (int)reg + offset;
  return (RV64Reg)xi;
}

class MOperand;
class MInst;
class MBasicBlock;
class MFunction;
enum class MOpTag {
  Invalid,
  Virt,
  Imm,
  PreColor,
  Colored,
  Glob,
};
inline const char *enum_to_string(const MOpTag &tag) {
  const static char *_str_tab[] = {"Invalid", "Virt", "Imm", "PreColor", "Colored", "Glob"};
  return _str_tab[(unsigned)tag];
}

class MOperand {
public:
  MOpTag op_type = MOpTag::Invalid;
  int value      = ~0;

  static inline std::unordered_map<sptr(GlobalVar), int> *mcmod_global_id_map = nullptr;
  static inline std::vector<sptr(GlobalVar)> *mcmod_global_list               = nullptr;

  static MOperand VirtReg(int reg_no) {
    auto mop = MOperand({MOpTag::Virt, reg_no});
    return mop;
  }

  static MOperand Imm(int imm_val) {
    auto mop = MOperand({MOpTag::Imm, imm_val});
    return mop;
  }

  static MOperand PreClrReg(RV64Reg phy_reg) {
    auto mop = MOperand({MOpTag::PreColor, (int)phy_reg});
    return mop;
  }

  static MOperand AllocReg(RV64Reg phy_reg) {
    return MOperand({MOpTag::Colored, (int)phy_reg});
  }

  static MOperand Glob(sptr(GlobalVar) glob) {
    return MOperand({MOpTag::Glob, mcmod_global_id_map->at(glob)});
  }

  bool operator==(const MOperand &op2) const {
    return op_type == op2.op_type && value == op2.value;
  }
  bool operator!=(const MOperand &op2) const {
    return op_type != op2.op_type || value != op2.value;
  }
  bool operator<(const MOperand &op2) const {
    return op_type == op2.op_type ? value < op2.value : op_type < op2.op_type;
  }
  bool is_precolored() const {
    return op_type == MOpTag::PreColor;
  }
  bool is_virt() const {
    return op_type == MOpTag::Virt;
  }
  bool is_imm() const {
    return op_type == MOpTag::Imm;
  }
  bool is_glob() const {
    return op_type == MOpTag::Glob;
  }
  bool is_reg() const {
    return op_type == MOpTag::PreColor || op_type == MOpTag::Virt || op_type == MOpTag::Colored;
  }
  std::string to_string() const;
};

struct MOperandHash {
  std::size_t operator()(MOperand const &op) const {
    using std::hash;
    auto h1 = std::hash<int>{}((int)op.op_type);
    auto h2 = std::hash<int>{}(op.value);
    return h1 ^ (h2 << 1);
  }
  std::size_t operator()(std::pair<MOperand, MOperand> const &op) const {
    using std::hash;
    auto h1 = MOperandHash{}(op.first);
    auto h2 = MOperandHash{}(op.second);
    return h1 ^ (h2 << 1);
  }
};

enum class MInstTag {
  Add,
  Sub,
  Mul,
  Div,
  Mod,
  Lt,
  Le,
  Ge,
  Gt,
  Eq,
  Ne,
  Lsh,  // sll
  Rsh,  // srl,sra
  Move, // actually a pseudo, mv = addi rt, rs, 0
  Branch,
  Jmp,
  Ret,
  Load,
  Store,
  Call,
  Comment,
};

inline MInstTag inverse_cond(MInstTag src_tag) {
  switch (src_tag) {
    case MInstTag::Lt: return MInstTag::Ge;
    case MInstTag::Le: return MInstTag::Gt;
    case MInstTag::Ge: return MInstTag::Lt;
    case MInstTag::Gt: return MInstTag::Le;
    case MInstTag::Eq: return MInstTag::Ne;
    case MInstTag::Ne: return MInstTag::Eq;
    default: assert(0);
  }
}

// Instruction insertion type
enum class InsType { Before, After, Head, Tail };

class MInst {
public:
  MInstTag inst_tag;
  // Notes about WIDTH:
  // width is important when load/store to an array address
  // in other cases, since stack is XLEN-padded, for single signed int variable, ld/lw sd/sw do not make much difference
  // Similarly, it makes difference when computing address
  int width = 4;

  std::list<sptr(MInst)>::iterator itr;
  sptr(MBasicBlock) parent_bb;
  MInst(MInstTag tag, sptr(MBasicBlock) parent_bb) : inst_tag(tag), parent_bb(parent_bb) {}
  virtual ~MInst()                              = default;
  virtual std::vector<MOperand> get_def() const = 0;
  virtual std::vector<MOperand> get_use() const = 0;
  virtual bool is_transfer() const {
    return false;
  }

  static void PostNew(sptr(MInst) inst, sptr(MBasicBlock) parent_bb, InsType ins_ty);
  static void PostNew(sptr(MInst) inst, sptr(MInst) rel_inst, InsType ins_ty);
};

class MBasicBlock {
public:
  int id;
  std::list<sptr(MBasicBlock)>::iterator itr;
  sptr(BasicBlock) ir_bb;
  sptr(MFunction) parent_func;
  std::list<sptr(MInst)> inst_list;
  std::list<sptr(MBasicBlock)> pred_list;
  std::list<sptr(MBasicBlock)> succ_list;

  std::set<MOperand> use;
  std::set<MOperand> def;
  std::set<MOperand> livein;
  std::set<MOperand> liveout;

  std::string to_string();
};

class MFunction {
public:
  int id;
  std::list<sptr(MBasicBlock)> bb_list;
  sptr(Function) ir_func;

  std::list<sptr(MInst)> stack_arg_reloc;
  std::unordered_set<RV64Reg> regs_to_save;
  unsigned stack_size; // only for local var, excluding the callee-saved regs

  unsigned virt_reg_cnt = 0;
};

class MCModule {
public:
  std::string file_name;
  std::list<sptr(MFunction)> function_list;
  std::vector<sptr(GlobalVar)> global_list;
  std::unordered_map<sptr(GlobalVar), int> global_id_map;
  void IR2MC(const Module &ir_module);
  void MC2ASM(std::ostream &ostr, bool debug = false) const;
};

class MInstBinary : public MInst {
public:
  MOperand dst;
  MOperand op1;
  MOperand op2;
  MInstBinary(MInstTag type, sptr(MBasicBlock) parent_bb) : MInst(type, parent_bb) {}

  static sptr(MInstBinary) New(MInstTag type, sptr(MBasicBlock) parent_bb, int width, InsType ins_ty = InsType::Tail) {
    auto inst   = make_shared<MInstBinary>(type, parent_bb);
    inst->width = width;
    MInst::PostNew(inst, parent_bb, ins_ty);
    return inst;
  }
  static sptr(MInstBinary) New(MInstTag type, sptr(MInst) rel_inst, int width, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstBinary>(type, parent_bb);
    inst->width       = width;
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
};

class MInstJump : public MInst {
public:
  sptr(MBasicBlock) target;
  MInstJump(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Jmp, parent_bb) {}

  static sptr(MInstJump) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstJump>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
  bool is_transfer() const override {
    return true;
  }
};

class MInstBranch : public MInst {
public:
  sptr(MBasicBlock) target;
  MOperand op1;
  MOperand op2;
  MInstTag branch_tag;
  MInstBranch(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Branch, parent_bb) {}
  static sptr(MInstBranch) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstBranch>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
  bool is_transfer() const override {
    return true;
  }
};

class MInstLoad : public MInst {
public:
  MOperand dst;
  MOperand addr;
  MOperand offset;
  MInstLoad(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Load, parent_bb) {}

  static sptr(MInstLoad) New(sptr(MBasicBlock) parent_bb, int width) {
    auto inst = make_shared<MInstLoad>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    inst->width = width;
    return inst;
  }
  static sptr(MInstLoad) New(sptr(MInst) rel_inst, int width, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstLoad>(parent_bb);
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    inst->width = width;
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
};

class MInstStore : public MInst {
public:
  MOperand data;
  MOperand addr;
  MOperand offset;
  MInstStore(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Store, parent_bb) {}

  static sptr(MInstStore) New(sptr(MBasicBlock) parent_bb, int width) {
    auto inst = make_shared<MInstStore>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    inst->width = width;
    return inst;
  }
  static sptr(MInstStore) New(sptr(MInst) rel_inst, int width, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstStore>(parent_bb);
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    inst->width = width;
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
};

class MInstMove : public MInst {
public:
  MOperand dst;
  MOperand src;
  MInstMove(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Move, parent_bb) {}
  static sptr(MInstMove) New(sptr(MBasicBlock) parent_bb, bool insert_begin = false) {
    auto inst = make_shared<MInstMove>(parent_bb);
    if (insert_begin)
      parent_bb->inst_list.push_front(inst);
    else
      parent_bb->inst_list.push_back(inst);
    return inst;
  }
  static sptr(MInstMove) New(sptr(MInst) rel_inst, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstMove>(parent_bb);
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
};

// class MInstSymbol : public MInst {
// public:
//   MOperand dst;
//   sptr(GlobalVar) symbol;
//   MInstSymbol(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Globsym, parent_bb) {}
//   static sptr(MInstSymbol) New(sptr(MBasicBlock) parent_bb, bool insert_begin = false) {
//     auto inst = make_shared<MInstSymbol>(parent_bb);
//     if (insert_begin)
//       parent_bb->inst_list.push_front(inst);
//     else
//       parent_bb->inst_list.push_back(inst);
//     return inst;
//   }
// };

class MInstReturn : public MInst {
public:
  MInstReturn(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Ret, parent_bb) {}
  static sptr(MInstReturn) New(sptr(MBasicBlock) parent_bb, bool insert_begin = false) {
    auto inst = make_shared<MInstReturn>(parent_bb);
    if (insert_begin)
      parent_bb->inst_list.push_front(inst);
    else
      parent_bb->inst_list.push_back(inst);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
  bool is_transfer() const override {
    return true;
  }
};

class MInstComment : public MInst {
public:
  std::string comment;
  MInstComment(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Comment, parent_bb) {}
  static sptr(MInstComment) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstComment>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
  static sptr(MInstComment) New(sptr(MInst) rel_inst, bool insert_after = false) {
    auto parent_bb    = rel_inst->parent_bb;
    auto inst         = make_shared<MInstComment>(parent_bb);
    auto rel_inst_itr = STD_FIND(parent_bb->inst_list, rel_inst);
    assert(rel_inst_itr != parent_bb->inst_list.end());
    if (insert_after) std::advance(rel_inst_itr, 1);
    parent_bb->inst_list.insert(rel_inst_itr, inst);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
};

class MInstCall : public MInst {
public:
  sptr(Function) ir_func;
  MInstCall(sptr(MBasicBlock) parent_bb) : MInst(MInstTag::Call, parent_bb) {}
  static sptr(MInstCall) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstCall>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand> get_use() const override;
};

void get_inst_defuse(sptr(MInst) inst, std::vector<MOperand *> &def, std::vector<MOperand *> &use);
} // namespace CompSysY