//
//  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>

namespace iterator {
template <typename Tuple, typename Projs, typename Indices> class RecursiveBase;
template <typename... It, typename... Projs, size_t... Is>
class RecursiveBase<std::tuple<It...>, detail::Projections<Projs...>,
                    std::index_sequence<Is...>>
    : public std::tuple<It...>, private detail::Projections<Projs...> {
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 EndAwareIterator<T>() const {
    return std::get<EndAwareIterator<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(RecursiveBase const & other) const { return *this == other; }

protected:
  RecursiveBase() = default;
  RecursiveBase(iterator_type<0> iter, Projs... projs)
      : detail::Projections<Projs...>(projs...) {
    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)) {
      auto && rval = (*this)(*iter, std::integral_constant<size_t, I>{});
      if constexpr (Assoc<std::decay_t<decltype(rval)>>) {
        // Needed for tuple-cat to preserve references
        return std::tie(rval.first, rval.second);
      } else {
        return std::tie(rval); // Cannot wrap tie(pair) and expect it to work
      }
    } 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(EndAwareIterator<T> it) {
    std::get<I>(*this) = it;
    if constexpr (I < LastIndex) {
      if (!it.at_end()) {
        assign<I + 1>((*this)(*it, std::integral_constant<size_t, I>{}));
      }
    }
  }

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

template <typename It, typename MaxDepth,
          typename Projs = detail::Projections<>>
struct RecursiveHelper {
  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 = RecursiveBase<iterator_tuple, Projections, 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 RecursiveIterator : public RecursiveHelper<It, MaxDepth>::type,
                          public Facade<RecursiveIterator<It, MaxDepth>> {
public:
  using sentinel_type = sentinel_t;

public:
  RecursiveIterator() = default;

  explicit RecursiveIterator(Range auto & range, MaxDepth = {})
      : RecursiveIterator(EndAwareIterator(range)) {}

  explicit RecursiveIterator(EndAwareIterator<It> iter, MaxDepth = {})
      : RecursiveIterator::RecursiveBase(iter) {}

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

  template <typename Ot>
  explicit RecursiveIterator(RecursiveIterator<Ot, MaxDepth> other)
      : RecursiveIterator(EndAwareIterator<Ot>(other)) {}
};

template <typename Range, typename MaxDepth>
RecursiveIterator(Range &, MaxDepth)
    -> RecursiveIterator<iterator_t<Range>, MaxDepth>;
template <typename Range>
RecursiveIterator(Range &) -> RecursiveIterator<iterator_t<Range>, unbounded>;
} // namespace iterator

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

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 = RecursiveIterator(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 = RecursiveIterator(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;
} // namespace iterator::views