#pragma once

#if !defined(MIJIN_UTIL_ITERATORS_HPP_INCLUDED)
#define MIJIN_UTIL_ITERATORS_HPP_INCLUDED 1

#include <cstddef>
#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 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<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>
{
    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
    {
        EnumeratingIterator 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 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));
}
}
} // namespace mijin

#endif // MIJIN_UTIL_ITERATORS_HPP_INCLUDED