#pragma once

#include <coroutine>

namespace aisa {

    inline auto suspend() noexcept { return std::suspend_always{}; }

    template<typename result_t> struct promise;

    template<typename result_t> struct task : public std::coroutine_handle<promise<result_t>> {
        using handle = std::coroutine_handle<struct promise<result_t>>;
        using promise_type = struct promise<result_t>;
        bool await_ready() const noexcept { return handle::done(); }
        result_t await_resume() const noexcept;
        template<typename other_t> void await_suspend(std::coroutine_handle<struct promise<other_t>> h) const noexcept;
        std::optional<result_t> operator()() noexcept;
    };

    template<> struct task<void> : public std::coroutine_handle<promise<void>> {
        using handle = std::coroutine_handle<struct promise<void>>;
        using promise_type = struct promise<void>;
        bool await_ready() const noexcept { return handle::done(); }
        void await_resume() const noexcept;
        template<typename other_t> void await_suspend(std::coroutine_handle<struct promise<other_t>> h) const noexcept;
        bool operator()() noexcept;
    };

    template<typename result_t> struct promise {
        std::coroutine_handle<> precursor;
        std::optional<result_t> result;
        promise() = default;
        promise(const promise &) = delete;
        task<result_t> get_return_object() noexcept { return task<result_t>{std::coroutine_handle<promise<result_t>>::from_promise(*this)}; }
        std::suspend_never initial_suspend() const noexcept { return {}; }
        std::suspend_always final_suspend() const noexcept { return {}; }
        void unhandled_exception() { }
        void return_value(result_t x) noexcept { result = std::move(x); }
    };

    template<> struct promise<void> {
        std::coroutine_handle<> precursor;
        promise() = default;
        promise(const promise &) = delete;
        task<void> get_return_object() noexcept { return task<void>{std::coroutine_handle<promise<void>>::from_promise(*this)}; }
        std::suspend_never initial_suspend() const noexcept { return {}; }
        std::suspend_always final_suspend() const noexcept { return {}; }
        void unhandled_exception() { }
        void return_void() noexcept { }
    };

    template<typename result_t> result_t task<result_t>::await_resume() const noexcept
    {
        auto x = std::move(handle::promise().result.value());
        handle::destroy();
        return std::move(x);
    }

    template<typename result_t> template<typename other_t> void task<result_t>::await_suspend(std::coroutine_handle<struct promise<other_t>> h) const noexcept
    {
        h.promise().precursor = *this;
    }

    template<typename result_t> std::optional<result_t> task<result_t>::operator()() noexcept
    {
        if (!handle::operator bool())
            return {};
        if (!handle::done()) {
            auto &precursor = handle::promise().precursor;
            if (precursor) {
                if (!precursor.done())
                    precursor.resume();
                if (precursor.done())
                    precursor = nullptr;
            }
            if (!precursor)
                handle::resume();
        }
        if (handle::done()) {
            auto x = await_resume();
            handle::operator=(nullptr);
            return std::move(x);
        }
        return {};
    }

    inline void task<void>::await_resume() const noexcept
    {
        handle::destroy();
    }

    template<typename other_t> void task<void>::await_suspend(std::coroutine_handle<struct promise<other_t>> h) const noexcept
    {
        h.promise().precursor = *this;
    }

    inline bool task<void>::operator()() noexcept
    {
        if (!handle::operator bool())
            return true;
        if (!handle::done()) {
            auto &precursor = handle::promise().precursor;
            if (precursor) {
                if (!precursor.done())
                    precursor.resume();
                if (precursor.done())
                    precursor = nullptr;
            }
            if (!precursor)
                handle::resume();
        }
        if (handle::done()) {
            await_resume();
            handle::operator=(nullptr);
            return true;
        }
        return false;
    }

}