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

#pragma once

#include <string>
#include <tuple>
#include <utility>

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

#include <iterator/detail/macro.h>

namespace iterator::recursive {
struct unbounded {
  template <typename It>
  static constexpr recursion_type const value = iter_typeclass<It>;
  using next = unbounded;
  static constexpr size_t size = std::numeric_limits<size_t>::max();
};

template <size_t N, size_t I> struct bounded {
  template <typename It>
  static constexpr recursion_type const value =
      I == N ? recursion_type::END : iter_typeclass<It>;
  using next = std::conditional_t<I == N, void, bounded<N, I + 1>>;
  static constexpr size_t size = N;
};

template <typename It, typename Bnd = unbounded,
          recursion_type = Bnd::template value<It>>
struct tuple_expander;

template <typename It, typename Bnd>
struct tuple_expander<It, Bnd, recursion_type::END> {
  using iter = std::tuple<end_aware_iterator<It>>;
  decltype(auto) get(It iter) const {
    if constexpr (associative_value<std::iter_value_t<It>>) {
      return std::tie(iter->first, iter->second);
    } else {
      return std::tie(*iter);
    }
  }
};

template <typename... Ts>
using tuple_cat_t = decltype(std::tuple_cat(VAL(Ts)...));

template <typename It, typename Bnd>
struct tuple_expander<It, Bnd, recursion_type::THRU> {
  using next = decltype(std::begin(*VAL(It)));
  using iter =
      tuple_cat_t<std::tuple<end_aware_iterator<It>>,
                  typename tuple_expander<next, typename Bnd::next>::iter>;
  auto get(It) const { return std::make_tuple(); }
};

template <typename It, typename Bnd>
struct tuple_expander<It, Bnd, recursion_type::ASSOC> {
  using next = decltype(std::begin(VAL(It)->second));
  using iter =
      tuple_cat_t<std::tuple<end_aware_iterator<It>>,
                  typename tuple_expander<next, typename Bnd::next>::iter>;
  auto get(It iter) const { return std::tie(iter->first); };
};

/**
 * @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 Bnd The bounding type, representing how many layers this iterator
 * is willing to delve in the parent object.
 */
template <typename It, typename Bnd>
class rimpl : public facade<rimpl<It, Bnd>> {
public:
  using sentinel_type = sentinel_t;
  using iters_t = typename tuple_expander<It, Bnd>::iter;
  static constexpr size_t n_iters = std::tuple_size_v<iters_t>;
  static constexpr size_t size = std::min(n_iters, Bnd::size);

private:
  iters_t impl_;

public:
  rimpl() = default;
  rimpl(end_aware_iterator<It> iter) { assign<0>(iter); }
  template <typename Ot>
  rimpl(end_aware_iterator<Ot> other) : rimpl(end_aware_iterator<It>(other)) {}
  template <typename Ot>
  rimpl(rimpl<Ot, Bnd> other) : rimpl(end_aware_iterator<Ot>(other)) {}

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

  decltype(auto) dereference() const {
    // 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(this->build_tuple())> == 1) {
      return *std::get<size - 1>(impl_);
    } else {
      return build_tuple();
    }
  }

  void increment() { increment_i(); }

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

  // Used by std::get, don't use elsewhere...
  auto const & impl() const { return impl_; }

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

  template <size_t I> decltype(auto) get() const {
    auto it = std::get<I>(impl_);
    // In the case of a bounded recursive iterator, I need to ensure that the
    // effectively terminal iterator is treated as such even if there is still
    // iteration to be had.
    if constexpr (I == size - 1) {
      return tuple_expander<decltype(it), unbounded, recursion_type::END>{}.get(
          it);
    } else {
      return tuple_expander<decltype(it)>{}.get(it);
    }
  }

  template <size_t... Is>
  decltype(auto) build_tuple(std::index_sequence<Is...>) const {
    return std::tuple_cat(get<Is>()...);
  }

  decltype(auto) build_tuple() const {
    return build_tuple(std::make_index_sequence<size>());
  }

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

  template <size_t I, typename C> void assign(C && collection) {
    assign<I>(end_aware_iterator(std::forward<C>(collection)));
  }

  template <size_t I, typename K, typename V>
  void assign(std::pair<K const, V> const & pair) {
    assign<I>(pair.second);
  }
  template <size_t I, typename K, typename V>
  void assign(std::pair<K const, V> & pair) {
    assign<I>(pair.second);
  }
};
}

namespace std {
template <std::size_t I, typename It>
auto get(::iterator::recursive_iterator<It> const & iter) {
  return ::std::get<I>(iter.impl());
}

template <std::size_t I, typename It, std::size_t N>
auto get(::iterator::recursive_iterator_n<It, N> const & iter) {
  static_assert(I < N, "Cannot get past bounding level");
  return ::std::get<I>(iter.impl());
}
}

template <typename C> auto make_recursive_iterator(C && collect) {
  return iterator::recursive_iterator<iterator::iterator_t<C>>(collect);
}

template <std::size_t Max, typename C>
auto make_recursive_iterator(C && collect) {
  return iterator::recursive_iterator_n<iterator::iterator_t<C>, Max>(collect);
}

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 =
        make_recursive_iterator<N>(std::views::all(std::forward<Rng>(rng)));
    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 =
        make_recursive_iterator(std::views::all(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>