/*
 * Copyright (c) 2021-2024 NVIDIA Corporation
 *
 * Licensed under the Apache License Version 2.0 with LLVM Exceptions
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *   https://llvm.org/LICENSE.txt
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include <any>
#include <cassert>
#include <exception>
#include <utility>

#include "../stdexec/__detail/__meta.hpp"
#include "../stdexec/__detail/__optional.hpp"
#include "../stdexec/__detail/__variant.hpp"
#include "../stdexec/coroutine.hpp"
#include "../stdexec/execution.hpp"
#include "../stdexec/functional.hpp"

#include "any_sender_of.hpp"
#include "at_coroutine_exit.hpp"
#include "completion_behavior.hpp"
#include "scope.hpp"

STDEXEC_PRAGMA_PUSH()
STDEXEC_PRAGMA_IGNORE_GNU("-Wundefined-inline")

namespace experimental::execution
{
  namespace __task
  {
    using namespace STDEXEC;

    // The required set_value_t() scheduler-sender completion signature is added in
    // any_receiver_ref::any_sender::any_scheduler.
    using __any_scheduler_completions_t =
      completion_signatures<set_value_t(), set_error_t(std::exception_ptr), set_stopped_t()>;

    using __any_scheduler_impl_t =
      any_scheduler<any_sender<any_receiver<__any_scheduler_completions_t>>>;

    // A scheduler concept that does not check for copyability since that creates a cycle
    // in the type system.
    template <class _Scheduler>
    concept __semi_scheduler = requires(_Scheduler& __sched) {
      typename _Scheduler::scheduler_concept;
      requires __std::derived_from<typename _Scheduler::scheduler_concept, scheduler_tag>;
      { schedule(__sched) } -> sender;
    };

    struct __any_scheduler
    {
      using scheduler_concept = scheduler_t;

      template <__not_same_as<__any_scheduler> _Scheduler>
        requires __semi_scheduler<_Scheduler>
      constexpr __any_scheduler(_Scheduler __sched) noexcept
        : __impl_(std::forward<_Scheduler>(__sched))
      {}

      bool operator==(__any_scheduler const & __other) const noexcept = default;

      [[nodiscard]]
      auto schedule() const
      {
        return __impl_.schedule();
      }

     private:
      __any_scheduler_impl_t __impl_;
    };

    static_assert(scheduler<__any_scheduler>);

    template <class _Ty>
    concept __stop_token_provider = requires(_Ty const & t) { get_stop_token(t); };

    template <class _Ty>
    concept __indirect_stop_token_provider = requires(_Ty const & t) {
      { get_env(t) } -> __stop_token_provider;
    };

    template <class _Ty>
    concept __indirect_start_scheduler_provider = requires(_Ty const & t) {
      { get_start_scheduler(get_env(t)) } -> scheduler;
    };

    template <class _ParentPromise>
    constexpr auto __parent_promise_has_start_scheduler() noexcept -> bool
    {
      static_assert(__indirect_start_scheduler_provider<_ParentPromise>,
                    "exec::task<T> cannot be co_await-ed in a coroutine that "
                    "does not have an associated start scheduler.");
      return __indirect_start_scheduler_provider<_ParentPromise>;
    }

    template <class _ParentPromise>
    struct __default_awaiter_context;

    ////////////////////////////////////////////////////////////////////////////////
    // This is the context that is associated with basic_task's promise type
    // by default. It handles forwarding of stop requests from parent to child.
    enum class __scheduler_affinity
    {
      __none,
      __sticky
    };

    template <__scheduler_affinity _SchedulerAffinity>
    struct __optional_scheduler_storage
    {
      __optional_scheduler_storage() = default;
      constexpr explicit __optional_scheduler_storage(__ignore) noexcept {}
    };

    template <>
    struct __optional_scheduler_storage<__scheduler_affinity::__sticky>
    {
      __optional_scheduler_storage() = default;

      template <scheduler _Scheduler>
      constexpr explicit __optional_scheduler_storage(_Scheduler __sched) noexcept
        : __scheduler_(__sched)
      {}

      __any_scheduler __scheduler_{STDEXEC::inline_scheduler{}};
    };

    template <__scheduler_affinity _SchedulerAffinity = __scheduler_affinity::__sticky>
    class __default_task_context_impl : private __optional_scheduler_storage<_SchedulerAffinity>
    {
      template <class _ParentPromise>
      friend struct __default_awaiter_context;

      static constexpr bool __with_affinity = _SchedulerAffinity == __scheduler_affinity::__sticky;

      inplace_stop_token __stop_token_;

     public:
      template <class _ParentPromise>
      constexpr explicit __default_task_context_impl(_ParentPromise& __parent) noexcept
      {
        if constexpr (__with_affinity)
        {
          if constexpr (__parent_promise_has_start_scheduler<_ParentPromise>())
          {
            // get_start_scheduler is used here to get the parent's "current" scheduler,
            // which is the one on which this task has been started (i.e., co_await-ed).
            auto __parent_sched = get_start_scheduler(get_env(__parent));
            this->__scheduler_  = __parent_sched;
          }
        }
      }

      template <scheduler _Scheduler>
      constexpr explicit __default_task_context_impl(_Scheduler&& __sched) noexcept
        : __optional_scheduler_storage<_SchedulerAffinity>{static_cast<_Scheduler&&>(__sched)}
      {}

      [[nodiscard]]
      constexpr auto query(get_start_scheduler_t) const noexcept -> __any_scheduler const &
        requires(__with_affinity)
      {
        return this->__scheduler_;
      }

      [[nodiscard]]
      constexpr auto query(get_stop_token_t) const noexcept -> inplace_stop_token
      {
        return __stop_token_;
      }

      [[nodiscard]]
      constexpr auto query(get_completion_behavior_t<set_value_t>) const noexcept
      {
        if constexpr (__with_affinity)
        {
          return completion_behavior::asynchronous_affine | completion_behavior::inline_completion;
        }
        else
        {
          return completion_behavior::unknown;
        }
      }

      [[nodiscard]]
      constexpr auto stop_requested() const noexcept -> bool
      {
        return __stop_token_.stop_requested();
      }

      template <scheduler _Scheduler>
      constexpr void set_scheduler(_Scheduler&& __sched)
        requires(__with_affinity)
      {
        this->__scheduler_ = static_cast<_Scheduler&&>(__sched);
      }

      template <class _ThisPromise>
      using promise_context_t = __default_task_context_impl;

      template <class _ThisPromise, class _ParentPromise = void>
        requires(!__with_affinity) || __indirect_start_scheduler_provider<_ParentPromise>
      using awaiter_context_t = __default_awaiter_context<_ParentPromise>;
    };

    template <class _Ty>
    using default_task_context = __default_task_context_impl<__scheduler_affinity::__sticky>;

    template <class _Ty>
    using inline_task_context = __default_task_context_impl<__scheduler_affinity::__none>;

    // This is the context associated with basic_task's awaiter. By default
    // it does nothing.
    template <class _ParentPromise>
    struct __default_awaiter_context
    {
      template <__scheduler_affinity _Affinity>
      constexpr explicit __default_awaiter_context(__default_task_context_impl<_Affinity>&,
                                                   _ParentPromise&) noexcept
      {}
    };

    ////////////////////////////////////////////////////////////////////////////////
    // This is the context to be associated with basic_task's awaiter when
    // the parent coroutine's promise type is known, is a __stop_token_provider,
    // and its stop token type is neither inplace_stop_token nor unstoppable.
    template <__indirect_stop_token_provider _ParentPromise>
    struct __default_awaiter_context<_ParentPromise>
    {
      using __stop_token_t    = stop_token_of_t<env_of_t<_ParentPromise>>;
      using __stop_callback_t = __stop_token_t::template callback_type<__forward_stop_request<>>;

      template <__scheduler_affinity _Affinity>
      constexpr explicit __default_awaiter_context(__default_task_context_impl<_Affinity>& __self,
                                                   _ParentPromise& __parent) noexcept
        // Register a callback that will request stop on this basic_task's
        // stop_source when stop is requested on the parent coroutine's stop
        // token.
        : __stop_callback_{get_stop_token(get_env(__parent)),
                           __forward_stop_request{__stop_source_}}
      {
        static_assert(std::is_nothrow_constructible_v<__stop_callback_t,
                                                      __stop_token_t,
                                                      __forward_stop_request<>>);
        __self.__stop_token_ = __stop_source_.get_token();
      }

      inplace_stop_source __stop_source_{};
      __stop_callback_t   __stop_callback_;
    };

    // If the parent coroutine's type has a stop token of type inplace_stop_token,
    // we don't need to register a stop callback.
    template <__indirect_stop_token_provider _ParentPromise>
      requires std::same_as<inplace_stop_token, stop_token_of_t<env_of_t<_ParentPromise>>>
    struct __default_awaiter_context<_ParentPromise>
    {
      template <__scheduler_affinity _Affinity>
      constexpr explicit __default_awaiter_context(__default_task_context_impl<_Affinity>& __self,
                                                   _ParentPromise& __parent) noexcept
      {
        __self.__stop_token_ = get_stop_token(get_env(__parent));
      }
    };

    // If the parent coroutine's stop token is unstoppable, there's no point
    // forwarding stop tokens or stop requests at all.
    template <__indirect_stop_token_provider _ParentPromise>
      requires unstoppable_token<stop_token_of_t<env_of_t<_ParentPromise>>>
    struct __default_awaiter_context<_ParentPromise>
    {
      template <__scheduler_affinity _Affinity>
      constexpr explicit __default_awaiter_context(__default_task_context_impl<_Affinity>&,
                                                   _ParentPromise&) noexcept
      {}
    };

    // Finally, if we don't know the parent coroutine's promise type, assume the
    // worst and save a type-erased stop callback.
    template <>
    struct __default_awaiter_context<void>
    {
      template <__scheduler_affinity _Affinity, class _ParentPromise>
      constexpr explicit __default_awaiter_context(__default_task_context_impl<_Affinity>&,
                                                   _ParentPromise&) noexcept
      {}

      template <__scheduler_affinity _Affinity, __indirect_stop_token_provider _ParentPromise>
      constexpr explicit __default_awaiter_context(__default_task_context_impl<_Affinity>& __self,
                                                   _ParentPromise&                         __parent)
      {
        // Register a callback that will request stop on this basic_task's
        // stop_source when stop is requested on the parent coroutine's stop
        // token.
        using __stop_token_t    = stop_token_of_t<env_of_t<_ParentPromise>>;
        using __stop_callback_t = stop_callback_for_t<__stop_token_t, __forward_stop_request<>>;

        if constexpr (std::same_as<__stop_token_t, inplace_stop_token>)
        {
          __self.__stop_token_ = get_stop_token(get_env(__parent));
        }
        else if (auto __token = get_stop_token(get_env(__parent)); __token.stop_possible())
        {
          __stop_callback_.emplace<__stop_callback_t>(std::move(__token),
                                                      __forward_stop_request{__stop_source_});
          __self.__stop_token_ = __stop_source_.get_token();
        }
      }

      inplace_stop_source __stop_source_{};
      std::any            __stop_callback_{};
    };

    template <class _Promise, class _ParentPromise = void>
    using awaiter_context_t =
      __decay_t<env_of_t<_Promise>>::template awaiter_context_t<_Promise, _ParentPromise>;

    ////////////////////////////////////////////////////////////////////////////////
    // In a base class so it can be specialized when _Ty is void:
    template <class _Ty>
    struct __promise_base
    {
      constexpr void return_value(_Ty value)
      {
        __data_.template emplace<0>(std::move(value));
      }

      __variant<_Ty, std::exception_ptr> __data_{__no_init};
    };

    template <>
    struct __promise_base<void>
    {
      struct __void
      {};

      constexpr void return_void()
      {
        __data_.template emplace<0>(__void{});
      }

      __variant<__void, std::exception_ptr> __data_{__no_init};
    };

    template <class _Sch>
    struct __just_void
    {
      using sender_concept        = sender_tag;
      using completion_signatures = STDEXEC::completion_signatures<set_value_t()>;

      template <class _Rcvr>
      [[nodiscard]]
      static constexpr auto connect(_Rcvr __rcvr) noexcept
      {
        return STDEXEC::connect(just(), static_cast<_Rcvr&&>(__rcvr));
      }

      [[nodiscard]]
      constexpr auto get_env() const noexcept
      {
        return prop{get_completion_scheduler<set_value_t>, __sch_};
      }

      _Sch __sch_;
    };

    enum class disposition : unsigned
    {
      stopped,
      succeeded,
      failed,
    };

    template <class _Promise>
    struct __reschedule_receiver
    {
      using receiver_concept = receiver_tag;

      void set_value() noexcept
      {
        // Resuming the continuation of the parent coroutine will cause it to continue
        // executing on the new scheduler.
        STDEXEC::__coroutine_resume_nothrow(__parent_);
      }

      void set_error(std::exception_ptr __eptr) noexcept
      {
        __eptr_ = std::move(__eptr);
        STDEXEC::__coroutine_resume_nothrow(__parent_);
      }

      void set_stopped() noexcept
      {
        // Resuming the stopped continuation unwinds the coroutine stack until we reach
        // a promise that can handle the stopped signal. The coroutine referred to by
        // __continuation_ will never be resumed.
        __std::coroutine_handle<> __unwind = __parent_.promise().unhandled_stopped();
        STDEXEC::__coroutine_resume_nothrow(__unwind);
      }

      [[nodiscard]]
      auto get_env() const noexcept -> env_of_t<_Promise>
      {
        return STDEXEC::get_env(__parent_.promise());
      }

      std::exception_ptr&               __eptr_;
      __std::coroutine_handle<_Promise> __parent_;
    };

    template <class _Scheduler, class _Promise>
    struct __reschedule_awaiter
    {
      using __sender_t   = __result_of<unstoppable, schedule_result_t<_Scheduler>>;
      using __receiver_t = __reschedule_receiver<_Promise>;
      using __opstate_t  = STDEXEC::connect_result_t<__sender_t, __receiver_t>;

      static constexpr auto await_ready() noexcept -> bool
      {
        return false;
      }

      auto await_suspend(__std::coroutine_handle<_Promise> __h) noexcept -> bool
      {
        STDEXEC_TRY
        {
          auto& __p = __h.promise();

          if (!std::exchange(__p.__rescheduled_, true))
          {
            // Create a cleanup action that transitions back onto the current scheduler:
            auto __sched = get_start_scheduler(*__p.__context_);
            auto __guard = at_coroutine_exit(__compose(unstoppable, STDEXEC::schedule),
                                             std::move(__sched));
            // Insert the cleanup action into the head of the continuation chain by
            // making direct calls to the cleanup task's awaiter member functions. See
            // type __at_coro_exit::__task in at_coroutine_exit.hpp:
            __guard.await_suspend(__h);
            (void) __guard.await_resume();
          }

          __p.__context_->set_scheduler(__new_sched_);
          auto& __op = __opstate_.__emplace_from(STDEXEC::connect,
                                                 unstoppable(schedule(__new_sched_)),
                                                 __receiver_t{__eptr_, __h});
          STDEXEC::start(__op);
          return true;  // suspend the coroutine until the scheduler operation completes
        }
        STDEXEC_CATCH_ALL
        {
          __eptr_ = std::current_exception();
          return false;
        }
      }

      void await_resume() noexcept
      {
        if (__eptr_)
        {
          std::rethrow_exception(std::move(__eptr_));
        }
      }

      _Scheduler              __new_sched_;
      std::exception_ptr      __eptr_{};
      __optional<__opstate_t> __opstate_{};
    };

    struct __reschedule_coroutine_on_t
    {
      template <class _Scheduler>
      struct __wrapper
      {
        _Scheduler __sched_;
      };

      template <scheduler _Scheduler>
      constexpr auto operator()(_Scheduler __sched) const noexcept -> __wrapper<_Scheduler>
      {
        return {static_cast<_Scheduler&&>(__sched)};
      }
    };

    ////////////////////////////////////////////////////////////////////////////////
    // basic_task
    template <class _Ty, class _Context = default_task_context<_Ty>>
    class [[nodiscard]] basic_task
    {
      struct __promise;

      template <class _ParentPromise>
      struct __task_awaiter;

      using __promise_context_t = _Context::template promise_context_t<__promise>;

     public:
      using promise_type = __promise;

      constexpr basic_task(basic_task&& __that) noexcept
        : __coro_(std::exchange(__that.__coro_, {}))
      {}

      // Make this task awaitable within a particular context:
      template <class _ParentPromise>
        requires __std::constructible_from<awaiter_context_t<__promise, _ParentPromise>,
                                           __promise_context_t&,
                                           _ParentPromise&>
      constexpr auto as_awaitable(_ParentPromise&) && noexcept -> __task_awaiter<_ParentPromise>
      {
        return __task_awaiter<_ParentPromise>{std::exchange(__coro_, {})};
      }

      // Make this task generally awaitable:
      constexpr auto operator co_await() && noexcept -> __task_awaiter<void>
        requires __minvocable_q<awaiter_context_t, __promise>
      {
        return __task_awaiter<void>{std::exchange(__coro_, {})};
      }

      constexpr ~basic_task()
      {
        if (__coro_)
          STDEXEC::__coroutine_destroy_nothrow(__coro_);
      }

     private:
      using __scheduler_t =
        __call_result_or_t<get_start_scheduler_t, STDEXEC::inline_scheduler, _Context>;

      struct __final_awaiter
      {
        static constexpr auto await_ready() noexcept -> bool
        {
          return false;
        }

        static constexpr auto await_suspend(__std::coroutine_handle<__promise> __h) noexcept  //
          -> __std::coroutine_handle<>
        {
          return __h.promise().continuation().handle();
        }

        static constexpr void await_resume() noexcept {}
      };

      struct __promise
        : __promise_base<_Ty>
        , with_awaitable_senders<__promise>
      {
        constexpr auto get_return_object() noexcept -> basic_task
        {
          return basic_task(__std::coroutine_handle<__promise>::from_promise(*this));
        }

        constexpr auto initial_suspend() noexcept -> __std::suspend_always
        {
          return {};
        }

        constexpr auto final_suspend() noexcept -> __final_awaiter
        {
          return {};
        }

        [[nodiscard]]
        constexpr auto disposition() const noexcept -> __task::disposition
        {
          switch (this->__data_.index())
          {
          case 0:
            return __task::disposition::succeeded;
          case 1:
            return __task::disposition::failed;
          default:
            return __task::disposition::stopped;
          }
        }

        constexpr void unhandled_exception() noexcept
        {
          this->__data_.template emplace<1>(std::current_exception());
        }

#ifndef __clang_analyzer__
        template <sender _CvSender>
          requires __start_scheduler_provider<_Context>
        auto await_transform(_CvSender&& __sndr) noexcept -> decltype(auto)
        {
          if constexpr (__completes_where_it_starts<set_value_t,
                                                    env_of_t<_CvSender>,
                                                    __promise_context_t&>)
          {
            return STDEXEC::as_awaitable(static_cast<_CvSender&&>(__sndr), *this);
          }
          else
          {
            return STDEXEC::as_awaitable(continues_on(static_cast<_CvSender&&>(__sndr),
                                                      get_start_scheduler(*__context_)),
                                         *this);
          }
        }

        template <class _Scheduler>
          requires __start_scheduler_provider<_Context>
        auto await_transform(__reschedule_coroutine_on_t::__wrapper<_Scheduler> __box) noexcept
          -> decltype(auto)
        {
          return __reschedule_awaiter<_Scheduler, __promise>{__box.__sched_};
        }
#endif

        template <__sender_adaptor_closure_for<__just_void<__scheduler_t>> _Closure>
        auto await_transform(_Closure&& __closure) noexcept -> decltype(auto)
        {
          return await_transform(static_cast<_Closure&&>(__closure)(__just_void<__scheduler_t>()));
        }

        template <class _Awaitable>
        constexpr auto await_transform(_Awaitable&& __awaitable) noexcept -> decltype(auto)
        {
          return with_awaitable_senders<__promise>::await_transform(
            static_cast<_Awaitable&&>(__awaitable));
        }

        constexpr auto get_env() const noexcept -> __promise_context_t const &
        {
          return *__context_;
        }

        __optional<__promise_context_t> __context_{};
        bool                            __rescheduled_{false};
      };

      template <class _ParentPromise>
      struct __task_awaiter
      {
        constexpr __task_awaiter(__std::coroutine_handle<__promise> __coro) noexcept
          : __coro_(__coro)
        {}
        STDEXEC_IMMOVABLE(__task_awaiter);

        constexpr ~__task_awaiter()
        {
          if (__coro_)
            STDEXEC::__coroutine_destroy_nothrow(__coro_);
        }

        static constexpr auto await_ready() noexcept -> bool
        {
          return false;
        }

        template <class _ParentPromise2>
        constexpr auto await_suspend(__std::coroutine_handle<_ParentPromise2> __parent) noexcept
          -> __std::coroutine_handle<>
        {
          static_assert(__one_of<_ParentPromise, _ParentPromise2, void>);
          __coro_.promise().__context_.emplace(__parent.promise());
          __context_.emplace(*__coro_.promise().__context_, __parent.promise());
          __coro_.promise().set_continuation(__parent);
          if constexpr (requires { __coro_.promise().stop_requested() ? 0 : 1; })
          {
            if (__coro_.promise().stop_requested())
              return STDEXEC::__coroutine_unhandled_stopped(__parent);
          }
          return __coro_;
        }

        constexpr auto await_resume() -> _Ty
        {
          __context_.reset();
          scope_guard __on_exit{
            [this]() noexcept
            { STDEXEC::__coroutine_destroy_nothrow(std::exchange(__coro_, {})); }};

          if (__coro_.promise().__data_.index() == 1)
            std::rethrow_exception(std::move(__var::__get<1>(__coro_.promise().__data_)));

          if constexpr (!std::is_void_v<_Ty>)
            return std::move(__var::__get<0>(__coro_.promise().__data_));
        }

        __std::coroutine_handle<__promise>                       __coro_{};
        __optional<awaiter_context_t<__promise, _ParentPromise>> __context_{};
      };

     public:
      constexpr explicit basic_task(__std::coroutine_handle<promise_type> __coro) noexcept
        : __coro_(__coro)
      {}

      __std::coroutine_handle<promise_type> __coro_{};
    };
  }  // namespace __task

  using task_disposition = __task::disposition;

  template <class _Ty>
  using default_task_context = __task::default_task_context<_Ty>;

  template <class _Promise, class _ParentPromise = void>
  using awaiter_context_t = __task::awaiter_context_t<_Promise, _ParentPromise>;

  template <class _Ty, class _Context = default_task_context<_Ty>>
  using basic_task = __task::basic_task<_Ty, _Context>;

  template <class _Ty>
  using task = basic_task<_Ty, default_task_context<_Ty>>;

  inline constexpr __task::__reschedule_coroutine_on_t reschedule_coroutine_on{};
}  // namespace experimental::execution

namespace exec = experimental::execution;

namespace STDEXEC
{
  template <class _Ty, class _Context>
  inline constexpr bool enable_sender<exec::basic_task<_Ty, _Context>> = true;
}  // namespace STDEXEC

STDEXEC_PRAGMA_POP()
