//
//  trie.hpp
//  trie
//
//  Created by Sam Jaffe on 6/16/17.
//

#pragma once

#include "pointers/const_propogating_ptr.hpp"
#include "pointers/not_null.hpp"
#include "pointers/value_ptr.hpp"
#include <map>

template <typename K, typename V, typename Compare = std::less<K>> class trie;

namespace detail {
  template <typename Trie, typename Iter> class trie_iterator_base;
  template <typename Iter, typename Trie, typename KS>
  Iter find_impl(Trie * tr, KS const & keys);
}

template <typename Trie, typename Iter> class trie_iterator;
template <typename Trie, typename Iter> class trie_reverse_iterator;

template <typename K, typename V, typename Compare> class trie {
private:
  using self_t = trie<K, V, Compare>;
  using impl_value_type = const_propogating_ptr<value_ptr<self_t>>;
  using backing_t = std::map<K, impl_value_type, Compare>;
  template <typename KS>
  using element_type = std::decay_t<decltype(*std::begin(std::declval<KS>()))>;
  template <typename KS>
  using is_collection_t =
      std::enable_if_t<std::is_same<K, element_type<KS>>::value>;

public:
  using key_type = K;
  using mapped_type = V;
  using key_compare = Compare;

  using local_iterator = typename backing_t::iterator;
  using local_const_iterator = typename backing_t::const_iterator;
  using local_reverse_iterator = typename backing_t::reverse_iterator;
  using local_const_reverse_iterator =
      typename backing_t::const_reverse_iterator;

  using iterator = trie_iterator<self_t, local_iterator>;
  using const_iterator = trie_iterator<self_t const, local_const_iterator>;
  using post_iterator = trie_iterator<self_t, local_reverse_iterator>;
  using const_post_iterator =
      trie_iterator<self_t const, local_const_reverse_iterator>;
  using reverse_iterator =
      trie_reverse_iterator<self_t, local_reverse_iterator>;
  using const_reverse_iterator =
      trie_reverse_iterator<self_t const, local_const_reverse_iterator>;

private:
  using impl_iterator = detail::trie_iterator_base<self_t, local_iterator>;
  using impl_const_iterator =
      detail::trie_iterator_base<self_t const, local_const_iterator>;

public:
  trie() {}
  trie(trie const & other);
  trie(trie && other) : trie() { swap(*this, other); }
  ~trie() { clear(); }
  self_t & operator=(mapped_type const & value);
  self_t & operator=(trie const & value);
  self_t & operator=(trie && value);

  operator mapped_type &() { return value(); }
  operator mapped_type const &() const { return value(); }
  mapped_type & value() { return value_; }
  mapped_type const & value() const { return value_; }

  template <typename KS, typename = is_collection_t<KS>>
  self_t & operator[](KS const & keys) {
    return *emplace(keys).first.stk.top();
  }
  self_t & operator[](key_type const & key) {
    return *emplace(key).first.stk.top();
  }
  self_t & operator[](std::initializer_list<key_type> keys) {
    return operator[]<>(keys);
  }

  template <typename KS>
  std::pair<iterator, bool> insert(KS const & keys, mapped_type const & value) {
    return emplace_impl(keys, value);
  }
  std::pair<iterator, bool> insert(key_type const & key,
                                   mapped_type const & value) {
    return emplace_impl({key}, value);
  }
  std::pair<iterator, bool> insert(std::initializer_list<key_type> keys,
                                   mapped_type const & value) {
    return emplace_impl(keys, value);
  }
  template <typename KS, typename... Args>
  std::pair<iterator, bool> emplace(KS && keys, Args &&... args) {
    return emplace_impl(keys, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(key_type const & key, Args &&... args) {
    return emplace_impl({key}, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(key_type && key, Args &&... args) {
    return emplace_impl({key}, std::forward<Args>(args)...);
  }
  template <typename... Args>
  std::pair<iterator, bool> emplace(std::initializer_list<key_type> keys,
                                    Args &&... args) {
    return emplace_impl(keys, std::forward<Args>(args)...);
  }

  iterator begin() { return {this}; }
  iterator end() { return {}; }
  const_iterator begin() const { return {this}; }
  const_iterator end() const { return {}; }
  const_iterator cbegin() const { return begin(); }
  const_iterator cend() const { return end(); }
  reverse_iterator rbegin() { return {this}; }
  reverse_iterator rend() { return {}; }
  const_reverse_iterator rbegin() const { return {this}; }
  const_reverse_iterator rend() const { return {}; }
  const_reverse_iterator crbegin() const { return rbegin(); }
  const_reverse_iterator crend() const { return rend(); }

  local_iterator local_begin() { return impl_.begin(); }
  local_iterator local_end() { return impl_.end(); }
  local_const_iterator local_begin() const { return impl_.begin(); }
  local_const_iterator local_end() const { return impl_.end(); }

  local_reverse_iterator local_rbegin() { return impl_.rbegin(); }
  local_reverse_iterator local_rend() { return impl_.rend(); }
  local_const_reverse_iterator local_rbegin() const { return impl_.rbegin(); }
  local_const_reverse_iterator local_rend() const { return impl_.rend(); }

  template <typename KS> iterator find(KS const & keys) {
    return detail::find_impl<impl_iterator>(this, keys);
  }
  iterator find(key_type const & key) {
    return find_impl<impl_iterator>(this, {key});
  }
  iterator find(std::initializer_list<key_type> keys) {
    return find_impl<impl_iterator>(this, keys);
  }
  template <typename KS> const_iterator find(KS const & keys) const {
    return detail::find_impl<impl_const_iterator>(this, keys);
  }
  const_iterator find(key_type const & key) const {
    return find_impl<impl_const_iterator>(this, {key});
  }
  const_iterator find(std::initializer_list<key_type> keys) const {
    return find_impl<impl_const_iterator>(this, keys);
  }

  template <typename KS> void erase(KS const & keys) { drop(find(keys)); }
  void erase(key_type const & key) { drop(find(key)); }
  void erase(std::initializer_list<key_type> keys) { drop(find(keys)); }

  void clear();

private:
  void drop(iterator it);

  template <typename Iter, typename Trie, typename KS>
  friend Iter detail::find_impl(Trie * tr, KS const & keys);
  template <typename Iter, typename Trie>
  static Iter find_impl(Trie * tr, std::initializer_list<key_type> keys) {
    return detail::find_impl<Iter>(tr, keys);
  }

  template <typename KS, typename... Args>
  std::pair<impl_iterator, bool> emplace_impl(KS && keys, Args &&... args);
  template <typename... Args>
  std::pair<impl_iterator, bool>
  emplace_impl(std::initializer_list<key_type> keys, Args &&... args) {
    return emplace_impl<std::initializer_list<key_type>>(
        std::move(keys), std::forward<Args>(args)...);
  }

  template <typename Key>
  void insert_impl(impl_iterator & out, bool & create, Key && key);

  friend bool operator==(trie const & lhs, trie const & rhs) {
    const_iterator it1 = lhs.begin(), it2 = rhs.begin();
    const_iterator const end1 = lhs.end(), end2 = lhs.end();
    for (; it1 != end1 && it2 != end2; ++it1, ++it2) {
      if (!it1.keys.empty() && !it2.keys.empty() &&
          it1.keys.back() != it2.keys.back()) {
        return false;
      } else if (it1.root().value_ != it2.root().value_) {
        return false;
      }
    }
    return it1 == end1 && it2 == end2;
  }

  friend void swap(trie & lhs, trie & rhs) {
    using std::swap;
    swap(lhs.value_, rhs.value_);
    swap(lhs.impl_, rhs.impl_);
  }

  mapped_type value_{};
  backing_t impl_{};
};

#include "trie_impl.hpp"