#include <variant>

// for clangd
#include "../include/global_info.hpp"
#include "../include/types.hpp"
#include "../include/error_handling.hpp"


namespace info {

void GlobalInfo::NamespaceVisitor::AddImport(definition::Import&& import_info,
                                             const std::optional<std::string>& name) {
  if (name.has_value()) {
    global_info_.usages_[name.value()] = std::move(import_info);
  } else {
    global_info_.imports_.push_back(std::move(import_info));
  }
}

void GlobalInfo::NamespaceVisitor::AddEnterNamespace(const std::string& name,
                                                     std::optional<utils::IsConstModifier> modifier,
                                                     const interpreter::tokens::BaseNode& base_node) {
  if (type::ToInternalType(name).has_value()) {
    error_handling::HandleTypecheckError("Can't define basic type namespace", base_node);
  }

  size_t id = global_info_.namespaces_.size();
  global_info_.namespaces_.emplace_back();
  definition::Namespace* namespace_info = &global_info_.namespaces_.back();

  if (modifier.has_value()) {
    if (modifier.value() == utils::IsConstModifier::Const) {
      global_info_.namespaces_[namespace_stack_.back()].const_namespaces[name] = id;
    } else {
      global_info_.namespaces_[namespace_stack_.back()].var_namespaces[name] = id;
    }

    namespace_info->modifier = modifier;
  } else {
    global_info_.namespaces_[namespace_stack_.back()].namespaces[name] = id;
  }

  namespace_info->parent_namespace = namespace_stack_.back();

  namespace_info->type_name = name;

  namespace_stack_.push_back(id);
  current_path_.push_back(name);
}

// better code ??
void GlobalInfo::NamespaceVisitor::EnterNamespace(const std::string& name,
                                                  std::optional<utils::IsConstModifier> modifier) {

  if (!modifier.has_value()) {
    for (ssize_t i = (ssize_t)namespace_stack_.size() - 1; i >= 0; --i) {
      auto namespace_iter = global_info_.namespaces_[namespace_stack_[i]].namespaces.find(name);
      if (namespace_iter != global_info_.namespaces_[namespace_stack_[i]].namespaces.end()) {
        namespace_stack_.push_back(namespace_iter->second);
        current_path_.push_back(name);
        return;
      }
    }
  } else if (modifier.has_value() && modifier.value() == utils::IsConstModifier::Const) {
    for (ssize_t i = (ssize_t)namespace_stack_.size() - 1; i >= 0; --i) {
    auto const_namespace_iter = global_info_.namespaces_[namespace_stack_[i]].const_namespaces.find(name);
      if (const_namespace_iter != global_info_.namespaces_[namespace_stack_[i]].const_namespaces.end()) {
        namespace_stack_.push_back(const_namespace_iter->second);
        current_path_.push_back(name);
        return;
      }
    }
  } else {
    for (ssize_t i = (ssize_t)namespace_stack_.size() - 1; i >= 0; --i) {
    auto var_namespace_iter = global_info_.namespaces_[namespace_stack_[i]].var_namespaces.find(name);
      if (var_namespace_iter != global_info_.namespaces_[namespace_stack_[i]].var_namespaces.end()) {
        namespace_stack_.push_back(var_namespace_iter->second);
        current_path_.push_back(name);
        return;
      }
    }
  }

  error_handling::HandleInternalError("Can't find namespace " + name,
                                      "GlobalInfo.NamespaceVisitor.EnterNamespace");
}

void GlobalInfo::NamespaceVisitor::ExitNamespace() {
  if (namespace_stack_.size() <= 1) {
    error_handling::HandleInternalError("Can't exit from global namespace",
                                        "GlobalInfo.NamespaceVisitor.ExitNamespace");
    return;
  }

  namespace_stack_.pop_back();
  current_path_.pop_back();
}

void GlobalInfo::NamespaceVisitor::ToGlobalNamespace() {
  namespace_stack_.clear();
  current_path_.clear();

  namespace_stack_.push_back(global_info_.GlobalNamespaceId);
}

utils::IdType GlobalInfo::NamespaceVisitor::AddFunctionDeclaration(
  const std::string& name,
  definition::FunctionDeclaration&& function_declaration_info) {

  size_t id = 0;

  auto function_id_iter = global_info_.namespaces_[namespace_stack_.back()].functions.find(name);

  if (function_id_iter == global_info_.namespaces_[namespace_stack_.back()].functions.end()) {
    id = global_info_.functions_.size();
    global_info_.namespaces_[namespace_stack_.back()].functions[name] = id;
    global_info_.functions_.emplace_back();
    global_info_.functions_.back().argument_count = function_declaration_info.argument_types.size(); // add return type
    global_info_.functions_.back().declaration = std::move(function_declaration_info);
  } else {
    id = function_id_iter->second;
    if (global_info_.functions_.back().argument_count != function_declaration_info.argument_types.size()) {
      error_handling::HandleInternalError("Not same argument count in function definition and declaration",
                                          "GlobalInfo.NamespaceVisitor. AddFunctionDeclaration");
    }
    global_info_.functions_[id].declaration = std::move(function_declaration_info);
  }

  return id;
}

utils::IdType GlobalInfo::NamespaceVisitor::AddFunctionDefinition(const std::string& name,
    definition::FunctionDefinition&& function_definition_info) {
  size_t id = 0;

  auto function_id_iter = global_info_.namespaces_[namespace_stack_.back()].functions.find(name);

  if (function_id_iter == global_info_.namespaces_[namespace_stack_.back()].functions.end()) {
    id = global_info_.functions_.size();
    global_info_.namespaces_[namespace_stack_.back()].functions[name] = id;
    global_info_.functions_.emplace_back();
    global_info_.functions_.back().argument_count = function_definition_info.argument_names.size() + 1;
    global_info_.functions_.back().definition = std::move(function_definition_info);
  } else {
    id = function_id_iter->second;
    if (global_info_.functions_.back().argument_count != function_definition_info.argument_names.size() + 1) {
      error_handling::HandleInternalError("Not same argument count in function definition and declaration", "GlobalInfo");
    }
    global_info_.functions_[id].definition = std::move(function_definition_info);
  }

  return id;
}

// TODO: internal types, etc.
// TODO: extended constructor names (point separated names)
utils::IdType GlobalInfo::NamespaceVisitor::AddType(const std::string& type,
                                                    definition::Type&& type_info,
                                                    const interpreter::tokens::BaseNode& base_node) {
  if (type::ToInternalType(type).has_value()) {
    error_handling::HandleTypecheckError("Can't redefine basic type", base_node);
  }

  size_t id = 0;

  auto type_id_iter = global_info_.namespaces_[namespace_stack_.back()].types.find(type);

  if (type_id_iter == global_info_.namespaces_[namespace_stack_.back()].types.end()) {
    id = global_info_.types_.size();
    global_info_.namespaces_[namespace_stack_.back()].types[type] = id;
    global_info_.types_.push_back(std::move(type_info));
  } else {
    error_handling::HandleTypecheckError("More then one type with the same name in namespace", base_node);
  } // TODO: typecheck error??

  definition::Type& moved_type_info = global_info_.types_.back();

  // TODO: constructors for tuple types, function types (??), array types (??), ...

  if (std::holds_alternative<definition::AnyType>(moved_type_info.type)) {
    definition::AnyType& any_type_info = std::get<definition::AnyType>(moved_type_info.type);
    if (std::holds_alternative<std::unique_ptr<interpreter::tokens::VariantType>>(*any_type_info.value)) {
      interpreter::tokens::VariantType& variant_type_info = *std::get<std::unique_ptr<interpreter::tokens::VariantType>>(*any_type_info.value);
      for (size_t i = 0; i < variant_type_info.constructors.size(); ++i) {
        std::string constructor_name;
        definition::Constructor constructor_info;

        constructor_info.type_id = id;
        constructor_info.order = i;

        if (std::holds_alternative<interpreter::tokens::Constructor>(variant_type_info.constructors[i])) {
          constructor_name = std::get<interpreter::tokens::Constructor>(variant_type_info.constructors[i]);
        } else if (std::holds_alternative<
            std::unique_ptr<interpreter::tokens::TupleType>>(variant_type_info.constructors[i])) {
          constructor_info.constructor_tuple_node =
            std::get<std::unique_ptr<interpreter::tokens::TupleType>>(variant_type_info.constructors[i]).get();

          auto maybe_constructor_name = constructor_info.constructor_tuple_node.value()->type;

          if (maybe_constructor_name.has_value()) {
            constructor_name = maybe_constructor_name.value();
          } else {
            constructor_name = type;
          }
        } else {
          // error
        }

        constructor_info.name = constructor_name;

        AddConstructor(constructor_name, std::move(constructor_info), base_node);
      }
    }
  }

  return id;
}

// TODO: link abstract type with let definitions
utils::IdType GlobalInfo::NamespaceVisitor::AddAbstractType(const std::string& abstract_type,
                                                            definition::AbstractType&& abstract_type_info,
                                                            const interpreter::tokens::BaseNode& base_node) {
  if (type::ToInternalType(abstract_type).has_value()) {
    error_handling::HandleTypecheckError("Can't redefine basic type as abstract type", base_node);
  }

  if (!FindAbstractType(abstract_type).has_value()) {
    utils::IdType id = global_info_.abstract_types_.size();
    global_info_.name_to_abstract_type_[abstract_type] = id;
    global_info_.abstract_types_.push_back(std::move(abstract_type_info));

    return id;
  }

  error_handling::HandleTypecheckError("More then one abstract type with the same name in namespace",
                                       base_node);
  return 0;
}

utils::IdType GlobalInfo::NamespaceVisitor::AddTypeclass(const std::string& typeclass,
                                                         definition::Typeclass&& typeclass_info,
                                                         const interpreter::tokens::BaseNode& base_node) {
  if (type::ToInternalType(typeclass).has_value()) {
    error_handling::HandleTypecheckError("Can't redefine basic type as typeclass", base_node);
  }

  if (!FindTypeclass(typeclass).has_value()) {
    size_t id = global_info_.typeclasses_.size();
    global_info_.name_to_typeclass_[typeclass] = id;
    global_info_.typeclasses_.push_back(std::move(typeclass_info));

    return id;
  }


  error_handling::HandleTypecheckError("More then one typeclass with the same name in namespace", base_node);
  return 0;
}

utils::IdType GlobalInfo::NamespaceVisitor::AddConstructor(const std::string& constructor,
                                                           definition::Constructor&& constructor_info,
                                                           const interpreter::tokens::BaseNode& base_node) {
  if (type::ToInternalType(constructor).has_value()) {
    error_handling::HandleTypecheckError("Can't redefine basic type as constructor", base_node);
  }

  auto constructor_id_iter = global_info_.namespaces_[namespace_stack_.back()].constructors.find(constructor);

  if (constructor_id_iter == global_info_.namespaces_[namespace_stack_.back()].constructors.end()) {
    size_t id = global_info_.constructors_.size();
    global_info_.namespaces_[namespace_stack_.back()].constructors[constructor] = id;
    global_info_.constructors_.push_back(std::move(constructor_info));

    return id;
  }

  error_handling::HandleTypecheckError("More then one constructor with the same name in namespace", base_node);
  exit(1);
}

utils::IdType GlobalInfo::NamespaceVisitor::AddPartition(const std::vector<std::string>& path,
                                                         interpreter::tokens::PartitionStatement* node) {
  PartitionInfo partition;

  partition.path.reserve(current_path_.size() + path.size() - 1);
  partition.path = current_path_;

  for (size_t i = 0; i + 1 < path.size(); ++i) {
    partition.path.push_back(path[i]);
  }

  partition.name = path.back();
  partition.node = node;

  utils::IdType id = global_info_.partitions_.size();

  global_info_.partitions_.push_back(partition);
  global_info_.partitions_trie_.Insert(partition.path, id);

  return id;
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindNamespace(const std::optional<std::vector<std::string>>& path) {
  return FindSomething<utils::IdType>(path,
      [] (utils::IdType current_namespace) -> std::optional<utils::IdType> {
    return current_namespace;
  });
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindFunction(
    const std::optional<std::vector<std::string>>& path,
    const std::string& name) {
  return FindSomething<utils::IdType>(path,
      [name, this] (utils::IdType current_namespace) -> std::optional<utils::IdType> {

    auto function_info_iter = global_info_.namespaces_[current_namespace].functions.find(name);
    if (function_info_iter == global_info_.namespaces_[current_namespace].functions.end()) {
      return std::nullopt;
    }

    return function_info_iter->second;
  });
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindMethod(
    const std::optional<std::vector<std::string>>& path,
    const std::string& type,
    const std::string& name,
    utils::IsConstModifier modifier) {
  // TODO: remove overhead

  return GlobalInfo::NamespaceVisitor::FindSomething<utils::IdType>(path,
      [type, name, modifier, this] (utils::IdType current_namespace) -> std::optional<utils::IdType> {


    auto variable_namespace_iter =
      (modifier == utils::IsConstModifier::Const
        ? global_info_.namespaces_[current_namespace].const_namespaces.find(type)
        : global_info_.namespaces_[current_namespace].var_namespaces.find(type));
    if (variable_namespace_iter ==
      (modifier == utils::IsConstModifier::Const
        ? global_info_.namespaces_[current_namespace].const_namespaces.end()
        : global_info_.namespaces_[current_namespace].var_namespaces.end())) {
      return std::nullopt;
    }

    auto method_iter = global_info_.namespaces_[variable_namespace_iter->second].functions.find(name);
    if (method_iter == global_info_.namespaces_[variable_namespace_iter->second].functions.end()) {
        return std::nullopt;
    }

    return method_iter->second;
  });
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindType(
    const std::optional<std::vector<std::string>>& path,
    const std::string& type) {
  return FindSomething<utils::IdType>(path,
      [type, this] (utils::IdType current_namespace) -> std::optional<utils::IdType> {

    auto type_info_iter = global_info_.namespaces_[current_namespace].types.find(type);
    if (type_info_iter == global_info_.namespaces_[current_namespace].types.end()) {
      return std::nullopt;
    }

    return type_info_iter->second;
  });
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindLocalType(const std::string& type) {
  auto type_id_iter = global_info_.namespaces_[namespace_stack_.back()].types.find(type);

  if (type_id_iter != global_info_.namespaces_[namespace_stack_.back()].types.end()) {
    return type_id_iter->second;
  }

  return std::nullopt;
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindAbstractType(const std::string& abstract_type) {
  auto abstract_type_id_iter = global_info_.name_to_abstract_type_.find(abstract_type);

  if (abstract_type_id_iter != global_info_.name_to_abstract_type_.end()) {
    return abstract_type_id_iter->second;
  }

  return std::nullopt;
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindTypeclass(const std::string& typeclass) {
  auto typeclass_id_iter = global_info_.name_to_typeclass_.find(typeclass);

  if (typeclass_id_iter != global_info_.name_to_typeclass_.end()) {
    return typeclass_id_iter->second;
  }

  return std::nullopt;
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindConstructor(
    const std::optional<std::vector<std::string>>& path,
    const std::string& constructor) {
  return FindSomething<utils::IdType>(path,
      [constructor, this] (utils::IdType current_namespace) -> std::optional<utils::IdType> {

    auto constructor_info_iter = global_info_.namespaces_[current_namespace].constructors.find(constructor);
    if (constructor_info_iter == global_info_.namespaces_[current_namespace].constructors.end()) {
      return std::nullopt;
    }

    return constructor_info_iter->second;
  });
}

template<typename T>
std::optional<T> GlobalInfo::NamespaceVisitor::FindSomething(
    const std::optional<std::vector<std::string>>& path,
    std::function<std::optional<T>(utils::IdType)> search_func) {
  for (ssize_t i = namespace_stack_.size() - 1; i >= 0; --i) {
    utils::IdType current_namespace;
    if (path.has_value()) {
      auto maybe_namespace = FindNamespaceIn(namespace_stack_[i], path.value());

      if (!maybe_namespace.has_value()) {
        continue;
      }

      current_namespace = maybe_namespace.value();
    } else {
      current_namespace = namespace_stack_[i];
    }

    std::optional<T> result = search_func(current_namespace);

    if (result.has_value()) {
      return result.value();
    }
  }

  return std::nullopt;
}

std::optional<utils::IdType> GlobalInfo::NamespaceVisitor::FindNamespaceIn(
    utils::IdType current_namespace,
    const std::vector<std::string>& path) {
  utils::IdType next_namespace = current_namespace;
  for (auto& name : path) {
    auto next_namespace_iter = global_info_.namespaces_[next_namespace].namespaces.find(name);
    if (next_namespace_iter == global_info_.namespaces_[next_namespace].namespaces.end()) {
      return std::nullopt;
    }
    next_namespace = next_namespace_iter->second;
  }

  return next_namespace;
}

} // namespace info