#include <typecheck.hh>

namespace xlang {

bool TypeCheckVisitor::inScope(const std::string &name) {
  for (auto it = scope.end() - 1;; it--) {
    if (it->find(name) != it->end()) {
      return true;
    }
    if (it == scope.begin()) {
      break;
    }
  }
  return false;
}

TypeCheckVisitor::TypeCheckVisitor(ErrorListener &errorlistener)
    : errorlistener{errorlistener} {}

std::any TypeCheckVisitor::visitFile(xlangParser::FileContext *ctx) {
  for (auto function : ctx->function()) {
    auto token = function->Identifier()->getSymbol();
    auto name = token->getText();
    if (function_arity.find(name) != function_arity.end()) {
      errorlistener.typeError(token->getLine(), token->getCharPositionInLine(),
                              "duplicate function '" + name + "'");
      continue;
    }
    if (auto arg_list = function->argumentList()) {
      function_arity[name] = arg_list->Identifier().size();
    } else {
      function_arity[name] = 0;
    }
  }
  visitChildren(ctx);
  return {};
}

std::any TypeCheckVisitor::visitFunction(xlangParser::FunctionContext *ctx) {
  scope.emplace_back();
  if (auto arg_list = ctx->argumentList()) {
    for (const auto arg : arg_list->Identifier()) {
      scope.back().emplace(arg->getSymbol()->getText());
    }
  }
  visitBlock(ctx->block());
  scope.pop_back();
  return {};
}

std::any TypeCheckVisitor::visitBlock(xlangParser::BlockContext *ctx) {
  scope.emplace_back();
  visitChildren(ctx);
  scope.pop_back();
  return {};
}

std::any TypeCheckVisitor::visitStatement(xlangParser::StatementContext *ctx) {
  if (auto identifier = ctx->Identifier()) {
    visitExpr(ctx->expr());
    scope.back().emplace(identifier->getSymbol()->getText());
    return {};
  }
  visitChildren(ctx);
  return {};
}

std::any TypeCheckVisitor::visitFactor(xlangParser::FactorContext *ctx) {
  if (auto identifier = ctx->Identifier()) {
    auto token = identifier->getSymbol();
    auto name = token->getText();
    if (ctx->LeftParen()) {
      if (function_arity.find(name) != function_arity.end()) {
        if (auto expr_list = ctx->exprList()) {
          auto arity = function_arity[name];
          auto arg_num = expr_list->expr().size();
          if (arity != arg_num) {
            errorlistener.typeError(
                token->getLine(), token->getCharPositionInLine(),
                "function '" + name + "' expects " + std::to_string(arity) +
                    " arguments, but got " + std::to_string(arg_num));
          }
        } else {
          if (auto arity = function_arity[name]) {
            errorlistener.typeError(
                token->getLine(), token->getCharPositionInLine(),
                "function '" + name + "' expects " + std::to_string(arity) +
                    " arguments, but got none");
          }
        }
      } else {
        errorlistener.typeError(token->getLine(),
                                token->getCharPositionInLine(),
                                "unknown function '" + name + "'");
      }
      visitChildren(ctx);
    } else {
      if (!inScope(name)) {
        errorlistener.typeError(token->getLine(),
                                token->getCharPositionInLine(),
                                "variable '" + name + "' is not in scope");
      }
    }
    return {};
  }
  visitChildren(ctx);
  return {};
}

}  // namespace xlang