// 将语法树翻译为 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:
  // const 变量名 -> 编译期整数值，用于数组维度折叠。
  using ConstEnv = std::unordered_map<std::string, int>;

  IRGenImpl(ir::Module& module, const SemanticContext& sema);

  std::any visitCompUnit(SysYParser::CompUnitContext* ctx) override;
  std::any visitFuncDef(SysYParser::FuncDefContext* ctx) override;
  std::any visitBlockStmt(SysYParser::BlockStmtContext* 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 visitVarDecl(SysYParser::VarDeclContext* ctx) override;
  std::any visitStmt(SysYParser::StmtContext* ctx) override;
  std::any visitVarDef(SysYParser::VarDefContext* 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 visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) override;
  std::any visitNumber(SysYParser::NumberContext* ctx) override;
  std::any visitLValue(SysYParser::LValueContext* ctx) override;
  std::any visitUnaryExp(SysYParser::UnaryExpContext* 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;

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

  struct LoopTargets {
    ir::BasicBlock* continue_target;
    ir::BasicBlock* break_target;
  };

  // 判断当前是否处于全局作用域（函数外部）。
  bool IsGlobalScope() const { return func_ == nullptr; }

  BlockFlow VisitBlockItemResult(SysYParser::BlockItemContext& item);
  ir::Value* EvalExpr(SysYParser::ExpContext& expr);
  ir::Value* EvalCond(SysYParser::CondContext& cond);
  ir::Value* ToBoolValue(ir::Value* v);
  std::string NextBlockName();

  // 预声明 SysY runtime 外部函数。
  void DeclareRuntimeFunctions();

  // 根据 sema 绑定或 name 查找局部/全局存储槽位（返回 i32* Value）。
  // 如果 lvalue 有下标，还会生成 GEP 指令并返回元素指针。
  ir::Value* ResolveStorage(SysYParser::LValueContext* lvalue);

  // 编译期常量整数求值（用于数组维度）。
  int EvalConstExpr(SysYParser::ConstExpContext* ctx) const;
  // 将 ExpContext（即 addExp）按编译期常量求值（用于 funcFParam 维度）。
  int EvalExpAsConst(SysYParser::ExpContext* ctx) const;

  // 查找变量的数组维度（先查局部，再查全局）。
  const std::vector<int>* FindArrayDims(const std::string& name) const;

  // 将一组数组下标表达式（已求值为 ir::Value*）折叠为线性偏移 ir::Value*。
  ir::Value* ComputeLinearIndex(const std::vector<int>& dims,
                                const std::vector<SysYParser::ExpContext*>& subs);

  // 扁平化 constInitValue 到整数数组（供 const 数组初始化使用）。
  void FlattenConstInit(SysYParser::ConstInitValueContext* ctx,
                        const std::vector<int>& dims, int dim_idx,
                        std::vector<int>& out, int& pos);

  // 扁平化 initValue 到 ir::Value* 数组（供普通数组初始化使用）。
  void FlattenInit(SysYParser::InitValueContext* ctx,
                   const std::vector<int>& dims, int dim_idx,
                   std::vector<ir::Value*>& out, int& pos);

  ir::AllocaInst* CreateEntryAllocaI32(const std::string& name);
  ir::AllocaInst* CreateEntryAllocaArray(int count, const std::string& name);

  ir::Module& module_;
  const SemanticContext& sema_;
  ir::Function* func_;
  ir::IRBuilder builder_;
  // 声明 -> 存储槽位（局部 alloca 或全局变量，均为 i32*）。
  std::unordered_map<SysYParser::VarDefContext*, ir::Value*> storage_map_;
  // 名称 -> 槽位（参数、const 变量等不经 sema binding 的后备查找）。
  std::unordered_map<std::string, ir::Value*> named_storage_;
  // 全局变量名 -> GlobalVariable*（跨函数持久）。
  std::unordered_map<std::string, ir::Value*> global_storage_;
  // 编译期 const 整数环境（全局 + 当前函数）。
  ConstEnv const_env_;
  // 数组维度信息：全局数组（跨函数持久）。
  std::unordered_map<std::string, std::vector<int>> global_array_dims_;
  // 数组维度信息：局部数组/参数（每函数清空）。
  std::unordered_map<std::string, std::vector<int>> local_array_dims_;
  // 逻辑与/或短路求值复用的函数级临时槽位，避免循环中动态 alloca 导致栈膨胀。
  ir::Value* short_circuit_slot_ = nullptr;
  std::vector<LoopTargets> loop_stack_;
};

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