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

#pragma once

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

#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 visitConstDecl(SysYParser::ConstDeclContext* ctx) override;
  std::any visitConstDef(SysYParser::ConstDefContext* ctx) override;
  std::any visitStmt(SysYParser::StmtContext* ctx) override;
  std::any visitVarDef(SysYParser::VarDefContext* ctx) override;
  std::any visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) override;
  std::any visitNumber(SysYParser::NumberContext* ctx) override;
  std::any visitLVal(SysYParser::LValContext* 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 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 visitCondUnaryExp(SysYParser::CondUnaryExpContext* ctx) override;
  std::any visitCond(SysYParser::CondContext* ctx) override;
  std::any visitFuncRParams(SysYParser::FuncRParamsContext* ctx) override;

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

  BlockFlow VisitBlockItemResult(SysYParser::BlockItemContext& item);
  ir::Value* EvalExpr(SysYParser::ExpContext& expr);
  ir::ConstantValue* EvaluateConst(antlr4::tree::ParseTree* tree);
  int EvaluateConstInt(SysYParser::ConstExpContext* ctx);
  int EvaluateConstInt(SysYParser::ExpContext* ctx);

  std::shared_ptr<ir::Type> GetGEPResultType(ir::Value* ptr, const std::vector<ir::Value*>& indices);

  // Flatten array initializers
  void FlattenInitVal(SysYParser::InitValContext* ctx, 
                      const std::vector<int>& dims, 
                      const std::vector<int>& sub_sizes,
                      int dim_idx,
                      size_t& current_pos,
                      std::vector<ir::Value*>& results,
                      bool is_float);
  
  void FlattenConstInitVal(SysYParser::ConstInitValContext* ctx, 
                           const std::vector<int>& dims, 
                           const std::vector<int>& sub_sizes,
                           int dim_idx,
                           size_t& current_pos,
                           std::vector<ir::ConstantValue*>& results,
                           bool is_float);

  ir::Module& module_;
  const SemanticContext& sema_;
  ir::Function* func_;
  ir::IRBuilder builder_;
  // 考虑到嵌套作用域（全局、函数、语句块），使用 vector 模拟栈来管理 storage_map_ 和 const_values_
  std::vector<std::unordered_map<std::string, ir::Value*>> storage_map_stack_;
  std::vector<std::unordered_map<std::string, ir::ConstantValue*>> const_values_stack_;

  // 用于在栈中查找变量
  ir::Value* FindStorage(const std::string& name) const {
    for (auto it = storage_map_stack_.rbegin(); it != storage_map_stack_.rend(); ++it) {
      if (it->count(name)) return it->at(name);
    }
    return nullptr;
  }
  
  ir::ConstantValue* FindConst(const std::string& name) const {
    for (auto it = const_values_stack_.rbegin(); it != const_values_stack_.rend(); ++it) {
      if (it->count(name)) return it->at(name);
    }
    return nullptr;
  }

  // 用于 break 和 continue 跳转的目标位置
  ir::BasicBlock* current_loop_cond_bb_ = nullptr;
  ir::BasicBlock* current_loop_exit_bb_ = nullptr;

  int bb_cnt_ = 0;
  std::string NextBlockName(const std::string& prefix = "bb") {
    return prefix + "_" + std::to_string(++bb_cnt_);
  }
};

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