#pragma once

#include <iterator>
#include "self_iterating_contianer.hpp"

namespace stream { namespace detail {  
  template <typename C> class join_iterator {
  public:
    using reference = typename C::const_iterator::reference;

  public:
    join_iterator(iterator<C> && f, iterator<C> && l)
        : start_(std::move(f)), finish_(std::move(l)) {
      if (start_ != finish_) {
        current_ = *start_;
        advance();
      }
    }

    reference operator*() { return *current_; }

    join_iterator & operator++() {
      ++current_;
      advance();
      return *this;
    }

    bool operator==(join_iterator const & rhs) const { return start_ == rhs.start_; }
  private:
    void advance() {
      while (!current_ && start_ != finish_) {
        if (++start_ == finish_) { break; }
        current_ = *start_;
      }
    }
    iterator<C> start_;
    iterator<C> finish_;
    self_iterating_container<C> current_;
  };

  template <typename C> class join_stream {
  public:
    using T = typename C::value_type;
    explicit join_stream(stream_base<C> const & sb) : source_(sb) {}

    iterator<T> begin() {
      return {join_iterator<C>{source_.begin(), source_.end()}};
    }
    iterator<T> end() {
      return {join_iterator<C>{source_.end(), source_.end()}};
    }

  private:
    stream_base<C> source_;
  };

  template <typename C>
  stream_base<typename C::value_type> make_join(stream_base<C> const & sb) {
    return std::make_shared<join_stream<C>>(sb);
  }

  template <typename T>
  template <typename F>
  stream_base<traits::fmapped_t<T, F>> stream_base<T>::flatmap(F && func) && {
    return make_join(std::move(*this).map(func));
  }

  template <typename T>
  template <typename F>
  stream_base<traits::fmapped_t<T, F>>
  stream_base<T>::flatmap(F && func) const & {
    return make_join(map(func));
  }
}}