//
//  join_iterator.hpp
//  iterator
//
//  Created by Sam Jaffe on 2/7/17.
//

#pragma once

#include "end_aware_iterator.hpp"

#include <iterator>
#include <utility>

namespace iterator {
  template <typename MetaIterator> class joining_iterator {
  public:
    using join_iter = MetaIterator;
    using joinable_type = typename join_iter::value_type;
    using iter_type = decltype(std::begin(*std::declval<join_iter>()));

    using value_type = typename std::iterator_traits<iter_type>::value_type;
    using reference = typename std::iterator_traits<iter_type>::reference;
    using pointer = typename std::iterator_traits<iter_type>::pointer;
    using difference_type =
        typename std::iterator_traits<iter_type>::difference_type;
    using iterator_category = std::forward_iterator_tag;

  public:
    explicit joining_iterator() = default;

    template <typename I>
    joining_iterator(joining_iterator<I> const & other)
        : joiner_(other.join_iterator()), iterator_(other.element_iterator()) {}

    joining_iterator(end_aware_iterator<join_iter> join)
        : joining_iterator(join, {}, true) {}
    joining_iterator(end_aware_iterator<join_iter> join,
                     end_aware_iterator<iter_type> elem, bool update = false)
        : joiner_(join), iterator_(elem) {
      if (update) { update_iterator(); }
    }

    joining_iterator & operator++() {
      if ((++iterator_).done()) {
        ++joiner_;
        update_iterator();
      }
      return *this;
    }

    joining_iterator operator++(int) {
      joining_iterator tmp{*this};
      operator++();
      return tmp;
    }

    reference operator*() const { return iterator_.operator*(); }
    pointer operator->() const { return iterator_.operator->(); }

    bool operator==(joining_iterator const & other) const {
      return joiner_ == other.joiner_ && iterator_ == other.iterator_;
    }

    bool operator!=(joining_iterator const & other) const {
      return !(operator==(other));
    }

    end_aware_iterator<join_iter> join_iterator() const { return joiner_; }
    end_aware_iterator<iter_type> element_iterator() const { return iterator_; }

  private:
    void update_iterator() {
      while (!joiner_.done() && std::begin(*joiner_) == std::end(*joiner_)) {
        ++joiner_;
      }
      if (!joiner_.done()) {
        iterator_ = {std::begin(*joiner_), std::end(*joiner_)};
      }
    }

    end_aware_iterator<join_iter> joiner_;
    end_aware_iterator<iter_type> iterator_;
  };
}