//
//  recursive_iterator.h
//  iterator
//
//  Created by Sam Jaffe on 2/17/17.
//

#pragma once

#include <ranges>
#include <tuple>
#include <utility>

#include <iterator/concepts.h>
#include <iterator/detail/projection_tuple.h>
#include <iterator/detail/recursive_expander.h>
#include <iterator/end_aware_iterator.h>
#include <iterator/facade.h>
#include <iterator/forwards.h>

#include <iterator/detail/macro.h>

namespace iterator {
template <typename Tuple, typename Indices> class recursive_iterator_base;
template <typename... It, size_t... Is>
class recursive_iterator_base<std::tuple<It...>, std::index_sequence<Is...>>
    : public std::tuple<It...> {
public:
  static constexpr size_t LastIndex = sizeof...(It) - 1;
  template <size_t I>
  using iterator_type = std::tuple_element_t<I, std::tuple<It...>>;
  template <size_t I> using value_type = std::iter_value_t<iterator_type<I>>;

public:
  template <typename T> operator end_aware_iterator<T>() const {
    return std::get<end_aware_iterator<T>>(*this);
  }

  decltype(auto) dereference() const {
    auto rval = std::tuple_cat(get<Is>()...);
    // Special Case Handling for circumstances where at least everything up to
    // the deepest nested container is non-associative. In this case, we don't
    // want to transmute our single element/association into a tuple, since
    // there's no benefit from that.
    if constexpr (std::tuple_size_v<decltype(rval)> == 1) {
      return std::get<0>(rval); // May be a reference
    } else {
      return rval; // Tuple-of-references
    }
  }

  void increment() { increment<>(); }

  bool at_end() const { return std::get<0>(*this).at_end(); }
  bool equal_to(recursive_iterator_base const & other) const {
    return *this == other;
  }

protected:
  recursive_iterator_base() = default;
  recursive_iterator_base(iterator_type<0> iter) { assign<0>(iter); }

private:
  template <size_t I = LastIndex> bool increment() {
    auto & iter = std::get<I>(*this);
    if (iter.at_end()) { return false; } // Make sure we don't go OOB
    ++iter;
    if constexpr (I > 0) {
      while (iter.at_end() && increment<I - 1>()) {
        assign<I>(*std::get<I - 1>(*this));
      }
    }
    return !iter.at_end();
  }

  template <size_t I> decltype(auto) get() const {
    auto iter = std::get<I>(*this);
    if constexpr (I + 1 == sizeof...(It)) {
      if constexpr (Assoc<value_type<I>>) {
        return std::tie(iter->first, iter->second);
      } else {
        return std::tie(*iter);
      }
    } else if constexpr (Assoc<value_type<I>>) {
      return std::tie(iter->first);
    } else {
      return std::make_tuple();
    }
  }

  template <size_t I, typename T> void assign(end_aware_iterator<T> it) {
    std::get<I>(*this) = it;
    if constexpr (I < LastIndex) {
      if (!it.at_end()) { assign<I + 1>(*it); }
    }
  }

  template <size_t I, typename T> void assign(T && value) {
    if constexpr (Range<T>) {
      assign<I>(end_aware_iterator(std::forward<T>(value)));
    } else {
      assign<I>(value.second);
    }
  }
};

template <typename It, typename MaxDepth,
          typename Projs = detail::Projections<>>
struct recursive_iterator_helper {
  static typename detail::ProjectionExpander<It, Projs>::type const & _projs;
  using Projections = decltype(detail::Projections(_projs));

  using iterator_tuple =
      typename detail::RecursiveExpander<It, Projections, MaxDepth>::type;

  static constexpr auto extent = std::tuple_size_v<iterator_tuple>;
  using indices = decltype(std::make_index_sequence<extent>());

  using type = recursive_iterator_base<iterator_tuple, indices>;
};

/**
 * @brief An iterator type for nested collections, allowing you to treat it as
 * a single-layer collection.
 *
 * In order to provide a simple interface, if an associative container is used
 * in the chain, the type returned by operator*() is a tuple. If multiple
 * associative containers are nested, then the tuple will be of the form
 * std::tuple<key1, key2, ..., keyN, value>. To avoid copies, and allow
 * editting of underlying values, the tuple contains references.
 *
 * @tparam It The iterator type of the top-level collection.
 * @tparam MaxDepth The bounding type, representing how many layers this
 * iterator is willing to delve in the parent object.
 */
template <typename It, typename MaxDepth>
class recursive_iterator : public recursive_iterator_helper<It, MaxDepth>::type,
                           public facade<recursive_iterator<It, MaxDepth>> {
public:
  using sentinel_type = sentinel_t;

public:
  recursive_iterator() = default;

  explicit recursive_iterator(Range auto & range, MaxDepth = {})
      : recursive_iterator(end_aware_iterator(range)) {}

  explicit recursive_iterator(end_aware_iterator<It> iter, MaxDepth = {})
      : recursive_iterator::recursive_iterator_base(iter) {}

  template <typename Ot>
  explicit recursive_iterator(end_aware_iterator<Ot> other, MaxDepth = {})
      : recursive_iterator(end_aware_iterator<It>(other)) {}

  template <typename Ot>
  explicit recursive_iterator(recursive_iterator<Ot, MaxDepth> other)
      : recursive_iterator(end_aware_iterator<Ot>(other)) {}
};

template <typename Range, typename MaxDepth>
recursive_iterator(Range &, MaxDepth)
    -> recursive_iterator<iterator_t<Range>, MaxDepth>;
template <typename Range>
recursive_iterator(Range &) -> recursive_iterator<iterator_t<Range>, unbounded>;
}

namespace std {
template <size_t I, typename It, typename MaxDepth>
auto get(::iterator::recursive_iterator<It, MaxDepth> const & iter) {
  using return_type = std::decay_t<decltype(iter)>::template iterator_type<I>;
  return static_cast<return_type>(iter);
}
}

namespace iterator::views {
template <size_t N>
struct recursive_n_fn : std::ranges::range_adaptor_closure<recursive_n_fn<N>> {
  template <std::ranges::range Rng> auto operator()(Rng && rng) const {
    auto begin = recursive_iterator(std::forward<Rng>(rng), bounded<N>{});
    return std::ranges::subrange(begin, decltype(begin)());
  }
};

struct recursive_fn : std::ranges::range_adaptor_closure<recursive_fn> {
  template <std::ranges::range Rng> auto operator()(Rng && rng) const {
    auto begin = recursive_iterator(std::forward<Rng>(rng));
    return std::ranges::subrange(begin, decltype(begin)());
  }
};

template <size_t N> constexpr recursive_n_fn<N> recursive_n{};
constexpr recursive_fn recursive;
}

#include <iterator/detail/undef.h>