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

#pragma once

#include "trie.hpp"

#include <stack>
#include <vector>

#include "iterator/end_aware_iterator.hpp"

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 {
  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;

  private:
    std::stack<Trie *> parent_trie_;
    std::vector<typename Trie::key_type> keys_;
    std::stack<iterator::end_aware_iterator<Iter>> iterators_;
    bool done_{false};

  public:
    trie_iterator_base() : done_{true} {}
    trie_iterator_base(Trie * tr) { parent_trie_.push(tr); }
    auto operator*() const -> reference { return root(); }
    auto operator-> () const -> pointer { return std::addressof(operator*()); }

    Trie & root() const {
      if (parent_trie_.empty()) {
        throw std::runtime_error("Dereferencing invalid iterator");
      }
      return *parent_trie_.top();
    }

    trie_iterator_base parent() const {
      trie_iterator_base tmp{*this};
      tmp.pop();
      return tmp;
    }

  protected:
    iterator::end_aware_iterator<Iter> & current() { return iterators_.top(); }
    bool empty() const { return iterators_.empty(); }
    void done(bool new_done) { done_ = new_done; }
    bool done() const { return done_; }

    void push(iterator::end_aware_iterator<Iter> it) {
      if (it.done()) {
        done_ = true;
        return;
      }
      iterators_.push(it);
      keys_.push_back(it->first);
      parent_trie_.push(it->second.get());
    }

    void pop() {
      parent_trie_.pop();
      iterators_.pop();
      keys_.pop_back();
    }

    bool poping_empty() {
      if (!current().done()) { return false; }
      pop();
      return !empty();
    }

    void assign() {
      if (empty()) { return; }
      parent_trie_.top() = current()->second.get();
      keys_.back() = current()->first;
    }

    bool can_recurse() { return !next().done(); }

    void recurse() { push(next()); }

    iterator::end_aware_iterator<Iter> next() {
      return trie_iterator_next<Iter>()(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::empty() || super::can_recurse()) {
      super::recurse();
    } else {
      advance();
    }
    return *this;
  }

private:
  void advance() {
    ++super::current();
    while (super::poping_empty()) {
      ++super::current();
    }
    super::assign();
    super::done(super::empty());
  }
};

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::empty()) {
      super::done(true);
    } else {
      advance();
    }
    return *this;
  }

private:
  void recurse() {
    do {
      super::recurse();
    } while (super::can_recurse());
  }

  void advance() {
    ++super::current();
    if (super::current().done()) {
      while (super::poping_empty())
        ;
      super::assign();
    } else {
      super::assign();
      if (super::can_recurse()) recurse();
    }
  }
};