#include "name_tree.hpp"

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

namespace names {

// --- AnyStatementProxy

nodes::Statement *AnyStatementProxy::get() {
  return name_tree_->nodes_.at(id_).get_statement().value();
}

const nodes::Statement *AnyStatementProxy::get() const {
  return name_tree_->nodes_.at(id_).get_statement().value();
}

// --- NameTree

// --- public

std::optional<AnyStatementProxy>
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 std::nullopt;
  }

  nodes_[node_id].set_statement(std::move(statement));
  return AnyStatementProxy(*this, node_id);
}

std::pair<std::optional<AnyStatementProxy>, 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 {AnyStatementProxy(*this, node_id), nodes::CombineResult::OK};
  }

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

std::optional<AnyStatementProxy> 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 AnyStatementProxy(*this, node_id.value());
  }

  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