#pragma once

#if !defined(MIJIN_UTIL_ITERATORS_HPP_INCLUDED)
#define MIJIN_UTIL_ITERATORS_HPP_INCLUDED 1

#include <cstddef>
#include <functional>
#include <optional>
#include <span>
#include <string_view>
#include <tuple>
#include <variant>

namespace mijin
{
struct RangeAdapter {};

template<typename T>
struct RangeRef
{
    T& range;

    decltype(auto) begin() const
    {
        return std::begin(range);
    }

    decltype(auto) end() const
    {
        return std::end(range);
    }
};

template<typename TIterator>
struct RangeRef<std::pair<TIterator, TIterator>>
{
    std::pair<TIterator, TIterator> iterators;

    TIterator begin() const
    {
        return iterators.first;
    }

    TIterator end() const
    {
        return iterators.second;
    }
};

// in case of pass-by-reference
template<typename TIterator>
struct RangeRef<std::pair<TIterator, TIterator>&> : RangeRef<std::pair<TIterator, TIterator>> {};

static_assert(std::is_same_v<decltype(std::declval<RangeRef<std::pair<int*, int*>>>().iterators.first), int*>, "RangeRef for iterator pair not working.");
static_assert(std::is_same_v<decltype(std::declval<RangeRef<std::pair<std::string_view::iterator, std::string_view::iterator>&>>().iterators.first), std::string_view::iterator>, "RangeRef for iterator pair not working.");

template<typename T>
struct is_range_adapter_type : std::false_type {};

template<typename T> requires std::is_base_of_v<RangeAdapter, T>
struct is_range_adapter_type<T> : std::true_type {};

// allow passing in spans and string_views by value
template<typename T, std::size_t Extent>
struct is_range_adapter_type<std::span<T, Extent>> : std::true_type {};

template<typename CharT, typename Traits>
struct is_range_adapter_type<std::basic_string_view<CharT, Traits>> : std::true_type {};

template<typename T>
concept RangeAdapterType = is_range_adapter_type<std::decay_t<T>>::value;

static_assert(!is_range_adapter_type<std::tuple<int, int>>::value);
static_assert(is_range_adapter_type<RangeAdapter>::value);

template<RangeAdapterType T>
struct RangeRef<T>
{
    std::remove_reference_t<T> range;

    decltype(auto) begin() const
    {
        return range.begin();
    }

    decltype(auto) end() const
    {
        return range.end();
    }
};

template<typename TIdx, typename TIterator>
struct EnumeratingIterator
{
    TIdx idx;
    TIterator base;

    EnumeratingIterator(TIdx idx_, TIterator base_) noexcept : idx(idx_), base(base_) {}
    EnumeratingIterator(const EnumeratingIterator&) noexcept = default;

    EnumeratingIterator& operator=(const EnumeratingIterator&) noexcept = default;

    auto operator*() const noexcept
    {
        return std::tie(idx, *base);
    }

    EnumeratingIterator& operator++() noexcept
    {
        ++idx;
        ++base;
        return *this;
    }

    EnumeratingIterator operator++(int) noexcept
    {
        EnumeratingIterator copy(*this);
        ++(*this);
        return copy;
    }

    EnumeratingIterator& operator--() noexcept
    {
        --idx;
        --base;
        return *this;
    }

    EnumeratingIterator operator--(int) noexcept
    {
        EnumeratingIterator copy(*this);
        --(*this);
        return copy;
    }

    bool operator==(const EnumeratingIterator& other) const noexcept
    {
        return base == other.base; // note: ignoring idx so we don't have to find it out for end()
    }

    bool operator!=(const EnumeratingIterator& other) const noexcept
    {
        return base != other.base; // note: ignoring idx so we don't have to find it out for end()
    }
};
template<typename TIdx, typename TIterator>
EnumeratingIterator(TIdx, TIterator) -> EnumeratingIterator<TIdx, TIterator>;

template<typename TIdx, typename TIterable>
struct Enumeratable : RangeAdapter
{
    RangeRef<TIterable> base;

    auto begin() const noexcept
    {
        return EnumeratingIterator(TIdx(0), base.begin());
    }

    auto end() const noexcept
    {
        return EnumeratingIterator(TIdx(0), base.end());
    }
};

template<typename TIdx = std::size_t, typename TIterable>
Enumeratable<TIdx, TIterable> enumerate(TIterable&& iterable)
{
    return {.base = {std::forward<TIterable>(iterable)}};
}

template<typename TIterator>
struct ReplacingIterator
{
    using difference_type = std::ptrdiff_t;
    using value_type = typename std::iterator_traits<TIterator>::value_type;
    using pointer = std::add_const_t<value_type>*;
    using reference = std::add_const_t<value_type>&;
    using iterator_category = std::bidirectional_iterator_tag; // TODO?

    TIterator base;
    value_type what;
    value_type with;

    ReplacingIterator(TIterator base_, value_type what_, value_type with_) noexcept : base(base_), what(what_), with(with_) {}
    ReplacingIterator(const ReplacingIterator&) noexcept = default;

    ReplacingIterator& operator=(const ReplacingIterator&) noexcept = default;

    pointer operator->() const noexcept
    {
        if (*base == what) {
            return &with;
        }
        return &*base;
    }

    reference operator*() const noexcept
    {
        if (*base == what) {
            return with;
        }
        return *base;
    }

    ReplacingIterator& operator++() noexcept
    {
        ++base;
        return *this;
    }

    ReplacingIterator operator++(int) noexcept
    {
        ReplacingIterator copy(*this);
        ++(*this);
        return copy;
    }

    ReplacingIterator& operator--() noexcept
    {
        --base;
        return *this;
    }

    ReplacingIterator operator--(int) noexcept
    {
        ReplacingIterator copy(*this);
        --(*this);
        return copy;
    }

    bool operator==(const ReplacingIterator& other) const noexcept
    {
        return what == other.what && with == other.with && base == other.base;
    }

    bool operator!=(const ReplacingIterator& other) const noexcept
    {
        return !(*this == other);
    }
};

template<typename TIterable>
struct ReplacingRange : RangeAdapter
{
    using value_type = typename std::iterator_traits<decltype(std::begin(std::declval<TIterable>()))>::value_type;
    
    RangeRef<TIterable> base;
    value_type what;
    value_type with;

    auto begin() const noexcept
    {
        return ReplacingIterator(base.begin(), what, with);
    }

    auto end() const noexcept
    {
        return ReplacingIterator(base.end(), what, with);
    }
};

template<typename TIterable>
auto replace(
    TIterable&& iterable,
    typename std::iterator_traits<decltype(std::begin(std::declval<TIterable>()))>::value_type what,
    typename std::iterator_traits<decltype(std::begin(std::declval<TIterable>()))>::value_type with)
{
    return ReplacingRange<TIterable>{
        .base = {std::forward<TIterable>(iterable)},
        .what = std::move(what),
        .with = std::move(with)
    };
}

template<typename TIterator, typename TFunctor>
struct MappingIterator
{
    using orig_value_type = typename std::iterator_traits<TIterator>::value_type;
    using difference_type = std::ptrdiff_t;
    using value_type = std::invoke_result_t<TFunctor, orig_value_type>;
    using reference = value_type&;
    using iterator_category = std::bidirectional_iterator_tag; // TODO?

    static_assert(!std::is_same_v<value_type, void>, "Invalid mapping result type.");

    TIterator base;
    TFunctor functor;
    mutable std::optional<value_type> result;

    MappingIterator(TIterator base_, TFunctor functor_) noexcept : base(base_), functor(std::move(functor_)) {}
    MappingIterator(const MappingIterator&) noexcept = default;
    MappingIterator(MappingIterator&&) noexcept = default;

    MappingIterator& operator=(const MappingIterator&) noexcept = default;
    MappingIterator& operator=(MappingIterator&&) noexcept = default;

    reference operator*() const noexcept
    {
        if (!result.has_value()) {
            result = functor(*base);
        }
        return *result;
    }

    MappingIterator& operator++() noexcept
    {
        ++base;
        result = std::nullopt;
        return *this;
    }

    MappingIterator operator++(int) noexcept
    {
        MappingIterator copy(*this);
        ++(*this);
        return copy;
    }

    MappingIterator& operator--() noexcept
    {
        --base;
        result = std::nullopt;
        return *this;
    }

    MappingIterator operator--(int) noexcept
    {
        MappingIterator copy(*this);
        --(*this);
        return copy;
    }

    bool operator==(const MappingIterator& other) const noexcept
    {
        return base == other.base; // TODO: compare functor? -> doesn't always work and may be useless
    }

    bool operator!=(const MappingIterator& other) const noexcept
    {
        return !(*this == other);
    }
};

template<typename TIterable, typename TFunctor>
struct MappingRange : RangeAdapter
{
    // using value_type = typename std::iterator_traits<decltype(std::begin(std::declval<TIterable>()))>::value_type;
    using value_type = typename std::iterator_traits<decltype(std::declval<RangeRef<TIterable>>().begin())>::value_type;

    RangeRef<TIterable> base;
    TFunctor functor;

    auto begin() const noexcept
    {
        return MappingIterator(base.begin(), functor);
    }

    auto end() const noexcept
    {
        return MappingIterator(base.end(), functor);
    }
};

template<typename TIterable, typename TFunctor>
auto map(TIterable&& iterable, TFunctor&& functor)
{
    return MappingRange<TIterable, TFunctor>{
        .base = {std::forward<TIterable>(iterable)},
        .functor = std::forward<TFunctor>(functor)
    };
}

template<typename TIterator, typename TFunctor>
struct OptionalMappingIterator
{
    using orig_value_type = typename std::iterator_traits<TIterator>::value_type;
    using difference_type = std::ptrdiff_t;
    using optional_type = std::invoke_result_t<TFunctor, orig_value_type>;
    using value_type = std::decay_t<decltype(std::declval<optional_type>().value())>;
    using iterator_category = std::forward_iterator_tag; // TODO?
    using reference = value_type&;

    static_assert(!std::is_same_v<value_type, void>, "Invalid mapping result type.");

    TIterator base;
    TIterator end;
    TFunctor functor;
    mutable optional_type result; // must be mutable so dereferencing can stay const, not nice :/

    OptionalMappingIterator(TIterator base_, TIterator end_, TFunctor functor_) noexcept : base(base_), end(end_), functor(std::move(functor_)) {
        if (base != end)
        {
            result = functor(*base);
            if (!result.has_value()) {
                ++(*this);
            }
        }
    }
    OptionalMappingIterator(const OptionalMappingIterator&) noexcept = default;
    OptionalMappingIterator(OptionalMappingIterator&&) noexcept = default;

    OptionalMappingIterator& operator=(const OptionalMappingIterator&) noexcept = default;
    OptionalMappingIterator& operator=(OptionalMappingIterator&&) noexcept = default;

    reference operator*() const noexcept
    {
        return *result;
    }

    OptionalMappingIterator& operator++() noexcept
    {
        do
        {
            ++base;
            result = base != end ? functor(*base) : std::nullopt;
        } while (base != end && !result.has_value());
        return *this;
    }

    OptionalMappingIterator operator++(int) noexcept
    {
        OptionalMappingIterator copy(*this);
        ++(*this);
        return copy;
    }

    bool operator==(const OptionalMappingIterator& other) const noexcept
    {
        return base == other.base && end == other.end; // TODO: compare functor? -> doesn't always work and may be useless
    }

    bool operator!=(const OptionalMappingIterator& other) const noexcept
    {
        return !(*this == other);
    }
};

template<typename TIterable, typename TFunctor>
struct OptionalMappingRange : RangeAdapter
{
    // using value_type = typename std::iterator_traits<decltype(std::begin(std::declval<TIterable>()))>::value_type;
    using value_type = typename std::iterator_traits<decltype(std::declval<RangeRef<TIterable>>().begin())>::value_type;

    RangeRef<TIterable> base;
    TFunctor functor;

    auto begin() const noexcept
    {
        return OptionalMappingIterator(base.begin(), base.end(), functor);
    }

    auto end() const noexcept
    {
        return OptionalMappingIterator(base.end(), base.end(), functor);
    }
};

template<typename TIterable, typename TFunctor>
auto mapOptional(TIterable&& iterable, TFunctor&& functor)
{
    return OptionalMappingRange<TIterable, TFunctor>{
        .base = {std::forward<TIterable>(iterable)},
        .functor = std::forward<TFunctor>(functor)
    };
}

template<typename TIterator, typename TPredicate>
struct FilteringIterator
{
    using difference_type = typename std::iterator_traits<TIterator>::difference_type;
    using value_type = typename std::iterator_traits<TIterator>::value_type;
    using iterator_category = std::forward_iterator_tag; // TODO?
    using reference = typename std::iterator_traits<TIterator>::reference;

    TIterator base;
    TIterator end;
    TPredicate predicate;

    FilteringIterator(TIterator base_, TIterator end_, TPredicate predicate_) noexcept : base(base_), end(end_), predicate(std::move(predicate_)) {
        if (base != end && !predicate(*base)) {
            ++(*this);
        }
    }
    FilteringIterator(const FilteringIterator&) noexcept = default;
    FilteringIterator(FilteringIterator&&) noexcept = default;

    FilteringIterator& operator=(const FilteringIterator&) noexcept = default;
    FilteringIterator& operator=(FilteringIterator&&) noexcept = default;

    reference operator*() const noexcept
    {
        return *base;
    }

    FilteringIterator& operator++() noexcept
    {
        do
        {
            ++base;
        } while (base != end && !predicate(*base));
        return *this;
    }

    FilteringIterator operator++(int) noexcept
    {
        FilteringIterator copy(*this);
        ++(*this);
        return copy;
    }

    bool operator==(const FilteringIterator& other) const noexcept
    {
        return base == other.base && end == other.end; // TODO: compare predicate?
    }

    bool operator!=(const FilteringIterator& other) const noexcept
    {
        return !(*this == other);
    }
};

template<typename TIterable, typename TPredicate>
struct FilteringRange : RangeAdapter
{
    // using value_type = typename std::iterator_traits<decltype(std::begin(std::declval<TIterable>()))>::value_type;
    using value_type = typename std::iterator_traits<decltype(std::declval<RangeRef<TIterable>>().begin())>::value_type;

    RangeRef<TIterable> base;
    TPredicate predicate;

    auto begin() const noexcept
    {
        return FilteringIterator(base.begin(), base.end(), predicate);
    }

    auto end() const noexcept
    {
        return FilteringIterator(base.end(), base.end(), predicate);
    }
};

template<typename TIterable, typename TPredicate>
auto filter(TIterable&& iterable, TPredicate&& predicate)
{
    return FilteringRange<TIterable, TPredicate>{
        .base = {std::forward<TIterable>(iterable)},
        .predicate = std::forward<TPredicate>(predicate)
    };
}

template<typename TFirstIterator, typename TSecondIterator>
struct ChainingIterator
{
    using difference_type = std::ptrdiff_t;
    using value_type = typename std::iterator_traits<TFirstIterator>::value_type;
    using pointer = std::add_const_t<value_type>*;
    using reference = std::add_const_t<value_type>&;
    using iterator_category = std::bidirectional_iterator_tag; // TODO?

    TFirstIterator firstBase;
    TFirstIterator firstEnd;
    TSecondIterator secondBase;
    TSecondIterator secondBegin;

    ChainingIterator(TFirstIterator firstBase_, TFirstIterator firstEnd_, TSecondIterator secondBase_, TSecondIterator secondBegin_) noexcept
        : firstBase(firstBase_), firstEnd(firstEnd_), secondBase(secondBase_), secondBegin(secondBegin_) {}
    ChainingIterator(const ChainingIterator&) noexcept = default;

    ChainingIterator& operator=(const ChainingIterator&) noexcept = default;

    pointer operator->() const noexcept
    {
        if (firstBase == firstEnd)
        {
            return &*secondBase;
        }
        return &*firstBase;
    }

    reference operator*() const noexcept
    {
        if (firstBase == firstEnd)
        {
            return *secondBase;
        }
        return *firstBase;
    }

    ChainingIterator& operator++() noexcept
    {
        if (firstBase == firstEnd) {
            ++secondBase;
        }
        else {
            ++firstBase;
        }
        return *this;
    }

    ChainingIterator operator++(int) noexcept
    {
        ChainingIterator copy(*this);
        ++(*this);
        return copy;
    }

    ChainingIterator& operator--() noexcept
    {
        if (secondBase == secondBegin) {
            --firstBase;
        }
        else {
            --secondBase;
        }
        return *this;
    }

    ChainingIterator operator--(int) noexcept
    {
        ChainingIterator copy(*this);
        --(*this);
        return copy;
    }

    bool operator==(const ChainingIterator& other) const noexcept
    {
        return firstBase == other.firstBase && secondBase == other.secondBase; // should be enough
    }

    bool operator!=(const ChainingIterator& other) const noexcept
    {
        return !(*this == other);
    }
};

template<typename TFirstRange, typename TSecondRange>
struct ChainedRange : RangeAdapter
{
    using value_type = typename std::iterator_traits<decltype(std::begin(std::declval<TFirstRange>()))>::value_type;
    
    RangeRef<TFirstRange> first;
    RangeRef<TSecondRange> second;

    auto begin() const noexcept
    {
        return ChainingIterator(first.begin(), first.end(), second.begin(), second.begin());
    }

    auto end() const noexcept
    {
        return ChainingIterator(first.end(), first.end(), second.end(), second.begin());
    }
};

template<typename TFirstRange, typename TSecondRange>
auto chain(TFirstRange&& firstRange, TSecondRange&& secondRange)
{
    return ChainedRange<TFirstRange, TSecondRange>{
        .first = {std::forward<TFirstRange>(firstRange)},
        .second = {std::forward<TSecondRange>(secondRange)}
    };
}

template<typename TFirstRange, typename TSecondRange, typename... TMoreRanges>
auto chain(TFirstRange&& firstRange, TSecondRange&& secondRange, TMoreRanges&&... moreRanges)
{
    return chain(std::forward<TFirstRange>(firstRange), chain(std::forward<TSecondRange>(secondRange), std::forward<TMoreRanges>(moreRanges)...));
}

template<typename TAs, typename TIterable>
TAs collect(TIterable&& iterable)
{
    return TAs(iterable.begin(), iterable.end());
}

namespace pipe
{
template<typename T>
struct Replace
{
    T what;
    T with;

    Replace(T what_, T with_) : what(std::move(what_)), with(std::move(with_)) {}
};
template<typename T>
Replace(T, T) -> Replace<T>;

template<typename TIterable>
auto operator|(TIterable&& iterable, Replace<typename std::iterator_traits<decltype(std::begin(std::declval<TIterable>()))>::value_type> repl)
{
    return replace(std::forward<TIterable>(iterable), std::move(repl.what), std::move(repl.with));
}

template<typename T>
struct Collect {};

template<typename TIterable, typename T>
auto operator|(TIterable&& iterable, Collect<T>)
{
    return collect<T>(std::forward<TIterable>(iterable));
}

struct Chain {};

template<typename TIterable>
struct Chain2
{
    RangeRef<TIterable> range;
};

template<typename TIterable>
auto operator|(TIterable&& iterable, Chain)
{
    return Chain2<TIterable>{.range = {.range = std::forward<TIterable>(iterable)}};
}

template<typename TFirstIterable, typename TSecondIterable>
auto operator|(Chain2<TFirstIterable> chain2, TSecondIterable&& secondIterable)
{
    return chain(chain2.range.range, std::forward<TSecondIterable>(secondIterable));
}

template<typename T>
struct Map
{
    T functor;

    explicit Map(T functor_) : functor(std::move(functor_)) {}
};
template<typename T>
Map(T) -> Map<T>;

template<typename TIterable, typename TFunctor>
auto operator|(TIterable&& iterable, Map<TFunctor> mapper)
{
    return map(std::forward<TIterable>(iterable), std::move(mapper.functor));
}

template<typename T>
struct MapOptional
{
    T functor;

    explicit MapOptional(T functor_) : functor(std::move(functor_)) {}
};
template<typename T>
MapOptional(T) -> MapOptional<T>;

template<typename TIterable, typename TFunctor>
auto operator|(TIterable&& iterable, MapOptional<TFunctor> mapper)
{
    return mapOptional(std::forward<TIterable>(iterable), std::move(mapper.functor));
}

template<typename T>
struct Filter
{
    T predicate;

    explicit Filter(T predicate_) : predicate(std::move(predicate_)) {}
};
template<typename T>
Filter(T) -> Filter<T>;

template<typename TIterable, typename TFilter>
auto operator|(TIterable&& iterable, Filter<TFilter> filterer)
{
    return filter(std::forward<TIterable>(iterable), std::move(filterer.predicate));
}

template<std::size_t idx>
struct Xth {};

using First = Xth<0>;
using Second = Xth<1>;

template<typename TIterable, std::size_t idx>
auto operator|(TIterable&& iterable, Xth<idx>)
{
    return map(std::forward<TIterable>(iterable), [](auto&& element) { return std::get<idx>(element); } );
}
}
} // namespace mijin

#endif // MIJIN_UTIL_ITERATORS_HPP_INCLUDED