#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 param_list = function->parameterList()) {
      function_arity[name] = param_list->Identifier().size();
    } else {
      function_arity[name] = 0;
    }
  }
  visitChildren(ctx);
  return {};
}

std::any TypeCheckVisitor::visitFunction(xlangParser::FunctionContext *ctx) {
  scope.emplace_back();
  if (auto param_list = ctx->parameterList()) {
    for (const auto param : param_list->Identifier()) {
      scope.back().emplace(param->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()) {
    visitValue(ctx->value());
    scope.back().emplace(identifier->getSymbol()->getText());
    return {};
  }
  visitChildren(ctx);
  return {};
}

std::any TypeCheckVisitor::visitVariable(xlangParser::VariableContext *ctx) {
  auto token = ctx->Identifier()->getSymbol();
  auto name = token->getText();
  auto line = token->getLine();
  auto charPositionInLine = token->getCharPositionInLine();

  if (ctx->LeftParen()) {
    auto it = function_arity.find(name);

    if (it == function_arity.end()) {
      errorlistener.typeError(line, charPositionInLine,
                              "unknown function '" + name + "'");
    } else {
      auto arity = it->second;

      if (auto arg_list = ctx->argumentList()) {
        auto arg_num = arg_list->value().size();

        if (arity != arg_num) {
          errorlistener.typeError(
              line, charPositionInLine,
              "function '" + name + "' expects " + std::to_string(arity) +
                  " arguments, but got " + std::to_string(arg_num));
        }
      } else {
        if (arity) {
          errorlistener.typeError(line, charPositionInLine,
                                  "function '" + name + "' expects " +
                                      std::to_string(arity) +
                                      " arguments, but got none");
        }
      }
    }
    visitChildren(ctx);
  } else {
    if (!inScope(name)) {
      errorlistener.typeError(line, charPositionInLine,
                              "variable '" + name + "' is not in scope");
    }
  }
  return {};
}

}  // namespace xlang