#pragma once

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

#include "3rdparty/easylogging++.h"
#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;
  template <typename TT> static bool isType(std::shared_ptr<Value> ptr) {
    if (ptr.get()) {
      auto &r = *ptr.get();
      return typeid(r) == typeid(TT);
    }
    LOG(WARNING) << "Comparing with nullptr";
    return false;
  }
  virtual std::string to_string() {
    return name + ": " + type->to_string();
  }
};

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;
  // we use empty basic_block to represent lib function
  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);
  }
  bool is_libfunc() {
    return bb_list.size() == 0;
  }
  virtual std::string to_string() override {
    std::string params_str;
    for (auto fparam : fparam_list) {
      params_str += fparam->to_string() + ",";
    }
    return name + "(" + params_str + ") -> " + type->to_string();
  }
};

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;

// emmm, actually it is more of a Instruction type
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, PointerType::make_shared(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) {}
  virtual std::string to_string() override {
    switch (tag) {
      case InstTag::Add:
        return "Add";
      case InstTag::Sub:
        return "Sub";
      case InstTag::Mod:
        return "Mod";
      case InstTag::Mul:
        return "Mul";
      case InstTag::Div:
        return "Div";
      case InstTag::Lt:
        return "Lt";
      case InstTag::Le:
        return "Le";
      case InstTag::Ge:
        return "Ge";
      case InstTag::Gt:
        return "Gt";
      case InstTag::Eq:
        return "Eq";
      case InstTag::Ne:
        return "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";
    }
  }
};

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 (Value::isType<InstGEP>(pointer)) {
      aim_to = std::dynamic_pointer_cast<InstGEP>(pointer)->aim_to;
    }
    else if (Value::isType<InstAlloca>(pointer) || Value::isType<GlobalVar>(pointer)){
      aim_to = pointer;
    }
    else {
      LOG(WARNING) << "Unexpected pointer type" << pointer->to_string();
      aim_to = pointer;
      panic("message");
    }
    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;
    // ptr->array
    if (pointed_type->type_tag == Type::TypeTag::ArrayType) {
      for (int i = 1; i < indices.size(); ++i) {
        pointed_type = std::dynamic_pointer_cast<ArrayType>(pointed_type)->element_type;
      }
    }
    if (pointed_type->type_tag == Type::TypeTag::IntegerType) {
      return pointed_type;
    }
    else {
      LOG(WARNING) << "not returning an int-type";
      return pointed_type;
    }
  }
};

} // namespace antlrSysY