#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 InsertPos
{
  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;
  virtual ~MInst() = default;
  virtual std::vector<MOperand> get_def() const
  {
    return {};
  };
  virtual std::vector<MOperand *> get_def_ptr()
  {
    return {};
  };
  virtual std::vector<MOperand> get_use() const
  {
    return {};
  };
  virtual std::vector<MOperand *> get_use_ptr()
  {
    return {};
  };
  virtual bool is_transfer() const
  {
    return false;
  }

  static void PostNew(sptr(MInst) inst, MInstTag tag, sptr(MBasicBlock) parent_bb, InsertPos ins_pos);
  static void PostNew(sptr(MInst) inst, MInstTag tag, sptr(MInst) rel_inst, InsertPos ins_pos);
};

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;

  static sptr(MInstBinary)
    New(MInstTag type, sptr(MBasicBlock) parent_bb, int width, InsertPos ins_pos = InsertPos::Tail)
  {
    auto inst   = make_shared<MInstBinary>();
    inst->width = width;
    MInst::PostNew(inst, type, parent_bb, ins_pos);
    return inst;
  }
  static sptr(MInstBinary) New(MInstTag type, sptr(MInst) rel_inst, int width, InsertPos ins_pos = InsertPos::Before)
  {
    auto inst   = make_shared<MInstBinary>();
    inst->width = width;
    MInst::PostNew(inst, type, rel_inst, ins_pos);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand *> get_def_ptr() override;
  virtual std::vector<MOperand> get_use() const override;
  virtual std::vector<MOperand *> get_use_ptr() override;
};

class MInstJump : public MInst
{
public:
  sptr(MBasicBlock) target;

  static sptr(MInstJump) New(sptr(MBasicBlock) parent_bb)
  {
    auto inst = make_shared<MInstJump>();
    PostNew(inst, MInstTag::Jmp, parent_bb, InsertPos::Tail);
    return inst;
  }
  bool is_transfer() const override
  {
    return true;
  }
};

class MInstBranch : public MInst
{
public:
  sptr(MBasicBlock) target;
  MOperand op1;
  MOperand op2;
  MInstTag branch_tag;
  static sptr(MInstBranch) New(sptr(MBasicBlock) parent_bb)
  {
    auto inst = make_shared<MInstBranch>();
    PostNew(inst, MInstTag::Branch, parent_bb, InsertPos::Tail);
    return inst;
  }
  virtual std::vector<MOperand> get_use() const override;
  virtual std::vector<MOperand *> get_use_ptr() override;
  bool is_transfer() const override
  {
    return true;
  }
};

class MInstLoad : public MInst
{
public:
  MOperand dst;
  MOperand addr;
  MOperand offset;

  static sptr(MInstLoad) New(sptr(MBasicBlock) parent_bb, int width, InsertPos ins_pos)
  {
    auto inst   = make_shared<MInstLoad>();
    inst->width = width;
    PostNew(inst, MInstTag::Load, parent_bb, ins_pos);
    return inst;
  }
  static sptr(MInstLoad) New(sptr(MInst) rel_inst, int width, InsertPos ins_pos)
  {
    auto inst   = make_shared<MInstLoad>();
    inst->width = width;
    PostNew(inst, MInstTag::Load, rel_inst, ins_pos);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand *> get_def_ptr() override;
  virtual std::vector<MOperand> get_use() const override;
  virtual std::vector<MOperand *> get_use_ptr() override;
};

class MInstStore : public MInst
{
public:
  MOperand data;
  MOperand addr;
  MOperand offset;

  static sptr(MInstStore) New(sptr(MBasicBlock) parent_bb, int width, InsertPos ins_pos)
  {
    auto inst   = make_shared<MInstStore>();
    inst->width = width;
    PostNew(inst, MInstTag::Store, parent_bb, ins_pos);
    return inst;
  }
  static sptr(MInstStore) New(sptr(MInst) rel_inst, int width, InsertPos ins_pos)
  {
    auto inst   = make_shared<MInstStore>();
    inst->width = width;
    PostNew(inst, MInstTag::Store, rel_inst, ins_pos);
    return inst;
  }
  virtual std::vector<MOperand> get_use() const override;
  virtual std::vector<MOperand *> get_use_ptr() override;
};

class MInstMove : public MInst
{
public:
  MOperand dst;
  MOperand src;
  static sptr(MInstMove) New(sptr(MBasicBlock) parent_bb, InsertPos ins_pos)
  {
    auto inst = make_shared<MInstMove>();
    PostNew(inst, MInstTag::Move, parent_bb, ins_pos);
    return inst;
  }
  static sptr(MInstMove) New(sptr(MInst) rel_inst, InsertPos ins_pos)
  {
    auto inst = make_shared<MInstMove>();
    PostNew(inst, MInstTag::Move, rel_inst, ins_pos);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand *> get_def_ptr() override;
  virtual std::vector<MOperand> get_use() const override;
  virtual std::vector<MOperand *> get_use_ptr() override;
};

class MInstReturn : public MInst
{
public:
  static sptr(MInstReturn) New(sptr(MBasicBlock) parent_bb, InsertPos ins_pos)
  {
    auto inst = make_shared<MInstReturn>();
    PostNew(inst, MInstTag::Ret, parent_bb, ins_pos);
    return inst;
  }
  virtual std::vector<MOperand> get_use() const override;
  bool is_transfer() const override
  {
    return true;
  }
};

class MInstComment : public MInst
{
public:
  std::string comment;
  static sptr(MInstComment) New(sptr(MBasicBlock) parent_bb, InsertPos ins_pos)
  {
    auto inst = make_shared<MInstComment>();
    PostNew(inst, MInstTag::Comment, parent_bb, ins_pos);
    return inst;
  }
  static sptr(MInstComment) New(sptr(MInst) rel_inst, InsertPos ins_pos)
  {
    auto inst = make_shared<MInstComment>();
    PostNew(inst, MInstTag::Comment, rel_inst, ins_pos);
    return inst;
  }
};

class MInstCall : public MInst
{
public:
  sptr(Function) ir_func;
  static sptr(MInstCall) New(sptr(MBasicBlock) parent_bb, InsertPos ins_pos)
  {
    auto inst = make_shared<MInstCall>();
    PostNew(inst, MInstTag::Call, parent_bb, ins_pos);
    return inst;
  }
  virtual std::vector<MOperand> get_def() const override;
  virtual std::vector<MOperand *> get_def_ptr() override;
  virtual std::vector<MOperand> get_use() const override;
};
} // namespace CompSysY