#include "sem/Sema.h"

#include <stdexcept>

namespace {

SymbolType ParseType(const std::string& text) {
  if (text == "int") {
    return SymbolType::TYPE_INT;
  }
  if (text == "float") {
    return SymbolType::TYPE_FLOAT;
  }
  if (text == "void") {
    return SymbolType::TYPE_VOID;
  }
  return SymbolType::TYPE_UNKNOWN;
}

SymbolType MergeNumericType(SymbolType lhs, SymbolType rhs) {
  if (lhs == SymbolType::TYPE_FLOAT || rhs == SymbolType::TYPE_FLOAT) {
    return SymbolType::TYPE_FLOAT;
  }
  if (lhs == SymbolType::TYPE_INT && rhs == SymbolType::TYPE_INT) {
    return SymbolType::TYPE_INT;
  }
  if (lhs != SymbolType::TYPE_UNKNOWN) {
    return lhs;
  }
  return rhs;
}

}  // namespace

void SemaVisitor::RecordNodeError(antlr4::ParserRuleContext* ctx,
                                  const std::string& msg) {
  if (!ctx || !ctx->getStart()) {
    ir_ctx_.RecordError(ErrorMsg(msg, 0, 0));
    return;
  }
  ir_ctx_.RecordError(ErrorMsg(msg, ctx->getStart()->getLine(),
                               ctx->getStart()->getCharPositionInLine() + 1));
}

std::any SemaVisitor::visitCompUnit(SysYParser::CompUnitContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitDecl(SysYParser::DeclContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitConstDecl(SysYParser::ConstDeclContext* ctx) {
  current_decl_is_const_ = true;
  current_decl_type_ = SymbolType::TYPE_UNKNOWN;
  if (ctx && ctx->btype()) {
    current_decl_type_ = ParseType(ctx->btype()->getText());
  }
  std::any result = visitChildren(ctx);
  current_decl_is_const_ = false;
  current_decl_type_ = SymbolType::TYPE_UNKNOWN;
  return result;
}

std::any SemaVisitor::visitBtype(SysYParser::BtypeContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitConstDef(SysYParser::ConstDefContext* ctx) {
  if (!ctx || !ctx->ID()) {
    return {};
  }

  const std::string name = ctx->ID()->getText();
  auto& table = ir_ctx_.GetSymbolTable();
  if (table.CurrentScopeHasVar(name)) {
    RecordNodeError(ctx, "重复定义变量: " + name);
  } else {
    VarInfo info;
    info.type = current_decl_type_;
    info.is_const = true;
    info.decl_ctx = ctx;
    table.BindVar(name, info, ctx);
  }

  ir_ctx_.SetType(ctx, current_decl_type_);
  return visitChildren(ctx);
}

std::any SemaVisitor::visitConstInitValue(SysYParser::ConstInitValueContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitVarDecl(SysYParser::VarDeclContext* ctx) {
  current_decl_is_const_ = false;
  current_decl_type_ = SymbolType::TYPE_UNKNOWN;
  if (ctx && ctx->btype()) {
    current_decl_type_ = ParseType(ctx->btype()->getText());
  }
  std::any result = visitChildren(ctx);
  current_decl_type_ = SymbolType::TYPE_UNKNOWN;
  return result;
}

std::any SemaVisitor::visitVarDef(SysYParser::VarDefContext* ctx) {
  if (!ctx || !ctx->ID()) {
    return {};
  }

  const std::string name = ctx->ID()->getText();
  auto& table = ir_ctx_.GetSymbolTable();
  if (table.CurrentScopeHasVar(name)) {
    RecordNodeError(ctx, "重复定义变量: " + name);
  } else {
    VarInfo info;
    info.type = current_decl_type_;
    info.is_const = current_decl_is_const_;
    info.decl_ctx = ctx;
    table.BindVar(name, info, ctx);
  }

  ir_ctx_.SetType(ctx, current_decl_type_);
  return visitChildren(ctx);
}

std::any SemaVisitor::visitInitValue(SysYParser::InitValueContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitFuncDef(SysYParser::FuncDefContext* ctx) {
  if (!ctx || !ctx->ID() || !ctx->funcType()) {
    return {};
  }

  const std::string func_name = ctx->ID()->getText();
  SymbolType ret_type = ParseType(ctx->funcType()->getText());
  ir_ctx_.SetCurrentFuncReturnType(ret_type);

  auto& table = ir_ctx_.GetSymbolTable();
  if (table.CurrentScopeHasFunc(func_name)) {
    RecordNodeError(ctx, "重复定义函数: " + func_name);
  } else {
    FuncInfo info;
    info.name = func_name;
    info.ret_type = ret_type;
    info.decl_ctx = ctx;
    if (ctx->funcFParams()) {
      for (auto* param : ctx->funcFParams()->funcFParam()) {
        if (!param || !param->btype()) {
          info.param_types.push_back(SymbolType::TYPE_UNKNOWN);
        } else {
          info.param_types.push_back(ParseType(param->btype()->getText()));
        }
      }
    }
    table.BindFunc(func_name, info, ctx);
  }

  ir_ctx_.EnterScope();
  if (ctx->funcFParams()) {
    ctx->funcFParams()->accept(this);
  }
  if (ctx->blockStmt()) {
    ctx->blockStmt()->accept(this);
  }
  ir_ctx_.LeaveScope();

  ir_ctx_.SetCurrentFuncReturnType(SymbolType::TYPE_UNKNOWN);
  return {};
}

std::any SemaVisitor::visitFuncType(SysYParser::FuncTypeContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitFuncFParams(SysYParser::FuncFParamsContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitFuncFParam(SysYParser::FuncFParamContext* ctx) {
  if (!ctx || !ctx->ID() || !ctx->btype()) {
    return {};
  }
  const std::string name = ctx->ID()->getText();
  auto& table = ir_ctx_.GetSymbolTable();
  if (table.CurrentScopeHasVar(name)) {
    RecordNodeError(ctx, "重复定义形参: " + name);
    return {};
  }

  VarInfo info;
  info.type = ParseType(ctx->btype()->getText());
  info.is_const = false;
  info.decl_ctx = ctx;
  table.BindVar(name, info, ctx);
  ir_ctx_.SetType(ctx, info.type);
  return {};
}

std::any SemaVisitor::visitBlockStmt(SysYParser::BlockStmtContext* ctx) {
  ir_ctx_.EnterScope();
  std::any result = visitChildren(ctx);
  ir_ctx_.LeaveScope();
  return result;
}

std::any SemaVisitor::visitBlockItem(SysYParser::BlockItemContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitStmt(SysYParser::StmtContext* ctx) {
  if (!ctx) {
    return {};
  }

  if (ctx->WHILE()) {
    ir_ctx_.EnterLoop();
    std::any result = visitChildren(ctx);
    ir_ctx_.ExitLoop();
    return result;
  }

  if (ctx->BREAK() && !ir_ctx_.InLoop()) {
    RecordNodeError(ctx, "break 只能出现在循环语句中");
  }

  if (ctx->CONTINUE() && !ir_ctx_.InLoop()) {
    RecordNodeError(ctx, "continue 只能出现在循环语句中");
  }

  if (ctx->lValue() && ctx->exp()) {
    ctx->lValue()->accept(this);
    ctx->exp()->accept(this);
    SymbolType lhs = ir_ctx_.GetType(ctx->lValue());
    SymbolType rhs = ir_ctx_.GetType(ctx->exp());
    if (!IsTypeCompatible(lhs, rhs)) {
      RecordNodeError(ctx, "赋值两侧类型不兼容");
    }
    return {};
  }

  return visitChildren(ctx);
}

std::any SemaVisitor::visitReturnStmt(SysYParser::ReturnStmtContext* ctx) {
  if (!ctx) {
    return {};
  }

  SymbolType ret_type = ir_ctx_.GetCurrentFuncReturnType();
  if (ctx->exp()) {
    ctx->exp()->accept(this);
    SymbolType expr_type = ir_ctx_.GetType(ctx->exp());
    if (ret_type == SymbolType::TYPE_VOID) {
      RecordNodeError(ctx, "void 函数不应返回表达式");
    } else if (!IsTypeCompatible(ret_type, expr_type)) {
      RecordNodeError(ctx, "return 表达式类型与函数返回类型不匹配");
    }
  } else if (ret_type != SymbolType::TYPE_VOID &&
             ret_type != SymbolType::TYPE_UNKNOWN) {
    RecordNodeError(ctx, "非 void 函数 return 必须带表达式");
  }

  return {};
}

std::any SemaVisitor::visitExp(SysYParser::ExpContext* ctx) {
  if (!ctx || !ctx->addExp()) {
    return {};
  }
  ctx->addExp()->accept(this);
  ir_ctx_.SetType(ctx, ir_ctx_.GetType(ctx->addExp()));
  return {};
}

std::any SemaVisitor::visitCond(SysYParser::CondContext* ctx) {
  if (!ctx || !ctx->lOrExp()) {
    return {};
  }
  ctx->lOrExp()->accept(this);
  ir_ctx_.SetType(ctx, SymbolType::TYPE_INT);
  return {};
}

std::any SemaVisitor::visitLValue(SysYParser::LValueContext* ctx) {
  if (!ctx || !ctx->ID()) {
    return {};
  }

  VarInfo var;
  void* decl_ctx = nullptr;
  auto& table = ir_ctx_.GetSymbolTable();
  const std::string name = ctx->ID()->getText();
  if (!table.LookupVar(name, var, decl_ctx)) {
    RecordNodeError(ctx, "未定义变量: " + name);
    ir_ctx_.SetType(ctx, SymbolType::TYPE_UNKNOWN);
    return {};
  }

  ir_ctx_.SetType(ctx, var.type);

  if (sema_ctx_ && decl_ctx) {
    auto* rule = static_cast<antlr4::ParserRuleContext*>(decl_ctx);
    if (auto* var_def = dynamic_cast<SysYParser::VarDefContext*>(rule)) {
      sema_ctx_->BindVarUse(ctx, var_def);
    }
  }

  return {};
}

std::any SemaVisitor::visitPrimaryExp(SysYParser::PrimaryExpContext* ctx) {
  if (!ctx) {
    return {};
  }

  if (ctx->exp()) {
    ctx->exp()->accept(this);
    ir_ctx_.SetType(ctx, ir_ctx_.GetType(ctx->exp()));
    return {};
  }

  if (ctx->lValue()) {
    ctx->lValue()->accept(this);
    ir_ctx_.SetType(ctx, ir_ctx_.GetType(ctx->lValue()));
    return {};
  }

  if (ctx->number()) {
    ctx->number()->accept(this);
    ir_ctx_.SetType(ctx, ir_ctx_.GetType(ctx->number()));
  }

  return {};
}

std::any SemaVisitor::visitNumber(SysYParser::NumberContext* ctx) {
  if (!ctx) {
    return {};
  }

  if (ctx->ILITERAL()) {
    ir_ctx_.SetType(ctx, SymbolType::TYPE_INT);
    ir_ctx_.SetConstVal(ctx, std::any(0L));
  } else if (ctx->FLITERAL()) {
    ir_ctx_.SetType(ctx, SymbolType::TYPE_FLOAT);
    ir_ctx_.SetConstVal(ctx, std::any(0.0));
  }

  return {};
}

std::any SemaVisitor::visitUnaryExp(SysYParser::UnaryExpContext* ctx) {
  if (!ctx) {
    return {};
  }

  if (ctx->primaryExp()) {
    ctx->primaryExp()->accept(this);
    ir_ctx_.SetType(ctx, ir_ctx_.GetType(ctx->primaryExp()));
    return {};
  }

  if (ctx->unaryOp() && ctx->unaryExp()) {
    ctx->unaryExp()->accept(this);
    ir_ctx_.SetType(ctx, ir_ctx_.GetType(ctx->unaryExp()));
    return {};
  }

  if (ctx->ID()) {
    FuncInfo fn;
    void* decl_ctx = nullptr;
    if (!ir_ctx_.GetSymbolTable().LookupFunc(ctx->ID()->getText(), fn, decl_ctx)) {
      RecordNodeError(ctx, "未定义函数: " + ctx->ID()->getText());
      ir_ctx_.SetType(ctx, SymbolType::TYPE_UNKNOWN);
    } else {
      ir_ctx_.SetType(ctx, fn.ret_type);
    }
  }

  return visitChildren(ctx);
}

std::any SemaVisitor::visitUnaryOp(SysYParser::UnaryOpContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitFuncRParams(SysYParser::FuncRParamsContext* ctx) {
  return visitChildren(ctx);
}

std::any SemaVisitor::visitMulExp(SysYParser::MulExpContext* ctx) {
  if (!ctx) {
    return {};
  }

  if (ctx->mulExp()) {
    ctx->mulExp()->accept(this);
  }
  if (ctx->unaryExp()) {
    ctx->unaryExp()->accept(this);
  }

  SymbolType lhs = ctx->mulExp() ? ir_ctx_.GetType(ctx->mulExp())
                                 : ir_ctx_.GetType(ctx->unaryExp());
  SymbolType rhs = ir_ctx_.GetType(ctx->unaryExp());
  ir_ctx_.SetType(ctx, MergeNumericType(lhs, rhs));
  return {};
}

std::any SemaVisitor::visitAddExp(SysYParser::AddExpContext* ctx) {
  if (!ctx) {
    return {};
  }

  if (ctx->addExp()) {
    ctx->addExp()->accept(this);
  }
  if (ctx->mulExp()) {
    ctx->mulExp()->accept(this);
  }

  SymbolType lhs = ctx->addExp() ? ir_ctx_.GetType(ctx->addExp())
                                 : ir_ctx_.GetType(ctx->mulExp());
  SymbolType rhs = ir_ctx_.GetType(ctx->mulExp());
  ir_ctx_.SetType(ctx, MergeNumericType(lhs, rhs));
  return {};
}

std::any SemaVisitor::visitRelExp(SysYParser::RelExpContext* ctx) {
  if (ctx) {
    visitChildren(ctx);
    ir_ctx_.SetType(ctx, SymbolType::TYPE_INT);
  }
  return {};
}

std::any SemaVisitor::visitEqExp(SysYParser::EqExpContext* ctx) {
  if (ctx) {
    visitChildren(ctx);
    ir_ctx_.SetType(ctx, SymbolType::TYPE_INT);
  }
  return {};
}

std::any SemaVisitor::visitLAndExp(SysYParser::LAndExpContext* ctx) {
  if (ctx) {
    visitChildren(ctx);
    ir_ctx_.SetType(ctx, SymbolType::TYPE_INT);
  }
  return {};
}

std::any SemaVisitor::visitLOrExp(SysYParser::LOrExpContext* ctx) {
  if (ctx) {
    visitChildren(ctx);
    ir_ctx_.SetType(ctx, SymbolType::TYPE_INT);
  }
  return {};
}

std::any SemaVisitor::visitConstExp(SysYParser::ConstExpContext* ctx) {
  if (!ctx || !ctx->addExp()) {
    return {};
  }
  ctx->addExp()->accept(this);
  ir_ctx_.SetType(ctx, ir_ctx_.GetType(ctx->addExp()));
  return {};
}

void RunSemanticAnalysis(SysYParser::CompUnitContext* ctx, IRGenContext& ir_ctx) {
  if (!ctx) {
    throw std::invalid_argument("CompUnitContext is null");
  }

  ir_ctx.EnterScope();
  SemaVisitor visitor(ir_ctx, nullptr);
  visitor.visit(ctx);
  ir_ctx.LeaveScope();
}

SemanticContext RunSema(SysYParser::CompUnitContext& comp_unit) {
  IRGenContext ctx;
  SemanticContext sema_ctx;
  ctx.EnterScope();
  SemaVisitor visitor(ctx, &sema_ctx);
  visitor.visit(&comp_unit);
  ctx.LeaveScope();
  return sema_ctx;
}
