// 将语法树翻译为 IR。
// 实现拆分在 IRGenFunc/IRGenStmt/IRGenExp/IRGenDecl。

#pragma once

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

#include "SysYBaseVisitor.h"
#include "SysYParser.h"
#include "ir/IR.h"
#include "sem/Sema.h"

namespace ir {
class Module;
class Function;
class IRBuilder;
class Value;
}

class IRGenImpl final : public SysYBaseVisitor {
 public:
  IRGenImpl(ir::Module& module, const SemanticContext& sema);

  std::any visitCompUnit(SysYParser::CompUnitContext* ctx) override;
  std::any visitFuncDef(SysYParser::FuncDefContext* ctx) override;
  std::any visitBlock(SysYParser::BlockContext* ctx) override;
  std::any visitBlockItem(SysYParser::BlockItemContext* ctx) override;
  std::any visitDecl(SysYParser::DeclContext* ctx) override;
  std::any visitVarDecl(SysYParser::VarDeclContext* ctx) override;
  std::any visitStmt(SysYParser::StmtContext* ctx) override;
  std::any visitVarDef(SysYParser::VarDefContext* ctx) override;
  std::any visitExp(SysYParser::ExpContext* ctx) override;
  std::any visitAddExp(SysYParser::AddExpContext* ctx) override;
  std::any visitMulExp(SysYParser::MulExpContext* ctx) override;
  std::any visitUnaryExp(SysYParser::UnaryExpContext* ctx) override;
  std::any visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) override;
  std::any visitLVal(SysYParser::LValContext* ctx) override;
  std::any visitNumber(SysYParser::NumberContext* ctx) override;

 private:
  enum class BlockFlow {
    Continue,
    Terminated,
  };

  BlockFlow VisitBlockItemResult(SysYParser::BlockItemContext& item);
  ir::Value* EvalExpr(SysYParser::ExpContext& expr);
  ir::Value* EvalBinaryOrFold(ir::Opcode op, ir::Value* lhs, ir::Value* rhs);
  std::shared_ptr<ir::Type> ResolveBType(SysYParser::BTypeContext* btype) const;
    int EvalConstIntExpr(SysYParser::ExpContext& expr);
    int EvalConstIntExpr(SysYParser::ConstExpContext& expr);
    int EvalConstIntAddExp(SysYParser::AddExpContext& expr);
    int EvalConstIntMulExp(SysYParser::MulExpContext& expr);
    int EvalConstIntUnaryExp(SysYParser::UnaryExpContext& expr);
    int EvalConstIntPrimaryExp(SysYParser::PrimaryExpContext& expr);
    double EvalConstFloatExpr(SysYParser::ConstExpContext& expr);
    double EvalConstFloatAddExp(SysYParser::AddExpContext& expr);
    double EvalConstFloatMulExp(SysYParser::MulExpContext& expr);
    double EvalConstFloatUnaryExp(SysYParser::UnaryExpContext& expr);
    double EvalConstFloatPrimaryExp(SysYParser::PrimaryExpContext& expr);
    std::vector<int> EvalArrayExtents(
      const std::vector<SysYParser::ConstExpContext*>& dims);
    std::vector<int> GetArrayExtentsForDecl(SysYParser::VarDefContext* decl);
    std::vector<int> GetArrayExtentsForConstDecl(
      SysYParser::ConstDefContext* decl);
    std::vector<int> GetArrayExtentsForLVal(SysYParser::LValContext& lval,
                        bool& is_array);
    ir::Value* BuildLinearizedIndex(
      const std::vector<ir::Value*>& indices,
      const std::vector<int>& extents_with_first_dim) ;
  ir::Value* CastValueTo(ir::Value* value,
                         const std::shared_ptr<ir::Type>& target_type);
  ir::Value* GetLValAddress(SysYParser::LValContext& lval);
  ir::AllocaInst* CreateEntryBlockAlloca(std::shared_ptr<ir::Type> elem_ty,
                                         const std::string& name,
                                         ir::Value* count = nullptr);
  std::string NextBlockName(const std::string& prefix);
  void EmitCondBranch(SysYParser::CondContext& cond, ir::BasicBlock* true_bb,
                      ir::BasicBlock* false_bb);
  void EmitLOrBranch(SysYParser::LOrExpContext& expr, ir::BasicBlock* true_bb,
                     ir::BasicBlock* false_bb);
  void EmitLAndBranch(SysYParser::LAndExpContext& expr, ir::BasicBlock* true_bb,
                      ir::BasicBlock* false_bb);
  void EmitEqBranch(SysYParser::EqExpContext& expr, ir::BasicBlock* true_bb,
                    ir::BasicBlock* false_bb);
  void EmitRelBranch(SysYParser::RelExpContext& expr, ir::BasicBlock* true_bb,
                     ir::BasicBlock* false_bb);
  ir::Value* EvalEqValue(SysYParser::EqExpContext& expr);
  ir::Value* EvalRelValue(SysYParser::RelExpContext& expr);

  ir::Module& module_;
  const SemanticContext& sema_;
  ir::Function* func_;
  ir::IRBuilder builder_;
  std::unordered_map<std::string, ir::Function*> function_map_;
  std::unordered_map<std::string, int> const_value_map_;
  std::vector<std::unordered_map<std::string, int>> local_const_stack_;
  std::vector<std::unordered_map<std::string, int>> const_value_history_;
  std::unordered_map<SysYParser::VarDefContext*, std::vector<int>>
      array_extents_map_;
    std::unordered_map<SysYParser::ConstDefContext*, std::vector<int>>
      const_array_extents_map_;
  std::unordered_map<std::string, std::vector<int>> param_array_extents_map_;
  std::unordered_map<std::string, ir::Value*> param_storage_map_;
  std::unordered_map<std::string, ir::Value*> param_pointer_map_;
  std::unordered_map<SysYParser::VarDefContext*, ir::Value*> global_storage_map_;
    std::unordered_map<SysYParser::ConstDefContext*, ir::Value*>
      const_global_storage_map_;
  // 名称绑定由 Sema 负责；IRGen 只维护“声明 -> 存储槽位”的代码生成状态。
  std::unordered_map<SysYParser::VarDefContext*, ir::Value*> storage_map_;
    std::unordered_map<SysYParser::ConstDefContext*, ir::Value*>
      const_storage_map_;
  std::vector<std::pair<ir::BasicBlock*, ir::BasicBlock*>> loop_stack_;
  int block_index_ = 0;
};

std::unique_ptr<ir::Module> GenerateIR(SysYParser::CompUnitContext& tree,
                                       const SemanticContext& sema);