// 基于语法树的语义检查与名称绑定（Lab2 扩展）
#pragma once

#include <optional>
#include <unordered_map>
#include <vector>

#include "SysYParser.h"
#include "sem/SymbolTable.h"

struct FuncTypeDesc {
  TypeDesc ret;
  std::vector<TypeDesc> params;
};

struct BoundDecl {
  enum class Kind { Var, Const, Param } kind = Kind::Var;
  SysYParser::VarDefContext* var_decl = nullptr;
  SysYParser::ConstDefContext* const_decl = nullptr;
  SysYParser::FuncFParamContext* param_decl = nullptr;
};

class SemanticContext {
 public:
  void BindVarUse(SysYParser::LValContext* use, BoundDecl decl) {
    var_uses_[use] = decl;
  }

  BoundDecl ResolveVarUse(const SysYParser::LValContext* use) const {
    auto it = var_uses_.find(use);
    return it == var_uses_.end() ? BoundDecl{} : it->second;
  }

  void RegisterVarDecl(SysYParser::VarDefContext* decl, TypeDesc ty) {
    var_types_[decl] = std::move(ty);
  }

  void RegisterConstDecl(SysYParser::ConstDefContext* decl, TypeDesc ty) {
    const_types_[decl] = std::move(ty);
  }

  void RegisterParam(SysYParser::FuncFParamContext* decl, TypeDesc ty) {
    param_types_[decl] = std::move(ty);
  }

  void RegisterFunc(SysYParser::FuncDefContext* decl, FuncTypeDesc ty) {
    func_types_[decl] = std::move(ty);
  }

  const TypeDesc* GetVarType(const SysYParser::VarDefContext* decl) const {
    auto it = var_types_.find(decl);
    return it == var_types_.end() ? nullptr : &it->second;
  }

  const TypeDesc* GetConstType(const SysYParser::ConstDefContext* decl) const {
    auto it = const_types_.find(decl);
    return it == const_types_.end() ? nullptr : &it->second;
  }

  const TypeDesc* GetParamType(const SysYParser::FuncFParamContext* decl) const {
    auto it = param_types_.find(decl);
    return it == param_types_.end() ? nullptr : &it->second;
  }

  const FuncTypeDesc* GetFuncType(const SysYParser::FuncDefContext* decl) const {
    auto it = func_types_.find(decl);
    return it == func_types_.end() ? nullptr : &it->second;
  }

  void BindFuncCall(SysYParser::UnaryExpContext* call,
                    SysYParser::FuncDefContext* decl) {
    func_calls_[call] = decl;
  }

  SysYParser::FuncDefContext* ResolveFuncCall(
      const SysYParser::UnaryExpContext* call) const {
    auto it = func_calls_.find(call);
    return it == func_calls_.end() ? nullptr : it->second;
  }

 private:
  std::unordered_map<const SysYParser::LValContext*, BoundDecl> var_uses_;
  std::unordered_map<const SysYParser::VarDefContext*, TypeDesc> var_types_;
  std::unordered_map<const SysYParser::ConstDefContext*, TypeDesc> const_types_;
  std::unordered_map<const SysYParser::FuncFParamContext*, TypeDesc> param_types_;
  std::unordered_map<const SysYParser::FuncDefContext*, FuncTypeDesc> func_types_;
  std::unordered_map<const SysYParser::UnaryExpContext*, SysYParser::FuncDefContext*>
      func_calls_;
};

SemanticContext RunSema(SysYParser::CompUnitContext& comp_unit);