//
//  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 V, typename Iter>
class trie_iterator {
private:
  using impl_t = ::iterator::end_aware_iterator<Iter>;
public:
  using value_type = V;
  using reference = value_type &;
  using pointer = value_type *;
  using difference_type = std::ptrdiff_t;
  using iterator_category = std::forward_iterator_tag;
public:
  trie_iterator() : done{true} {}
  trie_iterator(Trie * tr) { stk.push(tr); }
  reference operator*() const { return *stk.top(); }
  trie_iterator & operator++() {
    if (done) return *this;
    if ( iters.empty() || can_recurse() ) { recurse(); }
    else { advance(); }
    return *this;
  }
private:
  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()); }
  void advance() {
    ++iters.top();
    while (iters.top().done()) {
      stk.pop();
      iters.pop();
      keys.pop_back();
      if ( iters.empty() ) break;
      ++iters.top();
    }
    if (!iters.empty()) {
      stk.top() = iters.top()->second.get();
      keys.back() = iters.top()->first;
    } else {
      done = true;
    }
  }
  friend bool operator!=(trie_iterator const & lhs, trie_iterator const & rhs) {
    return !(lhs == rhs);
  }
  friend bool operator==(trie_iterator const & lhs, trie_iterator 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;
  }
  friend Trie;
  std::stack<Trie *> stk;
  std::list<typename Trie::key_type> keys;
  std::stack<impl_t> iters;
  bool done{false};
};

template <typename Trie, typename V, typename Iter>
class trie_reverse_iterator {
private:
  using impl_t = ::iterator::end_aware_iterator<Iter>;
public:
  using value_type = V;
  using reference = value_type &;
  using pointer = value_type *;
  using difference_type = std::ptrdiff_t;
  using iterator_category = std::forward_iterator_tag;
public:
  trie_reverse_iterator() : done{true} {}
  trie_reverse_iterator(Trie * tr) {
    stk.push(tr);
    if (can_recurse()) recurse();
  }
  reference operator*() const { return *stk.top(); }
  trie_reverse_iterator & operator++() {
    if (done || iters.empty()) { done = true; return *this; }
    advance();
    return *this;
  }
private:
  bool can_recurse() { return !next().done(); }
  void recurse() {
    do {
      auto it = next();
      iters.push(it);
      keys.push_back(it->first);
      stk.push(it->second.get());
    } while ( can_recurse() );
  }
  impl_t next() { return detail::trie_iterator_next<Iter>()(*stk.top()); }
  void advance() {
    ++iters.top();
    if (iters.top().done()) {
      while (iters.top().done()) {
        stk.pop();
        iters.pop();
        keys.pop_back();
        if ( iters.empty() ) break;
      }
      if (!iters.empty()) {
        stk.top() = iters.top()->second.get();
        keys.back() = iters.top()->first;
      }
    } else {
      stk.top() = iters.top()->second.get();
      keys.back() = iters.top()->first;
      if (can_recurse()) recurse();
    }
  }
  friend bool operator!=(trie_reverse_iterator const & lhs, trie_reverse_iterator const & rhs) {
    return !(lhs == rhs);
  }
  friend bool operator==(trie_reverse_iterator const & lhs, trie_reverse_iterator 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;
  }
  friend Trie;
  std::stack<Trie *> stk;
  std::list<typename Trie::key_type> keys;
  std::stack<impl_t> iters;
  bool done{false};
};