///////////////////////////////////////////////////////////////////////////////
// Reference implementation of std::generator proposal P2168.
//
// See https://wg21.link/P2168 for details.
//
///////////////////////////////////////////////////////////////////////////////
// Copyright Lewis Baker, Corentin Jabot
//
// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0.
// (See accompanying file LICENSE or http://www.boost.org/LICENSE_1_0.txt)
///////////////////////////////////////////////////////////////////////////////

module;

#include <cassert>
#include <coroutine>
#include <exception>
#include <memory>
#include <ranges>
#include <type_traits>
#include <utility>

export module nihil.generator:generator;

import :elements_of;
import :manual_lifetime;
import :promise_base_alloc;
import :util;

namespace nihil {

export template <typename _Ref,
          typename _Value = std::remove_cvref_t<_Ref>,
          typename _Allocator = use_allocator_arg>
class generator;


template<typename _Ref>
struct __generator_promise_base
{
    template <typename _Ref2, typename _Value, typename _Alloc>
    friend class generator;

    __generator_promise_base* __root_;
    std::coroutine_handle<> __parentOrLeaf_;
    // Note: Using manual_lifetime here to avoid extra calls to exception_ptr
    // constructor/destructor in cases where it is not needed (i.e. where this
    // generator coroutine is not used as a nested coroutine).
    // This member is lazily constructed by the __yield_sequence_awaiter::await_suspend()
    // method if this generator is used as a nested generator.
    manual_lifetime<std::exception_ptr> __exception_;
    manual_lifetime<_Ref> __value_;

    explicit __generator_promise_base(std::coroutine_handle<> thisCoro) noexcept
        : __root_(this)
        , __parentOrLeaf_(thisCoro)
    {}

    ~__generator_promise_base() {
        if (__root_ != this) {
            // This coroutine was used as a nested generator and so will
            // have constructed its __exception_ member which needs to be
            // destroyed here.
            __exception_.destruct();
        }
    }

    std::suspend_always initial_suspend() noexcept {
        return {};
    }

    void return_void() noexcept {}

    void unhandled_exception() {
        if (__root_ != this) {
            __exception_.get() = std::current_exception();
        } else {
            throw;
        }
    }

    // Transfers control back to the parent of a nested coroutine
    struct __final_awaiter {
        bool await_ready() noexcept {
            return false;
        }

        template <typename _Promise>
        std::coroutine_handle<>
        await_suspend(std::coroutine_handle<_Promise> __h) noexcept {
            _Promise& __promise = __h.promise();
            __generator_promise_base& __root = *__promise.__root_;
            if (&__root != &__promise) {
                auto __parent = __promise.__parentOrLeaf_;
                __root.__parentOrLeaf_ = __parent;
                return __parent;
            }
            return std::noop_coroutine();
        }

        void await_resume() noexcept {}
    };

    __final_awaiter final_suspend() noexcept {
        return {};
    }

    std::suspend_always yield_value(_Ref&& __x)
            noexcept(std::is_nothrow_move_constructible_v<_Ref>) {
        __root_->__value_.construct((_Ref&&)__x);
        return {};
    }

    template <typename _T>
    requires
        (!std::is_reference_v<_Ref>) &&
        std::is_convertible_v<_T, _Ref>
    std::suspend_always yield_value(_T&& __x)
            noexcept(std::is_nothrow_constructible_v<_Ref, _T>) {
        __root_->__value_.construct((_T&&)__x);
        return {};
    }

    template <typename _Gen>
    struct __yield_sequence_awaiter {
        _Gen __gen_;

        __yield_sequence_awaiter(_Gen&& __g) noexcept
            // Taking ownership of the generator ensures frame are destroyed
            // in the reverse order of their execution.
            : __gen_((_Gen&&)__g) {
        }

        bool await_ready() noexcept {
            return false;
        }

        // set the parent, root and exceptions pointer and
        // resume the nested
        template<typename _Promise>
        std::coroutine_handle<>
        await_suspend(std::coroutine_handle<_Promise> __h) noexcept {
            __generator_promise_base& __current = __h.promise();
            __generator_promise_base& __nested = *__gen_.__get_promise();
            __generator_promise_base& __root = *__current.__root_;

            __nested.__root_ = __current.__root_;
            __nested.__parentOrLeaf_ = __h;

            // Lazily construct the __exception_ member here now that we
            // know it will be used as a nested generator. This will be
            // destroyed by the promise destructor.
            __nested.__exception_.construct();
            __root.__parentOrLeaf_ = __gen_.__get_coro();

            // Immediately resume the nested coroutine (nested generator)
            return __gen_.__get_coro();
        }

        void await_resume() {
            __generator_promise_base& __nestedPromise = *__gen_.__get_promise();
            if (__nestedPromise.__exception_.get()) {
                std::rethrow_exception(std::move(__nestedPromise.__exception_.get()));
            }
        }
    };

    template <typename _OValue, typename _OAlloc>
    __yield_sequence_awaiter<generator<_Ref, _OValue, _OAlloc>>
    yield_value(nihil::elements_of<generator<_Ref, _OValue, _OAlloc>> __g) noexcept {
        return std::move(__g).get();
    }

    template <std::ranges::range _Rng, typename _Allocator>
    __yield_sequence_awaiter<generator<_Ref, std::remove_cvref_t<_Ref>, _Allocator>>
    yield_value(nihil::elements_of<_Rng, _Allocator> && __x) {
        return [](std::allocator_arg_t, _Allocator, auto && __rng) -> generator<_Ref, std::remove_cvref_t<_Ref>, _Allocator> {
            for(auto && e: __rng)
                co_yield static_cast<decltype(e)>(e);
        }(std::allocator_arg, __x.get_allocator(), std::forward<_Rng>(__x.get()));
    }

    void resume() {
        __parentOrLeaf_.resume();
    }

    // Disable use of co_await within this coroutine.
    void await_transform() = delete;
};

template<typename _Generator, typename _ByteAllocator, bool _ExplicitAllocator = false>
struct __generator_promise;

template<typename _Ref, typename _Value, typename _Alloc, typename _ByteAllocator, bool _ExplicitAllocator>
struct __generator_promise<generator<_Ref, _Value, _Alloc>, _ByteAllocator, _ExplicitAllocator> final
    : public __generator_promise_base<_Ref>
    , public promise_base_alloc<_ByteAllocator> {
    __generator_promise() noexcept
    : __generator_promise_base<_Ref>(std::coroutine_handle<__generator_promise>::from_promise(*this))
    {}

    generator<_Ref, _Value, _Alloc> get_return_object() noexcept {
        return generator<_Ref, _Value, _Alloc>{
            std::coroutine_handle<__generator_promise>::from_promise(*this)
        };
    }

    using __generator_promise_base<_Ref>::yield_value;

    template <std::ranges::range _Rng>
    typename __generator_promise_base<_Ref>::template __yield_sequence_awaiter<generator<_Ref, _Value, _Alloc>>
    yield_value(nihil::elements_of<_Rng> && __x) {
        static_assert (!_ExplicitAllocator,
        "This coroutine has an explicit allocator specified with std::allocator_arg so an allocator needs to be passed "
        "explicitely to std::elements_of");
        return [](auto && __rng) -> generator<_Ref, _Value, _Alloc> {
            for(auto && e: __rng)
                co_yield static_cast<decltype(e)>(e);
        }(std::forward<_Rng>(__x.get()));
    }
};

template<typename _Alloc>
using __byte_allocator_t = typename std::allocator_traits<std::remove_cvref_t<_Alloc>>::template rebind_alloc<std::byte>;

} // namespace nihil

namespace std {

// Type-erased allocator with default allocator behaviour.
export template<typename _Ref, typename _Value, typename... _Args>
struct coroutine_traits<nihil::generator<_Ref, _Value>, _Args...> {
    using promise_type = nihil::__generator_promise<nihil::generator<_Ref, _Value>, std::allocator<std::byte>>;
};

// Type-erased allocator with std::allocator_arg parameter
export template<typename _Ref, typename _Value, typename _Alloc, typename... _Args>
struct coroutine_traits<nihil::generator<_Ref, _Value>, allocator_arg_t, _Alloc, _Args...> {
private:
    using __byte_allocator = nihil::__byte_allocator_t<_Alloc>;
public:
    using promise_type = nihil::__generator_promise<nihil::generator<_Ref, _Value>, __byte_allocator, true /*explicit Allocator*/>;
};

// Type-erased allocator with std::allocator_arg parameter (non-static member functions)
export template<typename _Ref, typename _Value, typename _This, typename _Alloc, typename... _Args>
struct coroutine_traits<nihil::generator<_Ref, _Value>, _This, allocator_arg_t, _Alloc, _Args...> {
private:
    using __byte_allocator = nihil::__byte_allocator_t<_Alloc>;
public:
    using promise_type = nihil::__generator_promise<nihil::generator<_Ref, _Value>, __byte_allocator,  true /*explicit Allocator*/>;
};

// Generator with specified allocator type
export template<typename _Ref, typename _Value, typename _Alloc, typename... _Args>
struct coroutine_traits<nihil::generator<_Ref, _Value, _Alloc>, _Args...> {
    using __byte_allocator = nihil::__byte_allocator_t<_Alloc>;
public:
    using promise_type = nihil::__generator_promise<nihil::generator<_Ref, _Value, _Alloc>, __byte_allocator>;
};

} // namespace std

namespace nihil {

// TODO :  make layout compatible promise casts possible
export template <typename _Ref, typename _Value, typename _Alloc>
class generator {
    using __byte_allocator = __byte_allocator_t<_Alloc>;
public:
    using promise_type = __generator_promise<generator<_Ref, _Value, _Alloc>, __byte_allocator>;
    friend promise_type;
private:
    using __coroutine_handle = std::coroutine_handle<promise_type>;
public:

    generator() noexcept = default;

    generator(generator&& __other) noexcept
        : __coro_(std::exchange(__other.__coro_, {}))
        , __started_(std::exchange(__other.__started_, false)) {
    }

    ~generator() noexcept {
        if (__coro_) {
            if (__started_ && !__coro_.done()) {
                __coro_.promise().__value_.destruct();
            }
            __coro_.destroy();
        }
    }

    generator& operator=(generator && g) noexcept {
        swap(g);
        return *this;
    }

    void swap(generator& __other) noexcept {
        std::swap(__coro_, __other.__coro_);
        std::swap(__started_, __other.__started_);
    }

    struct sentinel {};

    class iterator {
      public:
        using iterator_category = std::input_iterator_tag;
        using difference_type = std::ptrdiff_t;
        using value_type = _Value;
        using reference = _Ref;
        using pointer = std::add_pointer_t<_Ref>;

        iterator() noexcept = default;
        iterator(const iterator &) = delete;

        iterator(iterator&& __other) noexcept
        : __coro_(std::exchange(__other.__coro_, {})) {
        }

        iterator& operator=(iterator&& __other) {
            std::swap(__coro_, __other.__coro_);
            return *this;
        }

        ~iterator() {
        }

        friend bool operator==(const iterator &it, sentinel) noexcept {
            return it.__coro_.done();
        }

        iterator &operator++() {
            __coro_.promise().__value_.destruct();
            __coro_.promise().resume();
            return *this;
        }
        void operator++(int) {
            (void)operator++();
        }

        reference operator*() const noexcept {
            return static_cast<reference>(__coro_.promise().__value_.get());
        }

      private:
        friend generator;

        explicit iterator(__coroutine_handle __coro) noexcept
        : __coro_(__coro) {}

        __coroutine_handle __coro_;
    };

    iterator begin() {
        assert(__coro_);
        assert(!__started_);
        __started_ = true;
        __coro_.resume();
        return iterator{__coro_};
    }

    sentinel end() noexcept {
        return {};
    }

private:
    explicit generator(__coroutine_handle __coro) noexcept
        : __coro_(__coro) {
    }

public: // to get around access restrictions for __yield_sequence_awaitable
    std::coroutine_handle<> __get_coro() noexcept { return __coro_; }
    promise_type* __get_promise() noexcept { return std::addressof(__coro_.promise()); }

private:
    __coroutine_handle __coro_;
    bool __started_ = false;
};

// Specialisation for type-erased allocator implementation.
export template <typename _Ref, typename _Value>
class generator<_Ref, _Value, use_allocator_arg> {
    using __promise_base = __generator_promise_base<_Ref>;
public:

    generator() noexcept 
        : __promise_(nullptr)
        , __coro_()
        , __started_(false)
    {}

    generator(generator&& __other) noexcept
        : __promise_(std::exchange(__other.__promise_, nullptr))
        , __coro_(std::exchange(__other.__coro_, {}))
        , __started_(std::exchange(__other.__started_, false)) {
    }

    ~generator() noexcept {
        if (__coro_) {
            if (__started_ && !__coro_.done()) {
                __promise_->__value_.destruct();
            }
            __coro_.destroy();
        }
    }

    generator& operator=(generator g) noexcept {
        swap(g);
        return *this;
    }

    void swap(generator& __other) noexcept {
        std::swap(__promise_, __other.__promise_);
        std::swap(__coro_, __other.__coro_);
        std::swap(__started_, __other.__started_);
    }

    struct sentinel {};

    class iterator {
      public:
        using iterator_category = std::input_iterator_tag;
        using difference_type = std::ptrdiff_t;
        using value_type = _Value;
        using reference = _Ref;
        using pointer = std::add_pointer_t<_Ref>;

        iterator() noexcept = default;
        iterator(const iterator &) = delete;

        iterator(iterator&& __other) noexcept
        : __promise_(std::exchange(__other.__promise_, nullptr))
        , __coro_(std::exchange(__other.__coro_, {}))
        {}

        iterator& operator=(iterator&& __other) {
            __promise_ = std::exchange(__other.__promise_, nullptr);
            __coro_ = std::exchange(__other.__coro_, {});
            return *this;
        }

        ~iterator() = default;

        friend bool operator==(const iterator &it, sentinel) noexcept {
            return it.__coro_.done();
        }

        iterator& operator++() {
            __promise_->__value_.destruct();
            __promise_->resume();
            return *this;
        }

        void operator++(int) {
            (void)operator++();
        }

        reference operator*() const noexcept {
            return static_cast<reference>(__promise_->__value_.get());
        }

      private:
        friend generator;

        explicit iterator(__promise_base* __promise, std::coroutine_handle<> __coro) noexcept
        : __promise_(__promise)
        , __coro_(__coro)
        {}

        __promise_base* __promise_;
        std::coroutine_handle<> __coro_;
    };

    iterator begin() {
        assert(__coro_);
        assert(!__started_);
        __started_ = true;
        __coro_.resume();
        return iterator{__promise_, __coro_};
    }

    sentinel end() noexcept {
        return {};
    }

private:
    template<typename _Generator, typename _ByteAllocator, bool _ExplicitAllocator>
    friend struct __generator_promise;

    template<typename _Promise>
    explicit generator(std::coroutine_handle<_Promise> __coro) noexcept
        : __promise_(std::addressof(__coro.promise()))
        , __coro_(__coro)
    {}

public: // to get around access restrictions for __yield_sequence_awaitable
    std::coroutine_handle<> __get_coro() noexcept { return __coro_; }
    __promise_base* __get_promise() noexcept { return __promise_; }

private:
    __promise_base* __promise_;
    std::coroutine_handle<> __coro_;
    bool __started_ = false;
};

} // namespace nihil