#pragma once
#include "common.h"
#include "llir_type.h"
#include "llir_value.h"
#include <cassert>
#include <sstream>
#include <string>

namespace CompSysY
{

DEF_PTR_T(InstAlloca);
DEF_PTR_T(InstStore);
DEF_PTR_T(InstLoad);
DEF_PTR_T(InstBinary);
DEF_PTR_T(InstZext);
DEF_PTR_T(InstBranch);
DEF_PTR_T(InstReturn);
DEF_PTR_T(InstCall);
DEF_PTR_T(InstGEP);
DEF_PTR_T(InstPhi);

enum class InstTag
{
  Add,
  Sub,
  Mod,
  Mul,
  Div,
  Lt,
  Le,
  Ge,
  Gt,
  Eq,
  Ne,
  // And, Or,
  Br,
  Call,
  Ret,
  Alloca,
  Load,
  Store,
  GEP,
  Zext,
  Phi,
  // MemPhi,
  // LoadDep,
  // InsertEle,
  // ExtractEle
};

class Instruction : public User
{
public:
  int ir_seqno = -1;
  InstTag tag;
  BasicBlockPtr_t parent_bb;
  std::list<InstructionPtr_t>::iterator inst_itr;
  Instruction(InstTag inst_tag, TypePtr_t type, BasicBlockPtr_t parent_bb)
    : User("", type), tag(inst_tag), parent_bb(parent_bb)
  {
  }
  std::string tag_string()
  {
    switch (tag)
      {
        case InstTag::Add: return "add";
        case InstTag::Sub: return "sub";
        case InstTag::Mod: return "srem";
        case InstTag::Mul: return "mul";
        case InstTag::Div: return "sdiv";
        case InstTag::Lt: return "icmp slt";
        case InstTag::Le: return "icmp sle";
        case InstTag::Ge: return "icmp sge";
        case InstTag::Gt: return "icmp sgt";
        case InstTag::Eq: return "icmp eq";
        case InstTag::Ne: return "icmp ne";
        // case InstTag::And: return "and";
        // case InstTag::Or: return "or";
        case InstTag::Br: return "br";
        case InstTag::Call: return "call";
        case InstTag::Ret: return "ret";
        case InstTag::Alloca: return "alloca";
        case InstTag::Load: return "load";
        case InstTag::Store: return "store";
        case InstTag::GEP: return "GEP";
        case InstTag::Zext: return "zext";
        case InstTag::Phi:
          return "Phi";
          // case InstTag::MemPhi: return "MemPhi";
          // case InstTag::LoadDep: return "LoadDep";
          // case InstTag::InsertEle: return "InsertEle";
          // case InstTag::ExtractEle: return "ExtractEle";
      }
  }
  virtual std::string to_string() override
  {
    return type->to_string() + " " + tag_string();
  };
};

class InstAlloca : public Instruction
{
public:
  InstAlloca(TypePtr_t alloc_type, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Alloca, std::make_shared<PointerType>(alloc_type), parent_bb)
  {
  }
  virtual std::string to_IR_string() override
  {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstAlloca) New(const std::string &name, TypePtr_t type, sptr(BasicBlock) parent_bb)
  {
    auto inst         = std::make_shared<InstAlloca>(type, parent_bb);
    inst->name        = name;
    auto func_head_bb = parent_bb->parent_func->bb_list.front();
    // put all alloca in the head of each basic block
    inst->inst_itr = func_head_bb->inst_list.insert(func_head_bb->inst_list.begin(), inst);
    return inst;
  }
};

class InstStore : public Instruction
{
public:
  InstStore(std::shared_ptr<Value> value, std::shared_ptr<Value> pointer, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Store, TypeHelper::TYPE_VOID, parent_bb)
  {
    assert(value);
    add_operand(value);
    add_operand(pointer);
  }
  static sptr(InstStore) New(sptr(Value) value, sptr(Value) pointer, sptr(BasicBlock) parent_bb)
  {
    auto inst      = std::make_shared<InstStore>(value, pointer, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

class InstLoad : public Instruction
{
public:
  InstLoad(std::shared_ptr<Value> value, TypePtr_t type, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Load, type, parent_bb)
  {
    add_operand(value);
  }
  virtual std::string to_IR_string() override
  {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstLoad) New(std::shared_ptr<Value> value, TypePtr_t type, std::shared_ptr<BasicBlock> parent_bb)
  {
    auto inst      = std::make_shared<InstLoad>(value, type, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

class InstBinary : public Instruction
{
public:
  InstBinary(
    InstTag inst_tag,
    TypePtr_t val_type,
    std::shared_ptr<Value> op1,
    std::shared_ptr<Value> op2,
    std::shared_ptr<BasicBlock> parent_bb
  )
    : Instruction(inst_tag, val_type, parent_bb)
  {
    add_operand(op1);
    add_operand(op2);
  }
  virtual std::string to_IR_string() override
  {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstBinary)
    New(InstTag inst_tag, std::shared_ptr<Value> op1, std::shared_ptr<Value> op2, std::shared_ptr<BasicBlock> parent_bb)
  {
    std::shared_ptr<InstBinary> inst;
    if (InstTag::Lt <= inst_tag && inst_tag <= InstTag::Ne)
      {
        inst = std::make_shared<InstBinary>(inst_tag, TypeHelper::TYPE_I1, op1, op2, parent_bb);
      }
    else if (InstTag::Add <= inst_tag && inst_tag <= InstTag::Div)
      {
        inst = std::make_shared<InstBinary>(inst_tag, TypeHelper::TYPE_I32, op1, op2, parent_bb);
      }
    else
      {
        panic("Invalid Binary Operation");
      }
    assert(Type::isType<IntegerType>(op1->type));
    assert(Type::isType<IntegerType>(op2->type));
    if (!(Type::asType<IntegerType>(op1->type)->int_type == Type::asType<IntegerType>(op2->type)->int_type))
      {
        if (!(shared_cast<ConstantInt>(op1) && shared_cast<ConstantInt>(op1)->value == 0)
      && !(shared_cast<ConstantInt>(op2) && shared_cast<ConstantInt>(op2)->value == 0))
          {
            LOG(ERROR) << op1->to_string() << ",\t" << op2->to_string();
            assert(0);
          }
      }
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

//<result> = zext <ty> <value> to <ty2>
// zext is integer type, i32
class InstZext : public Instruction
{
public:
  InstZext(std::shared_ptr<Value> op, std::shared_ptr<BasicBlock> parent_bb)
    : Instruction(InstTag::Zext, TypeHelper::TYPE_I32, parent_bb)
  {
    add_operand(op);
  }
  virtual std::string to_IR_string() override
  {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static sptr(InstZext) New(std::shared_ptr<Value> op, std::shared_ptr<BasicBlock> parent_bb)
  {
    auto inst      = std::make_shared<InstZext>(op, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

// Operands: 0:cond, 1:true, 2:false | 0:target
class InstBranch : public Instruction
{
public:
  // conditional branch
  InstBranch(ValuePtr_t cond, BasicBlockPtr_t true_block, BasicBlockPtr_t false_block, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Br, TypeHelper::TYPE_VOID, parent_bb)
  {
    this->add_operand(cond);
    this->add_operand(true_block);
    this->add_operand(false_block);
  }
  static sptr(InstBranch)
    New(ValuePtr_t cond, BasicBlockPtr_t true_block, BasicBlockPtr_t false_block, BasicBlockPtr_t parent_bb)
  {
    auto inst      = std::make_shared<InstBranch>(cond, true_block, false_block, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
  // unconditional branch
  InstBranch(BasicBlockPtr_t target_block, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Br, TypeHelper::TYPE_VOID, parent_bb)
  {
    this->add_operand(target_block);
  }
  static sptr(InstBranch) New(BasicBlockPtr_t target_block, BasicBlockPtr_t parent_bb)
  {
    auto inst      = std::make_shared<InstBranch>(target_block, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

class InstReturn : public Instruction
{
public:
  InstReturn(ValuePtr_t ret_val, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Ret, TypeHelper::TYPE_VOID, parent_bb)
  {
    this->add_operand(ret_val);
  }
  InstReturn(BasicBlockPtr_t parent_bb) : Instruction(InstTag::Ret, TypeHelper::TYPE_VOID, parent_bb) {}
  static sptr(InstReturn) New(ValuePtr_t ret_val, BasicBlockPtr_t parent_bb)
  {
    auto inst      = std::make_shared<InstReturn>(ret_val, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
  static sptr(InstReturn) New(BasicBlockPtr_t parent_bb)
  {
    auto inst      = std::make_shared<InstReturn>(parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

// call's type is the function's return type
class InstCall : public Instruction
{
public:
  InstCall(FunctionPtr_t func, const std::vector<ValuePtr_t> &args, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Call, func->get_type()->return_type, parent_bb)
  {
    add_operand(func);
    for (auto arg : args)
      {
        add_operand(arg);
      }
  }
  virtual std::string to_IR_string() override
  {
    if (Type::isType<VoidType>(type))
      {
        panic("No ret val");
      }
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static std::shared_ptr<InstCall> New(
    FunctionPtr_t func,
    const std::vector<ValuePtr_t> &args,
    BasicBlockPtr_t parent_bb
  )
  {
    auto inst      = std::make_shared<InstCall>(func, args, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }
};

// getelementptr <ty>, ptr <ptrval>{, [inrange] <ty> <idx>}*
// Example: lea arr[2][3] from arr[5][4]:
//  &arr[2][3] = GEP [5x[4xi32]], [5x[4xi32]]* arr, i32 0, i32 2, i32 3
// For simplicity, we want to limit the use of gep of at most 3 ops
// op0: pointer to array;
// if there is only one index, the return value has the same type as op0
// if there is 2 index, the return value has type of arr's element
class InstGEP : public Instruction
{
public:
  ValuePtr_t aim_to;
  TypePtr_t element_type;
  InstGEP(ValuePtr_t pointer, const std::vector<ValuePtr_t> &indices, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::GEP, std::make_shared<PointerType>(extract_type(pointer, indices)), parent_bb)
  {
    if (shared_cast<InstGEP>(pointer))
      {
        aim_to = std::dynamic_pointer_cast<InstGEP>(pointer)->aim_to;
      }
    else if (shared_cast<InstAlloca>(pointer) || shared_cast<GlobalVar>(pointer))
      {
        aim_to = pointer;
      }
    else
      {
        // LOG(WARNING) << "Unexpected pointer type " << pointer->to_string();
        aim_to = nullptr;
      }
    assert(indices.size() <= 2);
    element_type = extract_type(pointer, indices);
    add_operand(pointer);
    for (auto index : indices)
      {
        add_operand(index);
      }
  }

  virtual std::string to_IR_string() override
  {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static std::shared_ptr<InstGEP> New(
    ValuePtr_t pointer,
    const std::vector<ValuePtr_t> &indices,
    BasicBlockPtr_t parent_bb
  )
  {
    auto inst      = std::make_shared<InstGEP>(pointer, indices, parent_bb);
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.end(), inst);
    return inst;
  }

  // get the inner
  static TypePtr_t extract_type(ValuePtr_t pointer, const std::vector<ValuePtr_t> &indices)
  {
    auto pointed_type = shared_cast<PointerType>(pointer->type)->pointed_type;
    if (Type::isType<IntegerType>(pointed_type))
      {
        return pointed_type;
      }
    else if (Type::isType<ArrayType>(pointed_type))
      {
        for (int i = 1; i < indices.size(); ++i)
          {
            pointed_type = shared_cast<ArrayType>(pointed_type)->elem_type;
          }
        return pointed_type;
      }
    LOG(ERROR) << "unexpected type: " << pointer->to_string();
    assert(0);
  }
};

class InstPhi : public Instruction
{
public:
  InstPhi(TypePtr_t type, const decltype(Function::bb_list) &incoming_vals, BasicBlockPtr_t parent_bb)
    : Instruction(InstTag::Phi, type, parent_bb)
  {
    operand_list.resize(incoming_vals.size(), nullptr);
  }
  void set_incoming_val(unsigned index, ValuePtr_t val)
  {
    auto old_op         = operand_list[index];
    operand_list[index] = val;
    if (val) val->use_list.push_back({/*val.get(),*/ this, index});
    if (old_op)
      {
        if (std::find(operand_list.begin(), operand_list.end(), old_op) == operand_list.end())
          {
            old_op->u_remove_use(this);
          }
      }
  }
  virtual std::string to_IR_string() override
  {
    std::string str = type->to_IR_string() + " %" + std::to_string(ir_seqno);
    return str;
  }
  static InstPhiPtr_t New(
    TypePtr_t type,
    const decltype(Function::bb_list) &incoming_vals,
    BasicBlockPtr_t parent_bb,
    const std::string &name
  )
  {
    auto inst      = std::make_shared<InstPhi>(type, incoming_vals, parent_bb);
    inst->name     = name;
    inst->inst_itr = parent_bb->inst_list.insert(parent_bb->inst_list.begin(), inst);
    return inst;
  }
};

} // namespace CompSysY