// 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) {
  node->type_id_ = namespace_visitor_.FindType({}, node->type);

  if (node->name.has_value() && !node->type_id_.has_value()) {
    error_handling::HandleTypecheckError("Variable namespace type not found");
  }

  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]);
  }
  Visitor::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) {
    Visitor::Visit(parameter.get());
    abstract_types_.DefineType(parameter->type, parameter->type_graph_id_);
  }
  Visitor::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) {
  Visitor::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) {
    Visitor::Visit(parameter.get());
    abstract_types_.DefineType(parameter->type, parameter->type_graph_id_);
  }
}

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

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

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

// Operators

// Other Expressions

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

  ////////////
  node->argument_graph_ids_.resize(node->arguments.size());
  for (size_t i = 0; i < node->arguments.size(); ++i) {
    node->argument_graph_ids_[i] = namespace_visitor_.GetAbstractTypeGraph()->AddVertex();
  }

  node->return_type_graph_id_ = namespace_visitor_.GetAbstractTypeGraph()->AddVertex();
  ///////////

  Visitor::Visit(node->expression);
  abstract_types_.ExitContext();
}

// Name

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

// Type

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");
  }
}

// Typeclass

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