#pragma once

namespace stream { namespace detail {
  template <typename T, typename R> class map_iterator {
  public:
    map_iterator(std::function<R(T const &)> f, iterator<T> && impl)
        : fun_(f), impl_(std::forward<iterator<T>>(impl)) {}

    R operator*() { return fun_(*impl_); }
    map_iterator & operator++() {
      ++impl_;
      return *this;
    }
    bool operator==(map_iterator const & rhs) const { return impl_ == rhs.impl_; }
  private:
    std::function<R(T const &)> fun_;
    iterator<T> impl_;
  };

  template <typename T, typename R> class map_stream {
  public:
    template <typename F>
    map_stream(F && func, stream_base<T> const & sb)
        : fun_(func), source_(sb) {}

    iterator<R> begin() { return {map_iterator<T, R>{fun_, source_.begin()}}; }

    iterator<R> end() { return {map_iterator<T, R>{fun_, source_.end()}}; }

  private:
    std::function<R(T const &)> fun_;
    stream_base<T> source_;
  };

  template <typename T>
  template <typename F>
  stream_base<traits::mapped_t<T, F>> stream_base<T>::map(F && func) && {
    return std::make_shared<map_stream<T, traits::mapped_t<T, F>>>(
        func, std::move(*this));
  }

  template <typename T>
  template <typename F>
  stream_base<traits::mapped_t<T, F>> stream_base<T>::map(F && func) const & {
    return std::make_shared<map_stream<T, traits::mapped_t<T, F>>>(func, *this);
  }
}}