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

#pragma once

#include "trie.hpp"

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

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

  template <typename Iter>
  struct trie_iterator_next<std::reverse_iterator<Iter>> {
    using iter_t = ::iterator::end_aware_iterator<std::reverse_iterator<Iter>>;
    template <typename Trie> iter_t operator()(Trie & tr) {
      return {tr.local_rbegin(), tr.local_rend()};
    }
  };

  template <typename Trie, typename Iter> class trie_iterator_base {
  protected:
    using helper = trie_iterator_next<Iter>;
    using impl_t = typename helper::iter_t;
    std::stack<Trie *> stk;
    std::list<typename Trie::key_type> keys;
    std::stack<impl_t> iters;
    bool done{false};

  public:
    using reference = decltype(std::declval<Trie>().value());
    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_base() : done{true} {}
    trie_iterator_base(Trie * tr) { stk.push(tr); }
    auto operator*() const -> reference { return root(); }
    auto operator-> () const -> pointer { return std::addressof(operator*()); }
    Trie & root() const {
      if (stk.empty()) {
        throw std::runtime_error("Dereferencing invalid iterator");
      }
      return *stk.top();
    }
    trie_iterator_base parent() const {
      trie_iterator_base tmp{*this};
      tmp.pop();
      return tmp;
    }

  protected:
    void push(impl_t it) {
      if (it.done()) {
        done = true;
        return;
      }
      iters.push(it);
      keys.push_back(it->first);
      stk.push(it->second.get());
    }
    void pop() {
      stk.pop();
      iters.pop();
      keys.pop_back();
    }
    bool poping_empty() {
      if (!iters.top().done()) { return false; }
      pop();
      return !iters.empty();
    }
    void assign() {
      if (iters.empty()) { return; }
      stk.top() = iters.top()->second.get();
      keys.back() = iters.top()->first;
    }
    bool can_recurse() { return !next().done(); }
    void recurse() { push(next()); }
    impl_t next() { return helper()(root()); }
    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) {
      return (lhs.done && rhs.done) ||
             (lhs.keys == rhs.keys && lhs.done == rhs.done && *lhs == *rhs);
    }
    friend Trie;
    template <typename Self, typename _Trie, typename KS>
    friend Self find_impl(_Trie * tr, KS const & keys);
  };
}

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:
  trie_iterator() : super() {}
  trie_iterator(Trie * tr) : super(tr) {}
  trie_iterator(super && take) : super(std::move(take)) {}
  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:
  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;
};