#include "3rdparty/easylogging++.h"
#include "common.h"
#include "llir.h"
#include "visitor.h"
#include <array>
#include <memory>
#include <ostream>
#include <sstream>
#include <string>
#include <vector>

namespace antlrSysY {

const std::array<std::string, 8> libfunc_list =
  {"getint", "getch", "getarray", "putint", "putch", "putarray", "starttime", "stoptime"};

// On generation, the array is flattened, thus we have to calculate them
std::shared_ptr<ConstantArr> gen_arr_hierarchy(
  const std::shared_ptr<ArrayType> array_type,
  const std::vector<ValuePtr_t> &const_array,
  int base,
  int length
) {
  int dim_n    = array_type->element_count;
  int dim_size = length / dim_n;
  std::vector<ValuePtr_t> value_list;
  if (array_type->element_type->type_tag == Type::TypeTag::IntegerType) {
    sysy_assert(dim_size == 1);
    for (int i = 0; i < dim_n; ++i) {
      if (const_array[base + i]) {
        value_list.push_back(const_array[base + i]);
      }
      else {
        break;
      }
    }
    if (!value_list.empty())
      return ConstantArr::make_shared("", value_list, array_type);
    else
      return nullptr;
  }
  else {
    int last_non_null = -1;
    for (int i = 0; i < dim_n; ++i) {
      auto sub_arr = gen_arr_hierarchy(
        std::dynamic_pointer_cast<ArrayType>(array_type->element_type), const_array, dim_size * i, dim_size
      );
      value_list.push_back(sub_arr);
      if (sub_arr) last_non_null = i;
    }
    if (last_non_null == -1) {
      return nullptr;
    }
    else {
      // if (last_non_null + 1 < value_list.size())
      value_list.erase(value_list.begin() + last_non_null + 1, value_list.end());
      return ConstantArr::make_shared("", value_list, array_type);
    }
  }
}

static void _build_arr_init_list(
  std::ostream &ostr,
  std::shared_ptr<ConstantArr> arr,
  std::shared_ptr<ArrayType> arr_type
) {
  if (arr_type->element_type->type_tag == Type::TypeTag::IntegerType) {
    ostr << arr_type->to_IR_string();
    ostr << " [";
    for (int i = 0; i < arr->value_list.size(); ++i) {
      ostr << "i32 " << std::dynamic_pointer_cast<ConstantInt>(arr->value_list[i])->value;
      if (i < arr->value_list.size() - 1) {
        ostr << ", ";
      }
    }
    for (int i = 0; i < arr_type->element_count - arr->value_list.size(); ++i) {
      ostr << ", i32 0";
    }
    ostr << "]";
  }
  else {
    ostr << arr_type->to_IR_string() << " ";
    ostr << "[";
    for (int i = 0; i < arr->value_list.size(); ++i) {
      if (arr->value_list[i])
        _build_arr_init_list(
          ostr,
          std::dynamic_pointer_cast<ConstantArr>(arr->value_list[i]),
          std::dynamic_pointer_cast<ArrayType>(arr_type->element_type)
        );
      else
        ostr << arr_type->element_type->to_IR_string() << " zeroinitializer";
      if (i < arr->value_list.size() - 1) {
        ostr << ", ";
      }
    }
    for (int i = 0; i < arr_type->element_count - arr->value_list.size(); ++i) {
      ostr << ", " << arr_type->element_type->to_IR_string() << " zeroinitializer";
    }
    ostr << "]";
  }
}

static void _gen_blocks(std::ostream &ostr, const std::list<BasicBlockPtr_t> &block_list, int reg_count) {
  // first pass, allocate sequence number
  for (auto block_itr = block_list.begin(); block_itr != block_list.end(); ++block_itr) {
    auto block = *block_itr;
    sysy_assert(block->ir_seqno == -1); // multi-alloc is error
    block->ir_seqno = reg_count++;
    for (auto inst : block->inst_list) {
      sysy_assert(inst->ir_seqno == -1); // multi-alloc is error
      switch (inst->tag) {
        // These are not to get a seqno
        case InstTag::Br:
        case InstTag::Ret:
        case InstTag::Store: break;
        // Call's seqno is dependent on its return type
        case InstTag::Call:
          if (inst->type->type_tag == Type::TypeTag::IntegerType) {
            inst->ir_seqno = reg_count++;
          }
          break;
        // These should have a seqno
        case InstTag::Add:
        case InstTag::Sub:
        case InstTag::Mod:
        case InstTag::Mul:
        case InstTag::Div:
        case InstTag::Lt:
        case InstTag::Le:
        case InstTag::Ge:
        case InstTag::Gt:
        case InstTag::Eq:
        case InstTag::Ne:
        case InstTag::And:
        case InstTag::Or:
        case InstTag::Load:
        case InstTag::GEP:
        case InstTag::Alloca:
        case InstTag::Phi:
        case InstTag::Zext: inst->ir_seqno = reg_count++; break;
        // These should not be generated in frontend stage
        default: panic("Unexpected instruction");
      }
    }
  }
  // second pass, generate IR
  for (auto block_itr = block_list.begin(); block_itr != block_list.end(); ++block_itr) {
    auto block = *block_itr;
    if (block_itr != block_list.begin()) {
      ostr << block->ir_seqno << ":" << std::endl;
    }
    for (auto _inst : block->inst_list) {
      ostr << "  ";
      VLOG(5) << "Build inst" << _inst->ir_seqno << ": " << _inst->to_string();
      switch (_inst->tag) {
        case InstTag::Br: {
          auto inst = Value::as<InstBranch>(_inst);
          assert(inst->ir_seqno == -1);
          ostr << "br ";
          if (inst->operand_list.size() == 1) {
            assert(Value::is<BasicBlock>(inst->operand_list[0]));
            auto bb_dest = Value::as<BasicBlock>(inst->operand_list[0]);
            assert(bb_dest->ir_seqno >= 0);
            ostr << "label %" << bb_dest->ir_seqno;
          }
          else {
            assert(Value::is<Instruction>(inst->operand_list[0]));
            assert(Type::isType<IntegerType>(inst->operand_list[0]->type));
            auto cond     = Value::as<Instruction>(inst->operand_list[0]);
            auto bb_true  = Value::as<BasicBlock>(inst->operand_list[1]);
            auto bb_false = Value::as<BasicBlock>(inst->operand_list[2]);
            if (!Type::asType<IntegerType>(cond->type)->isI1()) {
              LOG(ERROR) << "Expect cond evals to i1: " << cond->to_string();
              panic("Grammar check");
            }
            assert(cond->ir_seqno >= 0);
            assert(bb_true->ir_seqno >= 0);
            assert(bb_false->ir_seqno >= 0);
            ostr << "i1 %" << cond->ir_seqno << ", label %" << bb_true->ir_seqno << ", label %" << bb_false->ir_seqno;
          }
          break;
        }
        case InstTag::Ret: {
          auto inst = Value::as<InstReturn>(_inst);
          assert(inst->ir_seqno == -1);
          ostr << "ret ";
          if (inst->operand_list.size() == 0) {
            assert(Type::isType<VoidType>(inst->type));
            ostr << "void";
          }
          else {
            if (Value::is<Instruction>(inst->operand_list[0])) {
              auto op0 = Value::as<Instruction>(inst->operand_list[0]);
              assert(op0->ir_seqno >= 0);
              ostr << op0->type->to_IR_string() << " %" << op0->ir_seqno;
            }
            else if (Value::is<ConstantInt>(inst->operand_list[0])) {
              auto op0 = Value::as<ConstantInt>(inst->operand_list[0]);
              ostr << op0->to_IR_string();
            }
            else if (Value::is<FParam>(inst->operand_list[0])) {
              auto op0 = Value::as<FParam>(inst->operand_list[0]);
              ostr << op0->to_IR_string();
            }
            else {
              LOG(ERROR) << "Unexpected type of op0: " << inst->operand_list[0]->to_string();
              assert(0);
            }
          }
          break;
        }
        case InstTag::Store: {
          auto inst = Value::as<InstStore>(_inst);
          assert(inst->ir_seqno == -1);
          ostr << "store ";

          if (Value::is<Instruction>(inst->operand_list[0])) {
            auto op0 = Value::as<Instruction>(inst->operand_list[0]);
            assert(op0->ir_seqno >= 0);
            ostr << op0->type->to_IR_string() << " %" << op0->ir_seqno << ", ";
          }
          else if (Value::is<ConstantInt>(inst->operand_list[0])) {
            auto op0 = Value::as<ConstantInt>(inst->operand_list[0]);
            ostr << op0->to_IR_string() << ", ";
          }
          else if (Value::is<FParam>(inst->operand_list[0])) {
            auto op0 = Value::as<FParam>(inst->operand_list[0]);
            ostr << op0->to_IR_string() << ", ";
          }
          else {
            LOG(ERROR) << "Unexpected type of op0: " << inst->operand_list[0]->to_string();
            assert(0);
          }
          if (Value::is<GlobalVar>(inst->operand_list[1])) {
            auto op1 = Value::as<GlobalVar>(inst->operand_list[1]);
            ostr << op1->type->to_IR_string() << " @" << op1->name;
          }
          else if (Value::is<Instruction>(inst->operand_list[1])) {
            auto op1 = Value::as<Instruction>(inst->operand_list[1]);
            assert(op1->ir_seqno >= 0);
            ostr << op1->type->to_IR_string() << " %" << op1->ir_seqno;
          }
          else if (Value::is<FParam>(inst->operand_list[0])) {
            auto op0 = Value::as<FParam>(inst->operand_list[0]);
            ostr << op0->to_IR_string();
          }
          else {
            LOG(ERROR) << "Unexpected type of op1: " << inst->operand_list[1]->to_string();
            assert(0);
          }
          ostr << ", align 4";
          break;
        }
        // Call's seqno is dependent on its return type
        case InstTag::Call: {
          auto inst = Value::as<InstCall>(_inst);
          LOG(DEBUG) << inst->operand_list[0]->to_string();
          auto func      = Value::as<Function>(inst->operand_list[0]);
          auto func_type = Type::asType<FunctionType>(func->type);
          if (Type::isType<VoidType>(func_type->return_type)) {
            assert(inst->ir_seqno == -1);
            ostr << "call void @";
          }
          else {
            assert(inst->ir_seqno >= 0);
            ostr << "%" << inst->ir_seqno << " = call i32 @";
          }
          ostr << func->name << "(";
          if (inst->operand_list.size() > 1) {
            for (int i = 1; i < inst->operand_list.size(); ++i) {
              if (Value::is<Instruction>(inst->operand_list[i])) {
                auto op0 = Value::as<Instruction>(inst->operand_list[i]);
                assert(op0->ir_seqno >= 0);
                ostr << op0->type->to_IR_string() << " %" << op0->ir_seqno;
              }
              else if (Value::is<ConstantInt>(inst->operand_list[i])) {
                auto op0 = Value::as<ConstantInt>(inst->operand_list[i]);
                ostr << op0->to_IR_string();
              }
              else if (Value::is<FParam>(inst->operand_list[i])) {
                auto op0 = Value::as<FParam>(inst->operand_list[i]);
                ostr << op0->to_IR_string() << ", ";
              }
              else {
                LOG(ERROR) << "Unexpected type of op_i: " << inst->operand_list[i]->to_string();
                assert(0);
              }
              if (i < inst->operand_list.size() - 1) ostr << ", ";
            }
          }
          ostr << ")";
          break;
        }
        // These should have a seqno
        case InstTag::Add:
        case InstTag::Sub:
        case InstTag::Mod:
        case InstTag::Mul:
        case InstTag::Div:
        case InstTag::Lt:
        case InstTag::Le:
        case InstTag::Ge:
        case InstTag::Gt:
        case InstTag::Eq:
        case InstTag::Ne:
        case InstTag::And:
        case InstTag::Or: {
          auto inst = Value::as<InstBinary>(_inst);
          ostr << "%" << inst->ir_seqno << " = " << inst->to_string() << " ";
          if (Value::is<Instruction>(inst->operand_list[0])) {
            auto op0 = Value::as<Instruction>(inst->operand_list[0]);
            assert(op0->ir_seqno >= 0);
            ostr << op0->type->to_IR_string() << " %" << op0->ir_seqno << ", ";
          }
          else if (Value::is<ConstantInt>(inst->operand_list[0])) {
            auto op0 = Value::as<ConstantInt>(inst->operand_list[0]);
            ostr << op0->to_IR_string() << ", ";
          }
          else if (Value::is<FParam>(inst->operand_list[0])) {
            auto op0 = Value::as<FParam>(inst->operand_list[0]);
            ostr << op0->to_IR_string() << ", ";
          }
          else {
            LOG(ERROR) << "Unexpected type of op0: " << inst->operand_list[0]->to_string();
            assert(0);
          }
          if (Value::is<Instruction>(inst->operand_list[1])) {
            auto op1 = Value::as<Instruction>(inst->operand_list[1]);
            assert(op1->ir_seqno >= 0);
            ostr << "%" << op1->ir_seqno;
          }
          else if (Value::is<ConstantInt>(inst->operand_list[1])) {
            auto op1 = Value::as<ConstantInt>(inst->operand_list[1]);
            ostr << op1->value;
          }
          else if (Value::is<FParam>(inst->operand_list[1])) {
            auto op1 = Value::as<FParam>(inst->operand_list[1]);
            ostr << "%" << op1->ir_seqno;
          }
          else {
            LOG(ERROR) << "Unexpected type of op1: " << inst->operand_list[1]->to_string();
            assert(0);
          }
          break;
        }
        case InstTag::Load: {
          auto inst = Value::as<InstLoad>(_inst);
          assert(inst->ir_seqno != -1);
          assert(Type::isType<PointerType>(inst->operand_list[0]->type));
          ostr << "%" << inst->ir_seqno << " = load " << inst->type->to_IR_string() << ", ";
          if (Value::is<GlobalVar>(inst->operand_list[0])) {
            auto op = Value::as<GlobalVar>(inst->operand_list[0]);
            ostr << op->type->to_IR_string() << " @" << op->name;
          }
          else if (Value::is<Instruction>(inst->operand_list[0])) {
            auto op = Value::as<Instruction>(inst->operand_list[0]);
            assert(op->ir_seqno >= 0);
            ostr << op->type->to_IR_string() << " %" << op->ir_seqno;
          }
          else if (Value::is<FParam>(inst->operand_list[0])) {
            auto op0 = Value::as<FParam>(inst->operand_list[0]);
            ostr << op0->to_IR_string();
          }
          else {
            LOG(ERROR) << "Unexpected type of op0: " << inst->operand_list[1]->to_string();
            assert(0);
          }
          break;
        }
        case InstTag::GEP: {
          auto inst = Value::as<InstGEP>(_inst);

          assert(inst->ir_seqno >= 0);
          auto pointer_type = Type::asType<PointerType>(inst->operand_list[0]->type);
          ostr << "%" << inst->ir_seqno << " = getelementptr " << pointer_type->pointed_type->to_IR_string() << ", ";
          if (Value::is<GlobalVar>(inst->operand_list[0])) {
            auto op0 = Value::as<GlobalVar>(inst->operand_list[0]);
            ostr << op0->to_IR_string();
          }
          else if (Value::is<Instruction>(inst->operand_list[0])) {
            auto op0 = Value::as<Instruction>(inst->operand_list[0]);
            ostr << op0->type->to_IR_string() << " %" << op0->ir_seqno;
          }
          else if (Value::is<FParam>(inst->operand_list[0])) {
            auto op0 = Value::as<FParam>(inst->operand_list[0]);
            ostr << op0->to_IR_string();
          }
          else {
            LOG(WARNING) << "Unexpected type of op0: " << inst->operand_list[0]->to_string();
            assert(0);
          }
          for (int i = 1; i < inst->operand_list.size(); ++i) {
            ostr << ", " << inst->operand_list[i]->to_IR_string();
          }
          break;
        }
        case InstTag::Alloca: {
          auto inst = Value::as<InstAlloca>(_inst);
          assert(inst->ir_seqno != -1);
          auto pointer_type = Type::asType<PointerType>(inst->type);
          ostr << "%" << inst->ir_seqno << " = alloca " << pointer_type->pointed_type->to_IR_string() << ", align 4";
          ostr << "    ;" << inst->name;
          break;
        }
        case InstTag::Zext: {
          auto inst = Value::as<InstZext>(_inst);
          assert(inst->ir_seqno >= 0);
          ostr << "%" << inst->ir_seqno << " = zext ";
          if (Value::is<Instruction>(inst->operand_list[0])) {
            auto op0 = Value::as<Instruction>(inst->operand_list[0]);
            assert(op0->ir_seqno >= 0);
            ostr << op0->type->to_IR_string() << " %" << op0->ir_seqno << "to ";
          }
          else if (Value::is<ConstantInt>(inst->operand_list[0])) {
            auto op0 = Value::as<ConstantInt>(inst->operand_list[0]);
            ostr << op0->to_IR_string() << "to ";
          }
          else if (Value::is<FParam>(inst->operand_list[0])) {
            auto op0 = Value::as<FParam>(inst->operand_list[0]);
            ostr << op0->to_IR_string() << "to ";
          }
          else {
            LOG(ERROR) << "Unexpected type of op0: " << inst->operand_list[0]->to_string();
            assert(0);
          }
          ostr << "i32";
          break;
        }
        case InstTag::Phi: {
          auto inst = Value::as<InstPhi>(_inst);
          assert(inst->ir_seqno >= 0);
          ostr << "%" << inst->ir_seqno << " = phi " << inst->type->to_IR_string() << " ";
          for (int i = 0; i < inst->operand_list.size(); ++i) {
            auto op = inst->operand_list[i];
            ostr << "[";
            if (Value::is<GlobalVar>(op)) {
              auto op0 = Value::as<GlobalVar>(op);
              ostr << "@" << op0->name;
            }
            else if (Value::is<Instruction>(op)) {
              auto op0 = Value::as<Instruction>(op);
              ostr << "%" << op0->ir_seqno;
            }
            else if (Value::is<FParam>(op)) {
              auto op0 = Value::as<FParam>(op);
              ostr << "%" << op0->ir_seqno;
            }
            else if (Value::is<ConstantInt>(op)) {
              auto op0 = Value::as<ConstantInt>(op);
              ostr << op0->value;
            }
            else {
              LOG(WARNING) << "Unexpected type of op: " << op->to_string();
              assert(0);
            }
            auto pred = inst->parent_bb->predecessors.begin();
            std::advance(pred, i);
            ostr << ", %" << (*pred)->ir_seqno << "]";
            if (i < inst->operand_list.size() - 1) {
              ostr << ", ";
            }
          }
          break;
        }
        // These should not be generated in frontend stage
        default: panic("Unexpected instruction");
      }
      ostr << "\n";
    }
  }
}

void Visitor::llir_gen(std::ostream &ostr) {
#pragma region GenLibFuncDecl
  for (auto &lib_func_name : libfunc_list) {
    VLOG(6) << "Gen LibFunc " << lib_func_name;
    auto lib_func      = _func_tab.get_name(lib_func_name).value();
    auto lib_func_type = std::dynamic_pointer_cast<FunctionType>(lib_func->type);
    ostr << "declare"
         << " " << lib_func_type->return_type->to_IR_string() << " "
         << "@" << lib_func_name << "(";
    auto &param_list = lib_func->fparam_list;
    for (int i = 0; i < (int)param_list.size() - 1; ++i) {
      ostr << param_list[i]->type->to_IR_string() << ", ";
    }
    if (param_list.size()) ostr << param_list.back()->type->to_IR_string();
    ostr << ")" << std::endl;
  }
  // ostr.flush();
#pragma endregion

#pragma region GenGlobDecl
  for (auto &global_var : module.global_var_list) {
    // both int and arr have the same leading part
    VLOG(5) << "Gen Global " << global_var->name;
    ostr << "@" << global_var->name << " = "
         << "dso_local"
         << " " << (global_var->is_const ? "constant" : "global") << " ";
    auto global_var_type = std::dynamic_pointer_cast<PointerType>(global_var->type);
    if (global_var_type->pointed_type->type_tag == Type::TypeTag::IntegerType) {
      auto init_value = std::dynamic_pointer_cast<ConstantInt>(global_var->init_value);
      ostr << "i32"
           << " " << init_value->value << std::endl;
    }
    else if (global_var_type->pointed_type->type_tag == Type::TypeTag::ArrayType) {
      // sysy_assert(global_var->)
      auto array_type = std::dynamic_pointer_cast<ArrayType>(global_var_type->pointed_type);
      auto init_value = std::dynamic_pointer_cast<ConstantArr>(global_var->init_value);
      if (init_value != nullptr) {
        auto hierarchy_array = gen_arr_hierarchy(array_type, init_value->value_list, 0, init_value->value_list.size());
        if (hierarchy_array) {
          // LOG(DEBUG) << hierarchy_array->to_string();
          _build_arr_init_list(ostr, hierarchy_array, array_type);
        }
        else {
          ostr << array_type->to_IR_string() << " zeroinitializer";
        }
      }
      else {
        ostr << array_type->to_IR_string() << " zeroinitializer";
      }
      ostr << ", align 4" << std::endl;
    }
    else {
      LOG(ERROR) << global_var->to_string();
      panic("Invalid Global Declaration");
    }
  }
#pragma endregion

#pragma region GenFunction
  for (auto &func : module.function_list) {
    if (func->is_libfunc()) {
      LOG(WARNING) << "Lib func";
      continue;
    }
    int reg_count  = 0; // allocate unique register name from %0
    auto func_type = std::dynamic_pointer_cast<FunctionType>(func->type);
    ostr << "define dso_local " << func_type->return_type->to_IR_string() << " @" << func->name;
    ostr << "(";
    auto &param_list = func->fparam_list;
    for (int i = 0; i < (int)param_list.size() - 1; ++i) {
      param_list[i]->ir_seqno = reg_count++;
      ostr << param_list[i]->to_IR_string() << ", ";
    }
    if (param_list.size()) {
      param_list.back()->ir_seqno = reg_count++;
      ostr << param_list.back()->to_IR_string();
    }
    ostr << ") {" << std::endl;
    _gen_blocks(ostr, func->bb_list, reg_count);
    ostr << "}" << std::endl;
  }
#pragma endregion
}

} // namespace antlrSysY