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

#pragma once

#include <iterator>
#include <memory>
#include <utility>

#include <iterator/end_aware_iterator.h>
#include <iterator/facade.h>
#include <iterator/forwards.h>

#include <iterator/detail/macro.h>

namespace iterator::joining {
template <typename Iter, typename = void> class iterator {
protected:
  using inner_t = DEREF_TYPE(Iter);

public:
  iterator() = default;

  template <typename I>
  iterator(iterator<I> const & other)
      : iterator(other.joiner_, other.element_) {}

  template <typename C, typename = std::enable_if_t<detail::is_container_v<C>>>
  iterator(C && container)
      : iterator(end_aware_iterator(std::forward<C>(container))) {}

  iterator(end_aware_iterator<Iter> join) : iterator(join, {}, true) {}

  iterator(end_aware_iterator<Iter> join,
           end_aware_iterator<detail::iter<inner_t>> elem, bool update = false)
      : joiner_(join), element_(elem) {
    if (update) { update_iterator(); }
  }

protected:
  void update_iterator() {
    while (!joiner_.at_end() &&
           (element_ = end_aware_iterator(*joiner_)).at_end()) {
      ++joiner_;
    }
  }

protected:
  template <typename, typename> friend class iterator;
  end_aware_iterator<Iter> joiner_;
  end_aware_iterator<detail::iter<inner_t>> element_;
};

template <typename Iter>
class iterator<Iter, std::enable_if_t<detail::is_rvalue_iterator_v<Iter>>> {
protected:
  using inner_t = DEREF_TYPE(Iter);

public:
  iterator() = default;
  template <typename C, typename = std::enable_if_t<detail::is_container_v<C>>>
  iterator(C && container) : joiner_(std::forward<C>(container)) {
    update_iterator();
  }

  iterator(end_aware_iterator<Iter> join) : joiner_(join) { update_iterator(); }

protected:
  void update_iterator() {
    while (!joiner_.at_end() && sync().at_end()) {
      ++joiner_;
    }
  }

private:
  end_aware_iterator<detail::iter<inner_t>> const & sync() {
    if (joiner_.at_end()) { return element_ = {}; }
    cache_ = std::make_shared<inner_t>(*joiner_);
    element_ = end_aware_iterator(*cache_);
    return element_;
  }

protected:
  end_aware_iterator<Iter> joiner_;
  std::shared_ptr<inner_t> cache_;
  end_aware_iterator<detail::iter<inner_t>> element_;
};
}

namespace iterator {
template <typename Iter>
class joining_iterator : public facade<joining_iterator<Iter>>,
                         public joining::iterator<Iter> {
public:
  using super_t = joining::iterator<Iter>;
  using sentinel_type = sentinel_t;

public:
  using super_t::super_t;

  template <typename C, typename = std::enable_if_t<detail::is_container_v<C>>>
  joining_iterator(C && container) : super_t(std::forward<C>(container)) {}

  joining_iterator(end_aware_iterator<Iter> join) : super_t(join) {}

  void increment() {
    if ((++super_t::element_).at_end()) {
      ++super_t::joiner_;
      super_t::update_iterator();
    }
  }

  decltype(auto) dereference() const { return *super_t::element_; }

  bool at_end() const { return join_iterator().at_end(); }
  bool equal_to(joining_iterator const & other) const {
    return join_iterator() == other.join_iterator() &&
           element_iterator() == other.element_iterator();
  }

  auto const & join_iterator() const { return super_t::joiner_; }
  auto const & element_iterator() const { return super_t::element_; }
};

template <typename C>
joining_iterator(C &&) -> joining_iterator<detail::iter<C>>;
template <typename JI>
joining_iterator(end_aware_iterator<JI>) -> joining_iterator<JI>;
}

MAKE_ITERATOR_FACADE_TYPEDEFS_T(::iterator::joining_iterator);

#include <iterator/detail/undef.h>