#include "sem/Sema.h"

#include <any>
#include <stdexcept>
#include <string>

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

namespace {

class SemaVisitor final : public SysYBaseVisitor {
 public:
  explicit SemaVisitor() {
    // 预填标准库函数
    AddBuiltin("getint", Symbol::Kind::Function);
    AddBuiltin("getch", Symbol::Kind::Function);
    AddBuiltin("getfloat", Symbol::Kind::Function);
    AddBuiltin("getarray", Symbol::Kind::Function);
    AddBuiltin("getfarray", Symbol::Kind::Function);
    AddBuiltin("putint", Symbol::Kind::Function);
    AddBuiltin("putch", Symbol::Kind::Function);
    AddBuiltin("putfloat", Symbol::Kind::Function);
    AddBuiltin("putarray", Symbol::Kind::Function);
    AddBuiltin("putfarray", Symbol::Kind::Function);
    AddBuiltin("starttime", Symbol::Kind::Function);
    AddBuiltin("stoptime", Symbol::Kind::Function);
    AddBuiltin("putf", Symbol::Kind::Function);
  }

 private:
  void AddBuiltin(const std::string& name, Symbol::Kind kind) {
    Symbol sym;
    sym.kind = kind;
    sym.def_ctx = nullptr; // 内建函数没有语法树节点
    table_.Add(name, sym);
  }

  std::any visitCompUnit(SysYParser::CompUnitContext* ctx) override {
    if (!ctx) return {};
    for (auto* child : ctx->children) {
      if (auto* decl = dynamic_cast<SysYParser::DeclContext*>(child)) {
        decl->accept(this);
      } else if (auto* funcDef = dynamic_cast<SysYParser::FuncDefContext*>(child)) {
        funcDef->accept(this);
      }
    }
    return {};
  }

  std::any visitDecl(SysYParser::DeclContext* ctx) override {
    if (ctx->constDecl()) return ctx->constDecl()->accept(this);
    if (ctx->varDecl()) return ctx->varDecl()->accept(this);
    return {};
  }

  std::any visitConstDecl(SysYParser::ConstDeclContext* ctx) override {
    for (auto* def : ctx->constDef()) {
      const std::string name = def->ID()->getText();
      if (table_.IsInCurrentScope(name)) {
        throw std::runtime_error(FormatError("sema", "重复定义常量: " + name));
      }
      Symbol sym;
      sym.kind = Symbol::Kind::Constant;
      sym.def_ctx = def;
      sym.is_const = true;
      sym.is_array = !def->exp().empty();
      table_.Add(name, sym);

      for (auto* exp : def->exp()) exp->accept(this);
      def->initValue()->accept(this);
    }
    return {};
  }

  std::any visitVarDecl(SysYParser::VarDeclContext* ctx) override {
    for (auto* def : ctx->varDef()) {
      const std::string name = def->ID()->getText();
      if (table_.IsInCurrentScope(name)) {
        throw std::runtime_error(FormatError("sema", "重复定义变量: " + name));
      }
      Symbol sym;
      sym.kind = Symbol::Kind::Variable;
      sym.def_ctx = def;
      sym.is_const = false;
      sym.is_array = !def->exp().empty();
      table_.Add(name, sym);

      for (auto* exp : def->exp()) exp->accept(this);
      if (def->initValue()) def->initValue()->accept(this);
    }
    return {};
  }

  std::any visitFuncDef(SysYParser::FuncDefContext* ctx) override {
    const std::string name = ctx->ID()->getText();
    if (table_.IsInCurrentScope(name)) {
      throw std::runtime_error(FormatError("sema", "重复定义函数: " + name));
    }
    Symbol sym;
    sym.kind = Symbol::Kind::Function;
    sym.def_ctx = ctx;
    table_.Add(name, sym);

    table_.PushScope();
    if (ctx->funcFParams()) {
      ctx->funcFParams()->accept(this);
    }
    if (ctx->blockStmt()) {
      // Visit block items without pushing another scope to keep params in same scope
      for (auto* item : ctx->blockStmt()->blockItem()) {
        item->accept(this);
      }
    }
    table_.PopScope();
    return {};
  }

  std::any visitFuncFParam(SysYParser::FuncFParamContext* ctx) override {
    const std::string name = ctx->ID()->getText();
    if (table_.IsInCurrentScope(name)) {
      throw std::runtime_error(FormatError("sema", "函数参数名冲突: " + name));
    }
    Symbol sym;
    sym.kind = Symbol::Kind::Parameter;
    sym.def_ctx = ctx;
    sym.is_array = !ctx->LBRACK().empty();
    table_.Add(name, sym);

    for (auto* exp : ctx->exp()) exp->accept(this);
    return {};
  }

  std::any visitBlockStmt(SysYParser::BlockStmtContext* ctx) override {
    table_.PushScope();
    for (auto* item : ctx->blockItem()) {
      item->accept(this);
    }
    table_.PopScope();
    return {};
  }

  std::any visitAssignStmt(SysYParser::AssignStmtContext* ctx) override {
    ctx->lValue()->accept(this);
    const std::string name = ctx->lValue()->ID()->getText();
    Symbol* sym = table_.Lookup(name);
    if (sym && sym->is_const) {
      throw std::runtime_error(FormatError("sema", "试图给常量赋值: " + name));
    }
    ctx->exp()->accept(this);
    return {};
  }

  std::any visitLValue(SysYParser::LValueContext* ctx) override {
    const std::string name = ctx->ID()->getText();
    Symbol* sym = table_.Lookup(name);
    if (!sym) {
      throw std::runtime_error(FormatError("sema", "使用了未定义的标识符: " + name));
    }
    if (sym->kind == Symbol::Kind::Function) {
      throw std::runtime_error(FormatError("sema", "函数名不能作为左值: " + name));
    }
    sema_.BindLValue(ctx, sym->def_ctx);
    for (auto* exp : ctx->exp()) exp->accept(this);
    return {};
  }

  std::any visitFuncCallExp(SysYParser::FuncCallExpContext* ctx) override {
    const std::string name = ctx->ID()->getText();
    Symbol* sym = table_.Lookup(name);
    if (!sym) {
      throw std::runtime_error(FormatError("sema", "调用未定义的函数: " + name));
    }
    if (sym->kind != Symbol::Kind::Function) {
      throw std::runtime_error(FormatError("sema", "标识符不是函数: " + name));
    }
    sema_.BindFuncCall(ctx, dynamic_cast<SysYParser::FuncDefContext*>(sym->def_ctx));
    if (ctx->funcRParams()) {
      ctx->funcRParams()->accept(this);
    }
    return {};
  }

  // Visit expressions to ensure all sub-expressions are checked (e.g. for variable uses)
  std::any visitParenExp(SysYParser::ParenExpContext* ctx) override {
    return ctx->exp()->accept(this);
  }

  std::any visitLValueExp(SysYParser::LValueExpContext* ctx) override {
    return ctx->lValue()->accept(this);
  }

  std::any visitNumberExp(SysYParser::NumberExpContext* ctx) override {
    return ctx->number()->accept(this);
  }

  std::any visitNumber(SysYParser::NumberContext* ctx) override {
    return {};
  }

  std::any visitNotExp(SysYParser::NotExpContext* ctx) override {
    return ctx->exp()->accept(this);
  }

  std::any visitUnaryAddExp(SysYParser::UnaryAddExpContext* ctx) override {
    return ctx->exp()->accept(this);
  }

  std::any visitUnarySubExp(SysYParser::UnarySubExpContext* ctx) override {
    return ctx->exp()->accept(this);
  }

  std::any visitMulDivModExp(SysYParser::MulDivModExpContext* ctx) override {
    ctx->exp(0)->accept(this);
    ctx->exp(1)->accept(this);
    return {};
  }

  std::any visitAddSubExp(SysYParser::AddSubExpContext* ctx) override {
    ctx->exp(0)->accept(this);
    ctx->exp(1)->accept(this);
    return {};
  }

  std::any visitRelExp(SysYParser::RelExpContext* ctx) override {
    ctx->exp(0)->accept(this);
    ctx->exp(1)->accept(this);
    return {};
  }

  std::any visitEqNeExp(SysYParser::EqNeExpContext* ctx) override {
    ctx->exp(0)->accept(this);
    ctx->exp(1)->accept(this);
    return {};
  }

  std::any visitAndExp(SysYParser::AndExpContext* ctx) override {
    ctx->exp(0)->accept(this);
    ctx->exp(1)->accept(this);
    return {};
  }

  std::any visitOrExp(SysYParser::OrExpContext* ctx) override {
    ctx->exp(0)->accept(this);
    ctx->exp(1)->accept(this);
    return {};
  }

  std::any visitReturnStmt(SysYParser::ReturnStmtContext* ctx) override {
    if (ctx->exp()) ctx->exp()->accept(this);
    return {};
  }

  std::any visitIfStmt(SysYParser::IfStmtContext* ctx) override {
    ctx->exp()->accept(this);
    ctx->stmt(0)->accept(this);
    if (ctx->stmt(1)) ctx->stmt(1)->accept(this);
    return {};
  }

  std::any visitWhileStmt(SysYParser::WhileStmtContext* ctx) override {
    ctx->exp()->accept(this);
    ctx->stmt()->accept(this);
    return {};
  }

  std::any visitBreakStmt(SysYParser::BreakStmtContext* ctx) override {
    return {};
  }

  std::any visitContinueStmt(SysYParser::ContinueStmtContext* ctx) override {
    return {};
  }

  std::any visitExpStmt(SysYParser::ExpStmtContext* ctx) override {
    if (ctx->exp()) ctx->exp()->accept(this);
    return {};
  }
public:
 SemanticContext TakeSemanticContext() { return std::move(sema_); }

private:
  SymbolTable table_;
  SemanticContext sema_;
};

}  // namespace

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