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

export module nihil.generator:generator;

import nihil.std;
import :byte_allocator;
import :coroutine_traits;
import :elements_of;
import :forward;
import :generator_promise_base;
import :generator_promise;
import :manual_lifetime;
import :promise_base_alloc;
import :util;

namespace nihil {

// TODO :  make layout compatible promise casts possible
export template <typename Ref, typename Value, typename Alloc>
struct 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
		: m_coro(std::exchange(other.m_coro, {}))
		, m_started(std::exchange(other.m_started, false))
	{
	}

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

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

	void swap(generator &other) noexcept
	{
		std::swap(m_coro, other.m_coro);
		std::swap(m_started, other.m_started);
	}

	struct sentinel
	{
	};

	struct iterator
	{
		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
			: m_coro(std::exchange(other.m_coro, {}))
		{
		}

		auto operator=(iterator &&other) noexcept -> iterator &
		{
			std::swap(m_coro, other.m_coro);
			return *this;
		}

		~iterator() = default;

		friend auto operator==(const iterator &it, sentinel) noexcept -> bool
		{
			return it.m_coro.done();
		}

		auto operator++() -> iterator &
		{
			m_coro.promise().m_value.destruct();
			m_coro.promise().resume();
			return *this;
		}

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

		auto operator*() const noexcept -> reference
		{
			return static_cast<reference>(m_coro.promise().m_value.get());
		}

	private:
		friend generator;

		explicit iterator(coroutine_handle coro) noexcept
			: m_coro(coro)
		{
		}

		coroutine_handle m_coro;
	};

	auto begin() -> iterator
	{
		assert(m_coro);
		assert(!m_started);
		m_started = true;
		m_coro.resume();
		return iterator{m_coro};
	}

	auto end() noexcept -> sentinel
	{
		return {};
	}

private:
	explicit generator(coroutine_handle coro) noexcept
		: m_coro(coro)
	{
	}

public: // to get around access restrictions for __yield_sequence_awaitable
	auto get_coro() noexcept -> std::coroutine_handle<>
	{
		return m_coro;
	}
	auto get_promise() noexcept -> promise_type *
	{
		return std::addressof(m_coro.promise());
	}

private:
	coroutine_handle m_coro;
	bool m_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
		: m_promise(nullptr)
	{
	}

	generator(generator &&other) noexcept
		: m_promise(std::exchange(other.m_promise, nullptr))
		, m_coro(std::exchange(other.m_coro, {}))
		, m_started(std::exchange(other.m_started, false))
	{
	}

	~generator() noexcept
	{
		if (m_coro) {
			if (m_started && !m_coro.done()) {
				m_promise->m_value.destruct();
			}
			m_coro.destroy();
		}
	}

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

	void swap(generator &other) noexcept
	{
		std::swap(m_promise, other.m_promise);
		std::swap(m_coro, other.m_coro);
		std::swap(m_started, other.m_started);
	}

	struct sentinel
	{
	};

	struct iterator
	{
		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
			: m_promise(std::exchange(other.m_promise, nullptr))
			, m_coro(std::exchange(other.m_coro, {}))
		{
		}

		auto operator=(iterator &&other) noexcept -> iterator &
		{
			m_promise = std::exchange(other.m_promise, nullptr);
			m_coro = std::exchange(other.m_coro, {});
			return *this;
		}

		~iterator() = default;

		friend auto operator==(const iterator &it, sentinel) noexcept -> bool
		{
			return it.m_coro.done();
		}

		auto operator++() -> iterator &
		{
			m_promise->m_value.destruct();
			m_promise->resume();
			return *this;
		}

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

		auto operator*() const noexcept -> reference
		{
			return static_cast<reference>(m_promise->m_value.get());
		}

	private:
		friend generator;

		explicit iterator(promise_base *promise,
		                  std::coroutine_handle<> coro) noexcept
			: m_promise(promise)
			, m_coro(coro)
		{
		}

		promise_base *m_promise;
		std::coroutine_handle<> m_coro;
	};

	auto begin() -> iterator
	{
		assert(m_coro);
		assert(!m_started);
		m_started = true;
		m_coro.resume();
		return iterator{m_promise, m_coro};
	}

	auto end() noexcept -> sentinel
	{
		return {};
	}

private:
	template <typename Generator, typename ByteAllocator, bool ExplicitAllocator>
	friend struct generator_promise;

	template <typename Promise>
	explicit generator(std::coroutine_handle<Promise> coro) noexcept
		: m_promise(std::addressof(coro.promise()))
		, m_coro(coro)
	{
	}

public: // to get around access restrictions for __yield_sequence_awaitable
	auto get_coro() noexcept -> std::coroutine_handle<>
	{
		return m_coro;
	}

	auto get_promise() noexcept -> promise_base *
	{
		return m_promise;
	}

private:
	promise_base *m_promise;
	std::coroutine_handle<> m_coro;
	bool m_started = false;
};

} // namespace nihil