// 极简符号表：记录局部变量定义点。
#pragma once

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

#include "SysYParser.h"
#include "ir/IR.h"

// 符号种类
enum class SymbolKind {
    Variable,
    Function,
    Parameter,
    Constant
};

// 符号条目
// 符号条目
struct Symbol {
    // 基本信息
    std::string name;
    SymbolKind kind;
    std::shared_ptr<ir::Type> type;
    int scope_level = 0;
    int stack_offset = -1;
    bool is_initialized = false;
    bool is_builtin = false;

    // 数组参数相关
    std::vector<int> array_dims;
    bool is_array_param = false;

    // 函数相关
    std::vector<std::shared_ptr<ir::Type>> param_types;

    // 常量值存储
    union ConstantValue {
        int i32;
        float f32;
    };
    
    // 标量常量
    bool is_int_const = true;
    ConstantValue const_value;
    
    // 数组常量（扁平化存储）
    bool is_array_const = false;
    std::vector<ConstantValue> array_const_values;

    // 语法树节点
    SysYParser::VarDefContext* var_def_ctx = nullptr;
    SysYParser::ConstDefContext* const_def_ctx = nullptr;
    SysYParser::FuncFParamContext* param_def_ctx = nullptr;
    SysYParser::FuncDefContext* func_def_ctx = nullptr;
    
    // 辅助方法
    bool IsScalarConstant() const { 
        return kind == SymbolKind::Constant && !type->IsArray(); 
    }
    
    bool IsArrayConstant() const { 
        return kind == SymbolKind::Constant && type->IsArray(); 
    }
    
    int GetIntConstant() const {
        if (!IsScalarConstant()) {
            throw std::runtime_error("不是标量常量");
        }
        if (!is_int_const) {
            throw std::runtime_error("不是整型常量");
        }
        return const_value.i32;
    }
    
    float GetFloatConstant() const {
        if (!IsScalarConstant()) {
            throw std::runtime_error("不是标量常量");
        }
        if (is_int_const) {
            return static_cast<float>(const_value.i32);
        }
        return const_value.f32;
    }
    
    ConstantValue GetArrayElement(size_t index) const {
        if (!IsArrayConstant()) {
            throw std::runtime_error("不是数组常量");
        }
        if (index >= array_const_values.size()) {
            throw std::runtime_error("数组下标越界");
        }
        return array_const_values[index];
    }
    
    size_t GetArraySize() const {
        if (!IsArrayConstant()) return 0;
        return array_const_values.size();
    }
};
class SymbolTable {
  public:
    SymbolTable();
    ~SymbolTable() = default;
    // 添加调试方法
    size_t getScopeCount() const { return active_scope_stack_.size(); }
    
    void dump() const {
        std::cerr << "=== SymbolTable Dump ===" << std::endl;
        for (size_t i = 0; i < scopes_.size(); ++i) {
            std::cerr << "Scope " << i << " (depth=" << i << ")";
            bool active = std::find(active_scope_stack_.begin(), active_scope_stack_.end(), i) != active_scope_stack_.end();
            std::cerr << (active ? " [active]" : " [inactive]") << std::endl;
            for (const auto& [name, sym] : scopes_[i]) {
                std::cerr << "  " << name 
                          << " (kind=" << (int)sym.kind 
                          << ", level=" << sym.scope_level << ")" << std::endl;
            }
        }
    }
    // ----- 作用域管理 -----
    void enterScope();               // 进入新作用域
    void exitScope();                // 退出当前作用域
    int currentScopeLevel() const { return static_cast<int>(active_scope_stack_.size()) - 1; }

    // ----- 符号操作（推荐使用）-----
    bool addSymbol(const Symbol& sym);                     // 添加符号到当前作用域
    Symbol* lookup(const std::string& name);               // 从当前作用域向外查找
    Symbol* lookupCurrent(const std::string& name);        // 仅在当前作用域查找
    const Symbol* lookup(const std::string& name) const;
    const Symbol* lookupCurrent(const std::string& name) const;
    const Symbol* lookupAll(const std::string& name) const;                // 所有作用域查找，包括已结束的作用域
    const Symbol* lookupByVarDef(const SysYParser::VarDefContext* decl) const;   // 通过定义节点查找符号
    const Symbol* lookupByConstDef(const SysYParser::ConstDefContext* decl) const; // 通过常量定义节点查找符号

    // ----- 与原接口兼容（保留原有功能）-----
    void Add(const std::string& name, SysYParser::VarDefContext* decl);
    bool Contains(const std::string& name) const;
    SysYParser::VarDefContext* Lookup(const std::string& name) const;

    // ----- 辅助函数：从语法树节点构造 Type -----
    static std::shared_ptr<ir::Type> getTypeFromFuncDef(SysYParser::FuncDefContext* ctx);

    void registerBuiltinFunctions();
    
    // 对常量表达式求值（返回整数值，用于数组维度等）
    int EvaluateConstExp(SysYParser::ConstExpContext* ctx) const;

    // 对常量表达式求值（返回浮点值，用于全局初始化）
    float EvaluateConstExpFloat(SysYParser::ConstExpContext* ctx) const;

    // 对常量初始化列表求值，返回一系列常量值（扁平化）
    struct ConstValue {
        enum Kind { INT, FLOAT };
        Kind kind;
        union {
            int int_val;
            float float_val;
        };
    };
    void flattenInit(SysYParser::ConstInitValContext* ctx,
                              std::vector<ConstValue>& out,
                              std::shared_ptr<ir::Type> base_type) const;
    std::vector<ConstValue> EvaluateConstInitVal(
        SysYParser::ConstInitValContext* ctx,
        const std::vector<int>& dims,
        std::shared_ptr<ir::Type> base_type) const;
    
    int EvaluateConstExpression(SysYParser::ExpContext* ctx) const;

  private:
    // 作用域栈：每个元素是一个从名字到符号的映射
    std::vector<std::unordered_map<std::string, Symbol>> scopes_;
    std::vector<size_t> active_scope_stack_;

    static constexpr int GLOBAL_SCOPE = 0;  // 全局作用域索引
    
    ConstValue EvaluateAddExp(SysYParser::AddExpContext* ctx) const;
    ConstValue EvaluateMulExp(SysYParser::MulExpContext* ctx) const;
    ConstValue EvaluateUnaryExp(SysYParser::UnaryExpContext* ctx) const;
    ConstValue EvaluatePrimaryExp(SysYParser::PrimaryExpContext* ctx) const;

    std::shared_ptr<ir::Type> getTypeFromVarDef(SysYParser::VarDefContext* ctx) const;
};