// 将语法树翻译为 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"

#define IRGen_DEBUG 0

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;          // 注意：规则名为 Block
  std::any visitBlockItem(SysYParser::BlockItemContext* ctx) override;
  
  // 声明
  std::any visitDecl(SysYParser::DeclContext* ctx) override;
  std::any visitVarDef(SysYParser::VarDefContext* ctx) override;
  std::any visitConstDecl(SysYParser::ConstDeclContext* ctx) override;
  std::any visitConstDef(SysYParser::ConstDefContext* ctx) override;
  std::any visitInitVal(SysYParser::InitValContext* ctx) override;
  std::any visitConstInitVal(SysYParser::ConstInitValContext* ctx) override;
  
  // 语句
  std::any visitStmt(SysYParser::StmtContext* ctx) override;
  
  // 表达式
  // 基本表达式（变量、常量、括号表达式）直接翻译为 IR 中的值；函数调用和一元运算需要特殊处理。
  std::any visitExp(SysYParser::ExpContext* ctx) override;
  std::any visitCond(SysYParser::CondContext* ctx) override;
  std::any visitLVal(SysYParser::LValContext* ctx) override;
  std::any visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) override;
  
  
  std::any visitUnaryExp(SysYParser::UnaryExpContext* ctx) override;
  std::any visitFuncRParams(SysYParser::FuncRParamsContext* ctx) override;
  std::any visitMulExp(SysYParser::MulExpContext* ctx) override;
  // 加法表达式、乘法表达式、关系表达式、相等表达式、条件表达式分别对应不同的访问函数，按照优先级分层访问，最终调用 visitLVal 来处理变量访问
  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;

  // 辅助函数
  ir::Value* EvalExpr(SysYParser::ExpContext& expr);   // 只保留一处
  ir::Value* EvalCond(SysYParser::CondContext& cond);
  std::any visitCallExp(SysYParser::UnaryExpContext* ctx);
  std::vector<ir::Value*> ProcessNestedInitVals(SysYParser::InitValContext* ctx);
  int TryEvaluateConstInt(SysYParser::ConstExpContext* ctx);
  void AddRuntimeFunctions();
  ir::Function* CreateRuntimeFunctionDecl(const std::string& funcName);
private:
  // 辅助函数声明
  enum class BlockFlow{
    Continue,
    Terminated,
  };

  BlockFlow VisitBlockItemResult(SysYParser::BlockItemContext& item);

  BlockFlow HandleReturnStmt(SysYParser::StmtContext* ctx);
  BlockFlow HandleIfStmt(SysYParser::StmtContext* ctx);
  BlockFlow HandleWhileStmt(SysYParser::StmtContext* ctx);
  BlockFlow HandleBreakStmt(SysYParser::StmtContext* ctx);
  BlockFlow HandleContinueStmt(SysYParser::StmtContext* ctx);
  BlockFlow HandleAssignStmt(SysYParser::StmtContext* ctx);

  // 完全正确的左值判断函数
  bool IsLValueExpression(SysYParser::ExpContext* ctx);
  bool IsLValueInAddExp(SysYParser::AddExpContext* ctx);
  bool IsLValueInMulExp(SysYParser::MulExpContext* ctx);
  bool IsLValueInUnaryExp(SysYParser::UnaryExpContext* ctx);
  bool IsLValueInPrimaryExp(SysYParser::PrimaryExpContext* ctx);
  // 循环上下文结构
  struct LoopContext {
    ir::BasicBlock* condBlock;
    ir::BasicBlock* bodyBlock;
    ir::BasicBlock* exitBlock;
  };

  struct ArrayInfo {
    std::vector<ir::Value*> elements;
    std::vector<int> dimensions;
  };

  std::vector<LoopContext> loopStack_;

  ir::Module& module_;
  const SemanticContext& sema_;
  ir::Function* func_;
  ir::IRBuilder builder_;
  ir::Value* EvalAssign(SysYParser::StmtContext* ctx);
 
  // 按 VarDefContext 查找存储位置（用于数组访问等场景）
  std::unordered_map<SysYParser::VarDefContext*, ir::Value*> storage_map_;
  
  // 按变量名快速查找（用于 LVal 等场景）
  std::unordered_map<std::string, ir::Value*> local_var_map_;      // 局部变量
  std::unordered_map<std::string, ir::GlobalValue*> global_map_;   // 全局变量
  std::unordered_map<std::string, ir::Value*> param_map_;          // 函数参数
  
  // 常量映射：常量名 -> 常量值（标量常量）
  std::unordered_map<std::string, ir::ConstantValue*> const_value_map_;
  
  // 全局常量映射：常量名 -> 全局变量（数组常量）
  std::unordered_map<std::string, ir::GlobalValue*> const_global_map_;
  
  // 原有的常量存储映射（用于兼容）
  std::unordered_map<antlr4::ParserRuleContext*, ir::Value*> const_storage_map_;
  
  std::unordered_map<SysYParser::VarDefContext*, ArrayInfo> array_info_map_;

  // 新增：处理全局和局部变量的辅助函数
  std::any HandleGlobalVariable(SysYParser::VarDefContext* ctx, 
                                const std::string& varName, 
                                bool is_array);
  std::any HandleLocalVariable(SysYParser::VarDefContext* ctx, 
                               const std::string& varName, 
                               bool is_array);
  
  // 常量求值辅助函数
  int EvaluateConstAddExp(SysYParser::AddExpContext* ctx);
  int EvaluateConstMulExp(SysYParser::MulExpContext* ctx);
  int EvaluateConstUnaryExp(SysYParser::UnaryExpContext* ctx);
  int EvaluateConstPrimaryExp(SysYParser::PrimaryExpContext* ctx);
  int EvaluateConstExp(SysYParser::ExpContext* ctx);                             
};

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