#pragma once
#include "detail/type.h"
#include "detail/task_promise.h"
namespace co_context {
#define CO_CONTEXT_AWAIT_HINT nodiscard("Did you forget to co_await?")
    template<typename T = void>
    class [[CO_CONTEXT_AWAIT_HINT]] task {
    public:
        using promise_type = detail::task_promise<T>;
    private:
        struct awaiter_base {
            std::coroutine_handle<promise_type> handle;

            explicit awaiter_base(std::coroutine_handle<promise_type> current
            ) noexcept
                : handle(current) {}

            [[nodiscard]]
            inline bool await_ready() const noexcept {
                return !handle || handle.done();
            }

            std::coroutine_handle<>
                await_suspend(std::coroutine_handle<> awaiting_coro) noexcept {
                handle.promise().set_parent(awaiting_coro);
                return handle;
            }
        };
    private:
        std::coroutine_handle<promise_type> handle;
    public:
        explicit task(std::coroutine_handle<promise_type> current) noexcept
            : handle(current) {}
        std::coroutine_handle<promise_type> get_handle() noexcept { return handle; }
        void detach() noexcept {
            if constexpr (std::is_void_v<T>) {
                handle.promise().is_detached_flag = promise_type::is_detached;
            }
            handle = nullptr;
        }
        /**
         * @brief wait for the task<> to complete, and get the ref of the result
         */
        auto operator co_await() const& noexcept {
            struct awaiter : awaiter_base {
                using awaiter_base::awaiter_base;

                decltype(auto) await_resume() {
                    // if (!this->handle) [[unlikely]]
                    //     throw std::logic_error("broken_promise");
                    assert(this->handle && "broken_promise");

                    return this->handle.promise().result();
                }
            };

            return awaiter{ handle };
        }

        /**
         * @brief wait for the task<> to complete, and get the rvalue ref of the
         * result
         */
        auto operator co_await() const&& noexcept {
            struct awaiter : awaiter_base {
                using awaiter_base::awaiter_base;

                decltype(auto) await_resume() {
                    // if (!this->handle) [[unlikely]]
                    //     throw std::logic_error("broken_promise");
                    assert(this->handle && "broken_promise");

                    return std::move(this->handle.promise()).result();
                }
            };

            return awaiter{ handle };
        }
    };

    namespace detail {
        template<typename T>
        inline task<T> task_promise<T>::get_return_object() noexcept {
            return task<T>{
                std::coroutine_handle<task_promise>::from_promise(*this)
            };
        }

        inline task<void> task_promise<void>::get_return_object() noexcept {
            return task<void>{
                std::coroutine_handle<task_promise>::from_promise(*this)
            };
        }

        template<typename T>
        inline task<T&> task_promise<T&>::get_return_object() noexcept {
            return task<T&>{
                std::coroutine_handle<task_promise>::from_promise(*this)
            };
        }
    } // namespace detail
} // namespace co_context