#ifndef SEMANTIC_ANALYSIS_H
#define SEMANTIC_ANALYSIS_H

#include <any>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

#include "SymbolTable.h"
#include "SysYBaseVisitor.h"
#include "SysYParser.h"

struct ErrorMsg {
  std::string msg;
  int line;
  int column;

  ErrorMsg(std::string m, int l, int c) : msg(std::move(m)), line(l), column(c) {}
};

class IRGenContext {
 public:
  void RecordError(const ErrorMsg& err) { errors_.push_back(err); }
  const std::vector<ErrorMsg>& GetErrors() const { return errors_; }
  bool HasError() const { return !errors_.empty(); }
  void ClearErrors() { errors_.clear(); }

  void SetType(void* ctx, SymbolType type) { node_type_map_[ctx] = type; }

  SymbolType GetType(void* ctx) const {
    auto it = node_type_map_.find(ctx);
    return it == node_type_map_.end() ? SymbolType::TYPE_UNKNOWN : it->second;
  }

  void SetConstVal(void* ctx, const std::any& val) { const_val_map_[ctx] = val; }

  std::any GetConstVal(void* ctx) const {
    auto it = const_val_map_.find(ctx);
    return it == const_val_map_.end() ? std::any() : it->second;
  }

  void EnterLoop() { sym_table_.EnterLoop(); }
  void ExitLoop() { sym_table_.ExitLoop(); }
  bool InLoop() const { return sym_table_.InLoop(); }

  bool IsIntType(const std::any& val) const {
    return val.type() == typeid(long) || val.type() == typeid(int);
  }

  bool IsFloatType(const std::any& val) const {
    return val.type() == typeid(double) || val.type() == typeid(float);
  }

  SymbolType GetCurrentFuncReturnType() const { return current_func_ret_type_; }
  void SetCurrentFuncReturnType(SymbolType type) { current_func_ret_type_ = type; }

  SymbolTable& GetSymbolTable() { return sym_table_; }
  const SymbolTable& GetSymbolTable() const { return sym_table_; }

  void EnterScope() { sym_table_.EnterScope(); }
  void LeaveScope() { sym_table_.LeaveScope(); }
  size_t GetScopeDepth() const { return sym_table_.GetScopeDepth(); }

 private:
  SymbolTable sym_table_;
  std::unordered_map<void*, SymbolType> node_type_map_;
  std::unordered_map<void*, std::any> const_val_map_;
  std::vector<ErrorMsg> errors_;
  SymbolType current_func_ret_type_ = SymbolType::TYPE_UNKNOWN;
};

class SemanticContext {
 public:
  void BindVarUse(const SysYParser::LValueContext* use,
                  SysYParser::VarDefContext* decl) {
    var_uses_[use] = decl;
  }

  SysYParser::VarDefContext* ResolveVarUse(
      const SysYParser::LValueContext* use) const {
    auto it = var_uses_.find(use);
    return it == var_uses_.end() ? nullptr : it->second;
  }

 private:
  std::unordered_map<const SysYParser::LValueContext*, SysYParser::VarDefContext*>
      var_uses_;
};

inline std::string FormatErrMsg(const std::string& msg, int line, int col) {
  std::ostringstream oss;
  oss << "[行:" << line << ",列:" << col << "] " << msg;
  return oss.str();
}

class SemaVisitor : public SysYBaseVisitor {
 public:
  explicit SemaVisitor(IRGenContext& ctx, SemanticContext* sema_ctx = nullptr)
      : ir_ctx_(ctx), sema_ctx_(sema_ctx) {}

  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 visitConstInitValue(SysYParser::ConstInitValueContext* ctx) override;
  std::any visitVarDecl(SysYParser::VarDeclContext* ctx) override;
  std::any visitVarDef(SysYParser::VarDefContext* ctx) override;
  std::any visitInitValue(SysYParser::InitValueContext* ctx) override;
  std::any visitFuncDef(SysYParser::FuncDefContext* ctx) override;
  std::any visitFuncType(SysYParser::FuncTypeContext* ctx) override;
  std::any visitFuncFParams(SysYParser::FuncFParamsContext* ctx) override;
  std::any visitFuncFParam(SysYParser::FuncFParamContext* ctx) override;
  std::any visitBlockStmt(SysYParser::BlockStmtContext* ctx) override;
  std::any visitBlockItem(SysYParser::BlockItemContext* ctx) override;
  std::any visitStmt(SysYParser::StmtContext* ctx) override;
  std::any visitReturnStmt(SysYParser::ReturnStmtContext* ctx) override;
  std::any visitExp(SysYParser::ExpContext* ctx) override;
  std::any visitCond(SysYParser::CondContext* ctx) override;
  std::any visitLValue(SysYParser::LValueContext* ctx) override;
  std::any visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) override;
  std::any visitNumber(SysYParser::NumberContext* ctx) override;
  std::any visitUnaryExp(SysYParser::UnaryExpContext* ctx) override;
  std::any visitUnaryOp(SysYParser::UnaryOpContext* 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;

  IRGenContext& GetContext() { return ir_ctx_; }
  const IRGenContext& GetContext() const { return ir_ctx_; }

 private:
  void RecordNodeError(antlr4::ParserRuleContext* ctx, const std::string& msg);

  IRGenContext& ir_ctx_;
  SemanticContext* sema_ctx_ = nullptr;
  SymbolType current_decl_type_ = SymbolType::TYPE_UNKNOWN;
  bool current_decl_is_const_ = false;
};

void RunSemanticAnalysis(SysYParser::CompUnitContext* ctx, IRGenContext& ir_ctx);
SemanticContext RunSema(SysYParser::CompUnitContext& comp_unit);

#endif  // SEMANTIC_ANALYSIS_H