#ifndef PREFIX_SEARCH_H_A6672A1D
#define PREFIX_SEARCH_H_A6672A1D

#include <stdexcept>
#include <unordered_map>
#include <vector>

template <typename T>
class prefix_search {
 public:
  prefix_search(size_t depth = 0) : depth_(depth), count_(0) {}

  size_t getDepth() const { return depth_; }

  size_t getCount() const { return count_; }

  void add(std::string& key, T val, size_t i = 0) {
    count_++;

    if (i == key.length()) {
      elements_.push_back(val);
    } else {
      char next = key.at(i);

      if (!children_.count(next)) {
        children_.emplace(next, prefix_search<T>{depth_ + 1});
      }

      children_.at(next).add(key, val, i + 1);
    }
  }

  T at(size_t i) const {
    if (i < elements_.size()) {
      return elements_.at(i);
    } else {
      i -= elements_.size();

      for (const auto& mapping : children_) {
        if (mapping.second.count_ > i) {
          return mapping.second.at(i);
        } else {
          i -= mapping.second.count_;
        }
      }
    }

    throw std::out_of_range("Index out of range");
  }

  const prefix_search<T>& find(const std::string& val, size_t i) const {
    if (i == val.length() || !children_.count(val.at(i))) {
      return *this;
    }

    return children_.at(val.at(i)).find(val, i + 1);
  }

 private:
  size_t depth_;
  size_t count_;
  std::unordered_map<char, prefix_search<T>> children_;
  std::vector<T> elements_;
};

#endif /* end of include guard: PREFIX_SEARCH_H_A6672A1D */