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

#pragma once

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

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

//#define DEBUG_IRGen

#ifdef DEBUG_IRGen
#include <iostream>
#define DEBUG_MSG(msg) std::cerr << "[IRGen Debug] " << msg << std::endl
#else
#define DEBUG_MSG(msg)
#endif

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

class IRGenImpl final : public SysYBaseVisitor {
public:
  // 修改构造函数，添加 SymbolTable 参数
  IRGenImpl(ir::Module& module, 
            const SemanticContext& sema,
            const SymbolTable& sym_table);  // 新增

  // 顶层
  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);
  // 带维度感知的展平：按 C 语言花括号对齐规则填充 total_size 个槽位
  // dims[0] 是最外层维度，dims.back() 是最内层维度（元素层）
  // 返回已展平并补零的 total_size 大小的向量
  std::vector<ir::Value*> FlattenInitVal(SysYParser::InitValContext* ctx,
                                         const std::vector<int>& dims,
                                         bool is_float);
  int TryEvaluateConstInt(SysYParser::ConstExpContext* ctx);
  void AddRuntimeFunctions();
  ir::Function* CreateRuntimeFunctionDecl(const std::string& funcName);
  ir::BasicBlock* EnsureCleanupBlock();
  void RegisterCleanup(ir::Function* free_func, ir::Value* ptr);
  ir::AllocaInst* CreateEntryAlloca(std::shared_ptr<ir::Type> ty,
                                    const std::string& name);
  ir::AllocaInst* CreateEntryAllocaI32(const std::string& name);
  ir::AllocaInst* CreateEntryAllocaFloat(const std::string& name);
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_;
  const SymbolTable& symbol_table_;  // 新增成员
  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_set<std::string> pointer_param_names_;            // 指针/数组形参名
  std::unordered_set<std::string> heap_local_array_names_;         // 堆分配的局部数组名
  
  // 常量映射：常量名 -> 常量值（标量常量）
  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::string current_function_name_;
  bool current_function_is_recursive_ = false;
  ir::AllocaInst* function_return_slot_ = nullptr;
  ir::BasicBlock* function_cleanup_block_ = nullptr;
  std::vector<std::pair<ir::Function*, ir::Value*>> function_cleanup_actions_;

  // 新增：处理全局和局部变量的辅助函数
  // 修改处理函数的签名，使用 Symbol* 参数
  std::any HandleGlobalVariable(SysYParser::VarDefContext* ctx, 
                                const std::string& varName,
                                const Symbol* sym);
  
  std::any HandleLocalVariable(SysYParser::VarDefContext* ctx, 
                               const std::string& varName,
                               const Symbol* sym);                       
};

// 修改 GenerateIR 函数签名
std::unique_ptr<ir::Module> GenerateIR(SysYParser::CompUnitContext& tree,
                                       const SemaResult& sema_result);