#include "sem/Sema.h"

#include <any>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

#include "SysYBaseVisitor.h"
#include "sem/func.h"
#include "utils/Log.h"

using namespace sem;

namespace {

// ─── 作用域栈 ─────────────────────────────────────────────────────────────────
class ScopeStack {
 public:
  void PushScope() { scopes_.emplace_back(); }
  void PopScope() {
    if (!scopes_.empty()) scopes_.pop_back();
  }
  // 在当前作用域定义符号
  void Define(const std::string& name, antlr4::ParserRuleContext* decl) {
    if (scopes_.empty()) return;
    scopes_.back()[name] = decl;
  }
  // 检查当前作用域是否已定义（用于重复定义检查）
  bool ContainsInCurrent(const std::string& name) const {
    return !scopes_.empty() && scopes_.back().count(name) > 0;
  }
  // 向上查找
  antlr4::ParserRuleContext* Lookup(const std::string& name) const {
    for (auto it = scopes_.rbegin(); it != scopes_.rend(); ++it) {
      auto found = it->find(name);
      if (found != it->end()) return found->second;
    }
    return nullptr;
  }

 private:
  std::vector<std::unordered_map<std::string, antlr4::ParserRuleContext*>>
      scopes_;
};

// ─── 语义分析访问者 ────────────────────────────────────────────────────────────
class SemaVisitor final : public SysYBaseVisitor {
 public:
  // 顶层：global scope → decl (global var/const) → funcDef
  std::any visitCompUnit(SysYParser::CompUnitContext* ctx) override {
    if (!ctx) {
      throw std::runtime_error(FormatError("sema", "缺少编译单元"));
    }
    scope_.PushScope();  // global scope
    // 先处理全局声明
    for (auto* decl : ctx->decl()) {
      if (decl) decl->accept(this);
    }
    // 再处理函数定义
    bool has_main = false;
    for (auto* func : ctx->funcDef()) {
      if (!func) continue;
      if (func->Ident() && func->Ident()->getText() == "main") {
        has_main = true;
      }
      func->accept(this);
    }
    if (!has_main) {
      throw std::runtime_error(FormatError("sema", "缺少main函数定义"));
    }
    scope_.PopScope();
    return {};
  }

  // 函数定义：参数入作用域 → 函数体
  std::any visitFuncDef(SysYParser::FuncDefContext* ctx) override {
    if (!ctx) return {};
    scope_.PushScope();
    if (auto* params = ctx->funcFParams()) {
      for (auto* param : params->funcFParam()) {
        if (!param || !param->Ident()) continue;
        std::string name = param->Ident()->getText();
        if (scope_.ContainsInCurrent(name)) {
          throw std::runtime_error(
              FormatError("sema", "参数重复定义: " + name));
        }
        scope_.Define(name, param);
      }
    }
    if (ctx->block()) ctx->block()->accept(this);
    scope_.PopScope();
    return {};
  }

  // 块：新作用域
  std::any visitBlock(SysYParser::BlockContext* ctx) override {
    if (!ctx) return {};
    scope_.PushScope();
    for (auto* item : ctx->blockItem()) {
      if (item) item->accept(this);
    }
    scope_.PopScope();
    return {};
  }

  // 块内项：分发
  std::any visitBlockItem(SysYParser::BlockItemContext* ctx) override {
    if (!ctx) return {};
    if (ctx->decl()) ctx->decl()->accept(this);
    if (ctx->stmt()) ctx->stmt()->accept(this);
    return {};
  }

  // 声明：分发到 varDecl / constDecl
  std::any visitDecl(SysYParser::DeclContext* ctx) override {
    if (!ctx) return {};
    if (ctx->varDecl()) ctx->varDecl()->accept(this);
    if (ctx->constDecl()) ctx->constDecl()->accept(this);
    return {};
  }

  // 变量声明
  std::any visitVarDecl(SysYParser::VarDeclContext* ctx) override {
    if (!ctx) return {};
    for (auto* varDef : ctx->varDef()) {
      if (!varDef || !varDef->Ident()) continue;
      std::string name = varDef->Ident()->getText();
      if (scope_.ContainsInCurrent(name)) {
        throw std::runtime_error(
            FormatError("sema", "变量重复定义: " + name));
      }
      scope_.Define(name, varDef);
      // 访问初始化表达式中的变量引用
      if (varDef->initVal()) varDef->initVal()->accept(this);
    }
    return {};
  }

  // 常量声明
  std::any visitConstDecl(SysYParser::ConstDeclContext* ctx) override {
    if (!ctx) return {};
    for (auto* constDef : ctx->constDef()) {
      if (!constDef || !constDef->Ident()) continue;
      std::string name = constDef->Ident()->getText();
      if (scope_.ContainsInCurrent(name)) {
        throw std::runtime_error(
            FormatError("sema", "常量重复定义: " + name));
      }
      scope_.Define(name, constDef);
      // 初始化表达式中的引用
      if (constDef->constInitVal()) constDef->constInitVal()->accept(this);
    }
    return {};
  }

  // 语句：分发各种形式
  std::any visitStmt(SysYParser::StmtContext* ctx) override {
    if (!ctx) return {};
    // 赋值语句：lVar = exp;
    if (ctx->lVar() && ctx->Assign()) {
      ctx->lVar()->accept(this);
      if (ctx->exp()) ctx->exp()->accept(this);
      return {};
    }
    // 表达式语句：(exp)?;
    if (ctx->exp() && !ctx->Return() && !ctx->If() && !ctx->While()) {
      ctx->exp()->accept(this);
      return {};
    }
    // 块语句
    if (ctx->block()) {
      ctx->block()->accept(this);
      return {};
    }
    // if 语句
    if (ctx->If()) {
      if (ctx->cond()) ctx->cond()->accept(this);
      auto stmts = ctx->stmt();
      for (auto* s : stmts) {
        if (s) s->accept(this);
      }
      return {};
    }
    // while 语句
    if (ctx->While()) {
      if (ctx->cond()) ctx->cond()->accept(this);
      auto stmts = ctx->stmt();
      for (auto* s : stmts) {
        if (s) s->accept(this);
      }
      return {};
    }
    // return 语句
    if (ctx->Return()) {
      if (ctx->exp()) ctx->exp()->accept(this);
      return {};
    }
    // break / continue：无变量引用
    return {};
  }

  // lVar：绑定变量使用
  std::any visitLVar(SysYParser::LVarContext* ctx) override {
    if (!ctx || !ctx->Ident()) return {};
    std::string name = ctx->Ident()->getText();
    auto* decl = scope_.Lookup(name);
    if (!decl) {
      // 可能是外部函数或未声明变量，不强制报错（IRGen 会处理外部调用）
      // 但对变量使用报错
      throw std::runtime_error(
          FormatError("sema", "未声明变量: " + name));
    }
    sema_.BindVarUse(ctx->Ident(), decl);
    // 下标表达式也需要访问
    for (auto* e : ctx->exp()) {
      if (e) e->accept(this);
    }
    return {};
  }

  // 表达式：通过 visitChildren 自动递归
  std::any visitExp(SysYParser::ExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitAddExp(SysYParser::AddExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitMulExp(SysYParser::MulExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitUnaryExp(SysYParser::UnaryExpContext* ctx) override {
    // 函数调用：Ident L_PAREN ...
    if (ctx->Ident() && ctx->L_PAREN()) {
      if (ctx->funcRParams()) ctx->funcRParams()->accept(this);
      return {};
    }
    return visitChildren(ctx);
  }
  std::any visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitCond(SysYParser::CondContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitLOrExp(SysYParser::LOrExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitLAndExp(SysYParser::LAndExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitRelExp(SysYParser::RelExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitEqExp(SysYParser::EqExpContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitFuncRParams(SysYParser::FuncRParamsContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitInitVal(SysYParser::InitValContext* ctx) override {
    return visitChildren(ctx);
  }
  std::any visitConstInitVal(SysYParser::ConstInitValContext* ctx) override {
    // 常量初始化器中通常没有变量引用（只有字面量），
    // 但如果有 constExp 引用其他常量则需要访问
    return visitChildren(ctx);
  }
  std::any visitConstExp(SysYParser::ConstExpContext* ctx) override {
    return visitChildren(ctx);
  }

  SemanticContext TakeSemanticContext() { return std::move(sema_); }

 private:
  ScopeStack scope_;
  SemanticContext sema_;
};

}  // namespace

SemanticContext RunSema(SysYParser::CompUnitContext& comp_unit) {
  SemaVisitor visitor;
  comp_unit.accept(&visitor);
  return visitor.TakeSemanticContext();
}