#pragma once

#include <memory>
#include <shared_mutex>
#include <string>
#include <vector>

#include "common.h"
#include "llir_type.h"

namespace antlrSysY {

class Value {
public:
  std::string name;
  TypePtr_t type;
  Value(const std::string &name, TypePtr_t type) : name(name), type(type) {}
  virtual ~Value() = default;
};

class User : public Value {
public:
};

class BasicBlock;
class Instruction;

class FParam : public Value {
public:
  FParam(const std::string &name, TypePtr_t type) : Value(name, type) {}
};

class Function : public Value {
public:
  std::vector<std::shared_ptr<FParam>> fparam_list;
  std::vector<std::shared_ptr<BasicBlock>> bb_list;
  Function(const std::string &name, TypePtr_t ret_type) : Value(name, std::make_shared<FunctionType>(ret_type)) {}
  std::shared_ptr<FunctionType> get_type() {
    return std::dynamic_pointer_cast<FunctionType>(type);
  }
};

class BasicBlock : public Value {
public:
  std::vector<std::shared_ptr<Instruction>> inst_list;
  std::shared_ptr<Function> parent;
  std::vector<std::shared_ptr<BasicBlock>> successors;
  std::vector<std::shared_ptr<BasicBlock>> predecessors;
  BasicBlock(const std::string &name, std::shared_ptr<Function> parent) : Value(name, TypeHelper::TYPE_LABEL) {
    this->parent = parent;
  }
};

typedef std::shared_ptr<Value> ValuePtr_t;
typedef std::shared_ptr<BasicBlock> BasicBlockPtr_t;
typedef std::shared_ptr<Function> FunctionPtr_t;

class Constant : public Value {
public:
  Constant(const std::string &name, TypePtr_t type) : Value(name, type) {}
};

class ConstantInt : public Constant {
public:
  int value;
  ConstantInt(const std::string &name, int value) : Constant(name, TypeHelper::TYPE_I32), value(value) {}
  static std::shared_ptr<ConstantInt> make_shared(int value) {
    return std::make_shared<ConstantInt>("", value);
  }
};

class ConstantArr : public Constant {
public:
  std::vector<ValuePtr_t> value_list;
  ConstantArr(const std::string &name, const std::vector<ValuePtr_t> &value_list, std::shared_ptr<ArrayType> type)
      : Constant(name, type), value_list(value_list) {}
  static std::shared_ptr<ConstantArr> make_shared(const std::string &name, const std::vector<ValuePtr_t> &value_list,
                                                  std::shared_ptr<ArrayType> type) {
    return std::make_shared<ConstantArr>(name, value_list, type);
  }
};

typedef std::shared_ptr<Constant> ConstantPtr_t;

class GlobalVar : public Value {
public:
  ConstantPtr_t init_value;
  bool is_const;
  GlobalVar(const std::string &name, ConstantPtr_t init_value, bool is_const)
      : Value(name, init_value->type), init_value(init_value), is_const(is_const) {}
  static std::shared_ptr<GlobalVar> make_shared(const std::string &name, ConstantPtr_t init_value, bool is_const) {
    return std::make_shared<GlobalVar>(name, init_value, is_const);
  }
};

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 Value {
public:
  InstTag tag;
  std::shared_ptr<BasicBlock> parent_bb;
  std::vector<std::shared_ptr<Value>> operand_list;
  Instruction(InstTag inst_tag, TypePtr_t type, std::shared_ptr<BasicBlock> parent_bb)
      : Value("", type), tag(inst_tag), parent_bb(parent_bb) {}
};

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

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) {
    operand_list.push_back(value);
    operand_list.push_back(pointer);
  }
};

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) {
    operand_list.push_back(value);
  }
};

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) {
    operand_list.push_back(op1);
    operand_list.push_back(op2);
  }
};

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) {
    operand_list.push_back(op);
  }
};

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->operand_list.push_back(cond);
    this->operand_list.push_back(true_block);
    this->operand_list.push_back(false_block);
  }
  // unconditional branch
  InstBranch(BasicBlockPtr_t target_block, BasicBlockPtr_t parent_bb)
      : Instruction(InstTag::Br, TypeHelper::TYPE_VOID, parent_bb) {
    this->operand_list.push_back(target_block);
  }
};

class InstReturn : public Instruction {
public:
  InstReturn(ValuePtr_t ret_val, BasicBlockPtr_t parent_bb)
      : Instruction(InstTag::Ret, TypeHelper::TYPE_VOID, parent_bb) {
    this->operand_list.push_back(ret_val);
  }
  InstReturn(BasicBlockPtr_t parent_bb) : Instruction(InstTag::Ret, TypeHelper::TYPE_VOID, parent_bb) {}
};

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) {
    operand_list.push_back(func);
    operand_list.insert(operand_list.end(), args.begin(), args.end());
  }
};

// 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
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 (typeid(pointer) == typeid(std::shared_ptr<InstGEP>)) {
      aim_to = std::dynamic_pointer_cast<InstGEP>(pointer)->aim_to;
    }
    else {
      sysy_assert(typeid(pointer) == typeid(std::shared_ptr<InstAlloca>) ||
                  typeid(pointer) == typeid(std::shared_ptr<GlobalVar>));
      aim_to = pointer;
    }
    element_type = extract_type(pointer, indices);
    operand_list.push_back(pointer);
    operand_list.insert(operand_list.end(), indices.begin(), indices.end());
  }

  // get the inner
  static TypePtr_t extract_type(ValuePtr_t pointer, const std::vector<ValuePtr_t> &indices) {
    auto pointer_type = std::dynamic_pointer_cast<PointerType>(pointer->type);
    auto pointed_type = pointer_type->pointed_type;
    if (pointed_type->type_tag == Type::TypeTag::ArrayType) {
      for (int i = 1; i < indices.size() - 1; ++i) {
        pointed_type = std::dynamic_pointer_cast<ArrayType>(pointed_type)->element_type;
      }
    }
    if (pointed_type->type_tag == Type::TypeTag::IntegerType) {
      return pointed_type;
    }
    else {
      panic("Unmatch indices and array pointer");
    }
  }
};

} // namespace antlrSysY