#include "name_tree.hpp"

#include "statement_nodes.hpp"
#include "utils.hpp"

namespace names {

// --- public

bool NameTree::insert(const std::string &path, nodes::Statement &&statement) {
  auto name_ids = slice_path_to_name_ids(path);

  size_t node_id = add_node(get_root(), name_ids);

  if (nodes_[node_id].get_statement().has_value()) {
    return false;
  }

  nodes_[node_id].set_statement(std::move(statement));
  return true;
}

nodes::CombineResult NameTree::insert_combine(const std::string &path,
                                              nodes::Statement &&statement) {
  auto name_ids = slice_path_to_name_ids(path);

  size_t node_id = add_node(get_root(), name_ids);

  if (!nodes_[node_id].get_statement().has_value()) {
    nodes_[node_id].set_statement(std::move(statement));
    return nodes::CombineResult::OK;
  }

  return nodes_[node_id].get_statement().value()->combine(std::move(statement));
}

std::optional<nodes::Statement *> NameTree::find(const std::string &path) {
  auto name_ids = slice_path_to_existing_name_ids(path);

  if (!name_ids.has_value()) {
    return std::nullopt;
  }

  auto node_id = find_node(get_root(), name_ids.value());

  if (node_id.has_value()) {
    return nodes_[node_id.value()].get_statement();
  }

  return std::nullopt;
}

std::optional<const nodes::Statement *>
NameTree::find(const std::string &path) const {
  auto name_ids = slice_path_to_existing_name_ids(path);

  if (!name_ids.has_value()) {
    return std::nullopt;
  }

  auto node_id = find_node(get_root(), name_ids.value());

  if (node_id.has_value()) {
    return nodes_[node_id.value()].get_statement();
  }

  return std::nullopt;
}

void NameTree::print(printers::Printer &printer) const {
  nodes_[get_root()].print(nodes_, printer);
}

// TODO
void NameTree::add_statement_children_to_tree() {
  for (auto &node : nodes_) {
    if (!node.get_statement().has_value()) {
      continue;
    }

    switch (node.get_statement().value()->get_any()->index()) {
    case 0: // Import
      // TODO: link imported symbols, if named import, otherwise link imported
      // symbols to root
      break;
    case 1: // FunctionDefinition
      // TODO: link to result type
      break;
    case 2: // TypeDefinition
      // TODO: link methods + link children to types + connect to typeclasses
      break;
    case 3: // EmptyLines
      break;
    default:
      error_handling::handle_general_error(
          "Unexpected statement type in name tree");
      break;
    }
  }
}

// --- private

size_t NameTree::add_node(size_t current_node,
                          const std::vector<size_t> &name_ids,
                          size_t current_name) {
  if (current_name >= name_ids.size()) {
    return current_node;
  }

  auto next_node = nodes_[current_node].find(name_ids[current_name]);

  if (!next_node.has_value()) {
    next_node = nodes_.size();
    nodes_[current_node].insert(name_ids[current_name], next_node.value());
    nodes_.emplace_back();
  }

  return add_node(next_node.value(), name_ids, current_name + 1);
}

std::optional<size_t> NameTree::find_node(size_t current_node,
                                          const std::vector<size_t> &name_ids,
                                          size_t current_name) const {
  if (current_name >= name_ids.size()) {
    return current_node;
  }

  auto next_node = nodes_[current_node].find(name_ids[current_name]);

  if (!next_node.has_value()) {
    return std::nullopt;
  }

  return find_node(next_node.value(), name_ids, current_name + 1);
}

std::vector<size_t> NameTree::slice_path_to_name_ids(const std::string &path) {
  // check that path is / is not points only operator
  bool points_only = true;
  for (size_t i = 0; i < path.size(); ++i) {
    if (path[i] != '.') {
      points_only = true;
      break;
    }
  }

  if (points_only) {
    return {add_name_id(path)};
  }

  std::vector<size_t> name_ids;

  size_t last_name_start = 0;
  for (size_t i = 0; i < path.size() + 1; ++i) {
    if (path[i] == '.' || i == path.size()) { // name separator
      name_ids.push_back(
          add_name_id(path.substr(last_name_start, i - last_name_start)));
      last_name_start = i + 1;
    }
  }

  return name_ids;
}

std::optional<std::vector<size_t>>
NameTree::slice_path_to_existing_name_ids(const std::string &path) const {
  // check that path is / is not points only operator
  bool points_only = true;
  for (size_t i = 0; i < path.size(); ++i) {
    if (path[i] != '.') {
      points_only = true;
      break;
    }
  }

  if (points_only) {
    auto maybe_id = find_name_id(path);

    if (maybe_id.has_value()) {
      return std::vector<size_t>{maybe_id.value()};
    } else {
      return std::nullopt;
    }
  }

  std::vector<size_t> name_ids;

  size_t last_name_start = 0;
  for (size_t i = 0; i < path.size() + 1; ++i) {
    if (path[i] == '.' || i == path.size()) { // name separator
      auto maybe_id =
          find_name_id(path.substr(last_name_start, i - last_name_start));

      if (!maybe_id.has_value()) {
        return std::nullopt;
      }

      name_ids.push_back(maybe_id.value());
      last_name_start = i + 1;
    }
  }

  return name_ids;
}

size_t NameTree::add_name_id(const std::string &name) {
  auto name_iter = name_to_id_.find(name);

  if (name_iter == name_to_id_.end()) {
    size_t name_id = name_to_id_.size();
    return name_to_id_[name] = name_id;
  }

  return name_iter->second;
}

std::optional<size_t> NameTree::find_name_id(const std::string &name) const {
  auto name_iter = name_to_id_.find(name);

  if (name_iter == name_to_id_.end()) {
    return std::nullopt;
  }

  return name_iter->second;
}

} // namespace names