#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
  ra,     // ra, caller
  sp,     // sp, callee
  gp,     // gp
  tp,     // tp
  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
};

// 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 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;

class MOperand {
public:
  enum class OpTag {
    Virt,
    Imm,
    PreColor,
  } op_type = OpTag::Virt;
  int value = ~0;

  MOperand(OpTag tag, int val) : op_type(tag), value(val) {}
  static sptr(MOperand) NewVirtReg(int reg_no) {
    auto mop = std::make_shared<MOperand>(OpTag::Virt, reg_no);
    return mop;
  }

  static sptr(MOperand) NewImm(int imm_val) {
    auto mop = std::make_shared<MOperand>(OpTag::Imm, imm_val);
    return mop;
  }

  static sptr(MOperand) NewReg(RV64Reg phy_reg) {
    auto mop = std::make_shared<MOperand>(OpTag::PreColor, (int)phy_reg);
    return mop;
  }
};

enum class MInstTag {
  Add,
  Sub,
  Mul,
  Div,
  Mod,
  Lt,
  Le,
  Ge,
  Gt,
  Eq,
  Ne,
  And,
  Or,
  Lsh,  // sll
  Rsh,  // srl,sra
  Move, // actually a pseudo, mv = addi rt, rs, 0
  Branch,
  Jmp,
  Ret,
  Load,
  Store,
  Compare,
  Call,
  Globsym,
  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);
  }
}

class MInst {
public:
  MInstTag inst_tag;

  sptr(MBasicBlock) parent_bb;
  MInst(MInstTag tag, sptr(MBasicBlock) parent_bb) : inst_tag(tag), parent_bb(parent_bb) {}
};

class MBasicBlock {
public:
  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::unordered_set<sptr(MOperand)> livein;
  std::unordered_set<sptr(MOperand)> liveout;
};

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

  std::list<sptr(MInst)> stack_arg_reloc;
  unsigned stack_size;

  unsigned virt_reg_cnt = 0;
};

class MCModule {
public:
  std::list<sptr(MFunction)> function_list;
  std::list<sptr(GlobalVar)> global_list;
  void IR2MC(const Module &ir_module);
};

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

  static sptr(MInstBinary) New(MInstTag type, sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstBinary>(type, parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

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;
  }
};

class MInstBranch : public MInst {
public:
  sptr(MBasicBlock) target;
  sptr(MOperand) op1;
  sptr(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;
  }
};

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

  static sptr(MInstLoad) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstLoad>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

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

  static sptr(MInstStore) New(sptr(MBasicBlock) parent_bb) {
    auto inst = make_shared<MInstStore>(parent_bb);
    parent_bb->inst_list.push_back(inst);
    return inst;
  }
};

class MInstMove : public MInst {
public:
  sptr(MOperand) dst;
  sptr(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 = 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;
  }
};

class MInstSymbol : public MInst {
public:
  sptr(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;
  }
};

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;
  }
};

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;
  }
};

} // namespace CompSysY