#pragma once

#include "IR.h"
#include "IRBuilder.h"
#include "SysYBaseVisitor.h"
#include "SysYParser.h"
#include <memory>
#include <unordered_map>
#include <forward_list>

namespace sysy {

class SymbolTable{
private:
  enum Kind
  {
    kModule,
    kFunction,
    kBlock,
  };

  std::forward_list<std::pair<Kind, std::unordered_map<std::string, Value*>>> Scopes;

public:
  struct ModuleScope {
    SymbolTable& tables_ref;
    ModuleScope(SymbolTable& tables) : tables_ref(tables) {
      tables.enter(kModule);
    }
    ~ModuleScope() { tables_ref.exit(); }
  };
  struct FunctionScope {
    SymbolTable& tables_ref;
    FunctionScope(SymbolTable& tables) : tables_ref(tables) {
      tables.enter(kFunction);
    }
    ~FunctionScope() { tables_ref.exit(); }
  };
  struct BlockScope {
    SymbolTable& tables_ref;
    BlockScope(SymbolTable& tables) : tables_ref(tables) {
      tables.enter(kBlock);
    }
    ~BlockScope() { tables_ref.exit(); }
  };

  SymbolTable() = default;

  bool isModuleScope() const { return Scopes.front().first == kModule; }
  bool isFunctionScope() const { return Scopes.front().first == kFunction; }
  bool isBlockScope() const { return Scopes.front().first == kBlock; }
  Value *lookup(const std::string &name) const {
    for (auto &scope : Scopes) {
      auto iter = scope.second.find(name);
      if (iter != scope.second.end())
        return iter->second;
    }
    return nullptr;
  }
  auto insert(const std::string &name, Value *value) {
    assert(not Scopes.empty());
    return Scopes.front().second.emplace(name, value);
  }
private:
  void enter(Kind kind) {
    Scopes.emplace_front();
    Scopes.front().first = kind;
  }
  void exit() {
    Scopes.pop_front(); 
  }
  
};

class SysYIRGenerator : public SysYBaseVisitor {
private:
  std::unique_ptr<Module> module;
  IRBuilder builder;
  SymbolTable symbols_table;

  int trueBlockNum = 0, falseBlockNum = 0;

  int d = 0, n = 0;
  vector<int> path;
  bool isalloca;
  AllocaInst *current_alloca;
  GlobalValue *current_global;
  Type* current_type;
  int numdims = 0;

public:
  SysYIRGenerator() = default;

public:
  Module *get() const { return module.get(); }

public:
  std::any visitCompUnit(SysYParser::CompUnitContext *ctx) override;
  std::any visitDecl(SysYParser::DeclContext *ctx) override;
  std::any visitConstDecl(SysYParser::ConstDeclContext *ctx) override;
  std::any visitBType(SysYParser::BTypeContext *ctx) override;
  std::any visitConstDef(SysYParser::ConstDefContext *ctx) override;
  std::any visitConstInitVal(SysYParser::ConstInitValContext *ctx) override;
  std::any visitFuncType(SysYParser::FuncTypeContext* ctx) override;
  std::any visitFuncDef(SysYParser::FuncDefContext* ctx) override;
  std::any visitVarDecl(SysYParser::VarDeclContext *ctx) override;
  std::any visitVarDef(SysYParser::VarDefContext *ctx) override;
  std::any visitInitVal(SysYParser::InitValContext *ctx) override;
  // std::any visitFuncFParam(SysYParser::FuncFParamContext *ctx) override;
  // std::any visitFuncFParams(SysYParser::FuncFParamsContext *ctx) override;
  std::any visitBlockStmt(SysYParser::BlockStmtContext* ctx) override;
  // std::any visitStmt(SysYParser::StmtContext *ctx) override;
  std::any visitAssignStmt(SysYParser::AssignStmtContext *ctx) override;
  std::any visitIfStmt(SysYParser::IfStmtContext *ctx) override;
  std::any visitWhileStmt(SysYParser::WhileStmtContext *ctx) override;
  std::any visitBreakStmt(SysYParser::BreakStmtContext *ctx) override;
  std::any visitContinueStmt(SysYParser::ContinueStmtContext *ctx) override;
  std::any visitReturnStmt(SysYParser::ReturnStmtContext *ctx) override;
  // std::any visitExp(SysYParser::ExpContext *ctx) override;
  std::any visitLValue(SysYParser::LValueContext *ctx) override;
  std::any visitPrimExp(SysYParser::PrimExpContext *ctx) override;
  // std::any visitParenExp(SysYParser::ParenExpContext *ctx) override;
  std::any visitNumber(SysYParser::NumberContext *ctx) override;
  // std::any visitString(SysYParser::StringContext *ctx) override;
  std::any visitCall(SysYParser::CallContext *ctx) override;
  // std::any visitUnaryExp(SysYParser::UnaryExpContext *ctx) override;
  // std::any visitUnaryOp(SysYParser::UnaryOpContext *ctx) override;
  std::any visitUnExp(SysYParser::UnExpContext *ctx) override;
  // std::any visitFuncRParams(SysYParser::FuncRParamsContext *ctx) override;
  std::any visitMulExp(SysYParser::MulExpContext *ctx) override;
  std::any visitAddExp(SysYParser::AddExpContext *ctx) override;
  std::any visitRelExp(SysYParser::RelExpContext *ctx) override;
  std::any visitEqExp(SysYParser::EqExpContext *ctx) override;
  std::any visitLAndExp(SysYParser::LAndExpContext *ctx) override;
  std::any visitLOrExp(SysYParser::LOrExpContext *ctx) override;
  std::any visitConstExp(SysYParser::ConstExpContext *ctx) override;

private:
  std::any visitConstGlobalDecl(SysYParser::ConstDeclContext *ctx, Type* type);
  std::any visitVarGlobalDecl(SysYParser::VarDeclContext *ctx, Type* type);
  std::any visitConstLocalDecl(SysYParser::ConstDeclContext *ctx, Type* type);
  std::any visitVarLocalDecl(SysYParser::VarDeclContext *ctx, Type* type);
  Type *getArithmeticResultType(Type *lhs, Type *rhs) {
    assert(lhs->isIntOrFloat() and rhs->isIntOrFloat());
    return lhs == rhs ? lhs : Type::getFloatType();
  }

}; // class SysYIRGenerator

} // namespace sysy