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

#pragma once

#include "trie.hpp"

#include "iterator/end_aware_iterator.hpp"
#include <stack>

namespace detail {
  template <typename Iter>
  struct trie_iterator_next {
    template <typename Trie>
    ::iterator::end_aware_iterator<Iter> operator()(Trie & tr) {
      return { tr.local_begin(), tr.local_end() };
    }
  };

  template <typename Iter>
  struct trie_iterator_next<std::reverse_iterator<Iter>> {
    template <typename Trie>
    ::iterator::end_aware_iterator<std::reverse_iterator<Iter>> operator()(Trie & tr) {
      return { tr.local_rbegin(), tr.local_rend() };
    }
  };
  
  template <typename Trie, typename Iter>
  class trie_iterator_base {
  protected:
    using impl_t = ::iterator::end_aware_iterator<Iter>;
    std::stack<Trie *> stk;
    std::list<typename Trie::key_type> keys;
    std::stack<impl_t> iters;
    bool done{false};
  public:
    trie_iterator_base() : done{true} {}
    trie_iterator_base(Trie * tr) { stk.push(tr); }
    auto operator*() const -> decltype(this->stk.top()->value()) { return *stk.top(); }
  protected:
    bool poping_empty() {
      if (iters.top().done()) {
        stk.pop();
        iters.pop();
        keys.pop_back();
        return !iters.empty();
      }
      return false;
    }
    void assign() {
      if (!iters.empty()) {
        stk.top() = iters.top()->second.get();
        keys.back() = iters.top()->first;
      }
    }
    bool can_recurse() { return !next().done(); }
    void recurse() {
      auto it = next();
      iters.push(it);
      keys.push_back(it->first);
      stk.push(it->second.get());
    }
    impl_t next() { return detail::trie_iterator_next<Iter>()(*stk.top()); }
    friend bool operator!=(trie_iterator_base const & lhs, trie_iterator_base const & rhs) {
      return !(lhs == rhs);
    }
    friend bool operator==(trie_iterator_base const & lhs, trie_iterator_base const & rhs) {
      if ( lhs.done && rhs.done ) return true;
      else if (lhs.keys != rhs.keys) return false;
      else if ( lhs.done != rhs.done ) return false;
      else return *lhs == *rhs;
    }
  };
}

template <typename Trie, typename Iter>
class trie_iterator : public detail::trie_iterator_base<Trie, Iter> {
private:
  using super = detail::trie_iterator_base<Trie, Iter>;
public:
  using reference = decltype(*(std::declval<super>()));
  using value_type = std::remove_reference_t<reference>;
  using pointer = value_type *;
  using difference_type = std::ptrdiff_t;
  using iterator_category = std::forward_iterator_tag;
public:
  trie_iterator() : super() {}
  trie_iterator(Trie * tr) : super(tr) { }
  trie_iterator & operator++() {
    if (super::done) return *this;
    if ( super::iters.empty() || super::can_recurse() ) { super::recurse(); }
    else { advance(); }
    return *this;
  }
private:
  void advance() {
    ++super::iters.top();
    while (super::poping_empty()) {
      ++super::iters.top();
    }
    super::assign();
    super::done = super::iters.empty();
  }
  friend Trie;
};

template <typename Trie, typename Iter>
class trie_reverse_iterator : public detail::trie_iterator_base<Trie, Iter> {
private:
  using super = detail::trie_iterator_base<Trie, Iter>;
public:
  using reference = decltype(*(std::declval<super>()));
  using value_type = std::remove_reference_t<reference>;
  using pointer = value_type *;
  using difference_type = std::ptrdiff_t;
  using iterator_category = std::forward_iterator_tag;
public:
  trie_reverse_iterator() : super() {}
  trie_reverse_iterator(Trie * tr) : super(tr) {
    if (super::can_recurse()) recurse();
  }
  trie_reverse_iterator & operator++() {
    if (super::done || super::iters.empty()) { super::done = true; }
    else { advance(); }
    return *this;
  }
private:
  void recurse() {
    do { super::recurse(); } while ( super::can_recurse() );
  }
  void advance() {
    ++super::iters.top();
    if (super::iters.top().done()) {
      while (super::poping_empty());
      super::assign();
    } else {
      super::assign();
      if (super::can_recurse()) recurse();
    }
  }
  friend Trie;
};