// for clangd
#include "../include/link_symbols_visitor.hpp"
#include "../include/error_handling.hpp"

namespace interpreter {

// Sources -----------------

void LinkSymbolsVisitor::Visit(SourceFile* node) {
  for (auto& statement : node->statements) {
    if (std::holds_alternative<Partition>(statement)) {
      Visit(&std::get<Partition>(statement));
    } else if (std::holds_alternative<SourceStatement>(statement)) {
      Visitor::Visit(std::get<SourceStatement>(statement));
    } else {
      // error
    }
  }
}

void LinkSymbolsVisitor::Visit(Sources* node) {
  abstract_types_.EnterContext();
  for (auto& statement : node->statements) {
    Visitor::Visit(statement);
  }
  abstract_types_.ExitContext();
}

// Namespaces, partitions -----------------

void LinkSymbolsVisitor::Visit(Partition* node) {
  // TODO
  Visit(node->scope.get());
}

void LinkSymbolsVisitor::Visit(Namespace* node) {
  namespace_visitor_.EnterNamespace(node->type);
  Visit(node->scope.get());
  namespace_visitor_.ExitNamespace();
}

// Definitions -----------------

void LinkSymbolsVisitor::Visit(ImportStatement* node) {}

void LinkSymbolsVisitor::Visit(AliasDefinitionStatement* node) {
  abstract_types_.EnterContext();
  for (size_t i = 0; i <node->parameters.size(); ++i) {
    abstract_types_.DefineType(node->parameters[i], node->parameter_graph_ids_[i]);
  }
  Visit(node->value.get());
  abstract_types_.ExitContext();
}

void LinkSymbolsVisitor::Visit(VariableDefinitionStatement* node) {
  abstract_types_.EnterContext();
  Visitor::Visit(node->name);
  Visitor::Visit(node->value);
  abstract_types_.ExitContext();
}

void LinkSymbolsVisitor::Visit(FunctionDeclaration* node) {
  abstract_types_.EnterContext();
  for (auto& parameter : node->parameters) {
    Visit(parameter.get());
    abstract_types_.DefineType(parameter->type, parameter->type_graph_id_);
  }
  Visit(node->type.get());
  abstract_types_.ExitContext();
}

void LinkSymbolsVisitor::Visit(FunctionDefinitionStatement* node) {
  abstract_types_.EnterContext();
  Visit(node->definition.get());
  Visitor::Visit(node->value);
  abstract_types_.ExitContext();
}

void LinkSymbolsVisitor::Visit(TypeDefinitionStatement* node) {
  abstract_types_.EnterContext();
  Visit(node->definition.get());
  Visitor::Visit(node->value);
  abstract_types_.ExitContext();
}

void LinkSymbolsVisitor::Visit(AbstractTypeDefinitionStatement* node) {
  Visit(node->type.get());
  // TODO: can't be used before definition
  abstract_types_.DefineType(node->type->type, node->type->type_graph_id_);
}

void LinkSymbolsVisitor::Visit(TypeclassDefinitionStatement* node) {
  abstract_types_.EnterContext();
  Visit(node->definition.get());
  for (auto& requirement : node->requirements) {
    Visit(requirement.get());
  }
  abstract_types_.ExitContext();
}

// Definition parts

void LinkSymbolsVisitor::Visit(FunctionDefinition* node) {
  for (auto& parameter : node->parameters) {
    Visit(parameter.get());
    abstract_types_.DefineType(parameter->type, parameter->type_graph_id_);
  }
}

void LinkSymbolsVisitor::Visit(TypeDefinition* node) {
  Visit(node->type.get());
  for (auto& parameter : node->parameters) {
    Visit(parameter.get());
    abstract_types_.DefineType(parameter->type, parameter->type_graph_id_);
  }
}

void LinkSymbolsVisitor::Visit(AnyAnnotatedType* node) {
  // TODO check ??
  if (!node->typeclasses.empty() > 0) {
    for (auto& typeclass : node->typeclasses) {
      Visitor::Visit(typeclass);
    }
  }
}

// Flow control -----------------

void LinkSymbolsVisitor::Visit(MatchCase* node) {
  Visitor::Visit(node->value);
  if (node->condition.has_value()) {
    Visitor::Visit(node->condition.value());
  }
  if (node->statement.has_value()) {
    Visitor::Visit(node->statement.value());
  }
}

void LinkSymbolsVisitor::Visit(Match* node) {
  Visitor::Visit(node->value);
  for (auto& match_case : node->matches) {
    Visit(&match_case);
  }
}

void LinkSymbolsVisitor::Visit(Condition* node) {
  Visitor::Visit(node->conditions[0]);
  Visitor::Visit(node->statements[0]);
  for (size_t i = 1; i < node->conditions.size(); ++i) {
      Visitor::Visit(node->conditions[i]);
      Visitor::Visit(node->statements[i]);
  }
  if (node->statements.size() > node->conditions.size()) {
    Visitor::Visit(node->statements[node->conditions.size()]);
  }
}

void LinkSymbolsVisitor::Visit(DoWhileLoop* node) {
  Visitor::Visit(node->statement);
  Visitor::Visit(node->condition);
}

void LinkSymbolsVisitor::Visit(WhileLoop* node) {
  Visitor::Visit(node->condition);
  Visitor::Visit(node->statement);
}
void LinkSymbolsVisitor::Visit(ForLoop* node) {
  Visitor::Visit(node->variable);
  Visitor::Visit(node->interval);
  Visitor::Visit(node->statement);
}

void LinkSymbolsVisitor::Visit(LoopLoop* node) {
  Visitor::Visit(node->statement);
}

// Statements, expressions, blocks, etc. -----------------

void LinkSymbolsVisitor::Visit(Block* node) {
  for (auto& statement : node->statements) {
    Visitor::Visit(statement);
  }
}

void LinkSymbolsVisitor::Visit(ScopedStatement* node) {
  Visitor::Visit(node->statement);
}

void LinkSymbolsVisitor::Visit(LoopControlExpression& node) {}

// Operators

void LinkSymbolsVisitor::Visit(BinaryOperatorExpression* node) {
  Visitor::Visit(node->left_expression);
  Visitor::Visit(node->right_expression);
}

void LinkSymbolsVisitor::Visit(UnaryOperatorExpression* node) {
  Visitor::Visit(node->expression);
}

void LinkSymbolsVisitor::Visit(ReferenceExpression* node) {
  Visit(node->expression.get());
}

// Other Expressions

void LinkSymbolsVisitor::Visit(FunctionCallExpression* node) {
  Visit(node->name.get());
  for (auto& argument : node->arguments) {
    Visitor::Visit(argument);
  }
}

void LinkSymbolsVisitor::Visit(TupleExpression* node) {
  for (auto& expression : node->expressions) {
    Visitor::Visit(expression);
  }
}

void LinkSymbolsVisitor::Visit(VariantExpression* node) {
  for (auto& expression : node->expressions) {
    Visitor::Visit(expression);
  }
}

void LinkSymbolsVisitor::Visit(ReturnExpression* node) {
  Visitor::Visit(node->expression);
}

void LinkSymbolsVisitor::Visit(TypeConstructor* node) {
  Visit(node->type.get());
  for (auto& parameter : node->parameters) {
    Visitor::Visit(std::get<2>(parameter));
  }
}

void LinkSymbolsVisitor::Visit(LambdaFunction* node) {
  abstract_types_.EnterContext();
  for (auto& parameter : node->parameters) {
    Visit(parameter.get());
    abstract_types_.DefineType(parameter->type, parameter->type_graph_id_);
  }
  Visitor::Visit(node->expression);
  abstract_types_.ExitContext();
}

void LinkSymbolsVisitor::Visit(ArrayExpression* node) {
  for (auto& element : node->elements) {
    Visitor::Visit(element);
  }
}

// Name

void LinkSymbolsVisitor::Visit(NameExpression* node) {
  for (auto& variable_namespace : node->namespaces) {
    Visitor::Visit(variable_namespace);
  }
  for (auto& expression : node->expressions) {
    Visitor::Visit(expression);
  }
}

void LinkSymbolsVisitor::Visit(TupleName* node) {
  for (auto& name : node->names) {
    Visitor::Visit(name);
  }
}

void LinkSymbolsVisitor::Visit(VariantName* node) {
  for (auto& name : node->names) {
    Visitor::Visit(name);
  }
}

void LinkSymbolsVisitor::Visit(AnnotatedName* node) {
  if (node->type.has_value()) {
    Visitor::Visit(node->type.value());
  }
}

// Type, typeclass, etc. -----------------

// Type

void LinkSymbolsVisitor::Visit(FunctionType* node) {
  for (auto& type : node->types) {
    Visitor::Visit(type);
  }
}

void LinkSymbolsVisitor::Visit(TupleType* node) {
  for (auto& entity : node->entities) {
    Visit(entity.second.get());
  }
}

void LinkSymbolsVisitor::Visit(VariantType* node) {
  for (auto& constructor : node->constructors) {
    if (std::holds_alternative<Constructor>(constructor)) {
      // do nothing
    } else if (std::holds_alternative<std::unique_ptr<TupleType>>(constructor)) {
      Visit(std::get<std::unique_ptr<TupleType>>(constructor).get());
    } else {
      // error
    }
  }
}

void LinkSymbolsVisitor::Visit(ParametrizedType* node) {
  Visit(node->type_expression.get());
  for (auto& parameter : node->parameters) {
    Visitor::Visit(parameter);
  }
}

void LinkSymbolsVisitor::Visit(TypeExpression* node) {
  std::vector<std::string> path;
  path.reserve(node->namespaces.size());

  for (auto& type_namespace : node->namespaces) {
    Visitor::Visit(type_namespace);
    if (std::holds_alternative<std::unique_ptr<std::string>>(type_namespace)) {
      path.push_back(*std::get<std::unique_ptr<std::string>>(type_namespace));
    } else if (std::holds_alternative<std::unique_ptr<ParametrizedType>>(type_namespace)) {
      path.push_back(std::get<std::unique_ptr<ParametrizedType>>(type_namespace)->type_expression->type);
    } else {
      // error
    }
  }

  std::optional<utils::IdType> maybe_type = namespace_visitor_.FindType(path, node->type);
  std::optional<utils::IdType> maybe_abstract_type = std::nullopt;

  if (path.size() == 0) {
    maybe_abstract_type = abstract_types_.GetTypeId(node->type);
  }

  if (maybe_abstract_type.has_value()) {
    if (maybe_type.has_value()) {
      error_handling::HandleTypecheckError("Ambigious type");
    } else {
      node->type_id_ = maybe_abstract_type.value();
    }
  } else if (maybe_type.has_value()) {
    node->type_id_ = maybe_type.value();
  } else {
    error_handling::HandleTypecheckError("Type not found");
  }
}

void LinkSymbolsVisitor::Visit(ExtendedScopedAnyType* node) {
  Visitor::Visit(node->type);
}

// Typeclass

void LinkSymbolsVisitor::Visit(ParametrizedTypeclass* node) {
  Visit(node->typeclass_expression.get());
  for (auto& parameter : node->parameters) {
    Visitor::Visit(parameter);
  }
}

void LinkSymbolsVisitor::Visit(TypeclassExpression* node) {
  std::vector<std::string> path;
  path.reserve(node->namespaces.size());

  for (auto& typeclass_namespace : node->namespaces) {
    Visitor::Visit(typeclass_namespace);
    if (std::holds_alternative<std::unique_ptr<std::string>>(typeclass_namespace)) {
      path.push_back(*std::get<std::unique_ptr<std::string>>(typeclass_namespace));
    } else if (std::holds_alternative<std::unique_ptr<ParametrizedType>>(typeclass_namespace)) {
      path.push_back(std::get<std::unique_ptr<ParametrizedType>>(typeclass_namespace)->type_expression->type);
    } else {
      // error
    }
  }

  std::optional<utils::IdType> maybe_typeclass = namespace_visitor_.FindType(path, node->typeclass);
  std::optional<utils::IdType> maybe_abstract_typeclass = std::nullopt;

  if (path.size() == 0) {
    maybe_abstract_typeclass = abstract_types_.GetTypeId(node->typeclass);
  }

  if (maybe_abstract_typeclass.has_value()) {
    if (maybe_typeclass.has_value()) {
      error_handling::HandleTypecheckError("Ambigious type");
    } else {
      node->type_id_ = maybe_abstract_typeclass.value();
    }
  } else if (maybe_typeclass.has_value()) {
    node->type_id_ = maybe_typeclass.value();
  } else {
    error_handling::HandleTypecheckError("Type not found");
  }
}

} // namespace interpreter