class TrieElement
{
 public:
  using TrieElementPtr = std::unique_ptr<TrieElement>;
  using TrieElementRawPtr = TrieElement*;
  std::unordered_map<char, TrieElementPtr> children_;
  bool is_word{false};
  uint32_t words_underneath{0};
};

class Trie
{
 public:
  void AddWord(const std::string& word_str)
  {
    if (word_str.empty()) return;

    auto current_node = &root_;
    for (const char letter : word_str)
    {
      if (not current_node->children_.count(letter))
        current_node->children_[letter] = std::make_unique<TrieElement>();

      current_node = current_node->children_.at(letter).get();
      ++current_node->words_underneath;
    }
    current_node->is_word = true;
  }

  uint32_t FindPartialMatches(const std::string& query_str)
  {
    uint32_t matches{0};

    auto node = FindLastNodeForQuery(query_str);
    if (not node) return matches;

    // from here we start looking for partials
    std::queue<TrieElement::TrieElementRawPtr> search_queue;
    search_queue.push(node);
    while (not search_queue.empty())
    {
      const auto current = search_queue.front();
      search_queue.pop();
      if (current->is_word)
      {
        ++matches;
      }

      // get children and add them to queue
      for (const auto& pair : current->children_)
      {
        search_queue.push(pair.second.get());
      }
    }

    return matches;
  }

  bool FindExact(const std::string& query_str)
  {
    auto node = FindLastNodeForQuery(query_str);
    if (not node) return false;

    return node->is_word;
  }

  uint32_t FindPartialMatchesFast(const std::string& query_str)
  {
    auto node = FindLastNodeForQuery(query_str);
    if (not node) return 0;

    return node->words_underneath;
  }

 private:
  TrieElement::TrieElementRawPtr FindLastNodeForQuery(const std::string& query_str)
  {
    if (query_str.empty()) return nullptr;

    auto current_node = &root_;
    for (const char letter : query_str)
    {
      if (not current_node->children_.count(letter)) return nullptr;
      current_node = current_node->children_.at(letter).get();
    }

    return current_node;
  }

  TrieElement root_;
};