//  Copyright (c) 2007-2021 Hartmut Kaiser
//
//  SPDX-License-Identifier: BSL-1.0
//  Distributed under the Boost Software License, Version 1.0. (See accompanying
//  file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#pragma once

#include <hpx/config.hpp>
#include <hpx/assert.hpp>
#include <hpx/futures/future.hpp>
#include <hpx/futures/futures_factory.hpp>
#include <hpx/futures/traits/acquire_shared_state.hpp>
#include <hpx/futures/traits/future_access.hpp>
#include <hpx/futures/traits/future_traits.hpp>

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>

///////////////////////////////////////////////////////////////////////////////
namespace hpx { namespace lcos {

    namespace detail {
        template <typename Future>
        struct wait_acquire_future
        {
            template <typename R>
            HPX_FORCEINLINE hpx::future<R> operator()(
                hpx::future<R>& future) const
            {
                return HPX_MOVE(future);
            }

            template <typename R>
            HPX_FORCEINLINE hpx::shared_future<R> operator()(
                hpx::shared_future<R>& future) const
            {
                return future;
            }
        };

        ///////////////////////////////////////////////////////////////////////
        // This version has a callback to be invoked for each future when it
        // gets ready.
        template <typename Future, typename F>
        struct wait_each
        {
        protected:
            void on_future_ready_(hpx::execution_base::agent_ref ctx)
            {
                std::size_t oldcount = ready_count_.fetch_add(1);
                HPX_ASSERT(oldcount < lazy_values_.size());

                if (oldcount + 1 == lazy_values_.size())
                {
                    // reactivate waiting thread only if it's not us
                    if (ctx != hpx::execution_base::this_thread::agent())
                        ctx.resume();
                    else
                        goal_reached_on_calling_thread_ = true;
                }
            }

        public:
            using argument_type = std::vector<Future>;

            template <typename F_>
            wait_each(argument_type const& lazy_values, F_&& f,
                std::atomic<std::size_t>* success_counter)
              : lazy_values_(lazy_values)
              , ready_count_(0)
              , f_(HPX_FORWARD(F, f))
              , success_counter_(success_counter)
              , goal_reached_on_calling_thread_(false)
            {
            }

            template <typename F_>
            wait_each(argument_type&& lazy_values, F_&& f,
                std::atomic<std::size_t>* success_counter)
              : lazy_values_(HPX_MOVE(lazy_values))
              , ready_count_(0)
              , f_(HPX_FORWARD(F, f))
              , success_counter_(success_counter)
              , goal_reached_on_calling_thread_(false)
            {
            }

            wait_each(wait_each&& rhs)
              : lazy_values_(HPX_MOVE(rhs.lazy_values_))
              , ready_count_(rhs.ready_count_.load())
              , f_(HPX_MOVE(rhs.f_))
              , success_counter_(rhs.success_counter_)
              , goal_reached_on_calling_thread_(
                    rhs.goal_reached_on_calling_thread_)
            {
                rhs.success_counter_ = nullptr;
                rhs.goal_reached_on_calling_thread_ = false;
            }

            wait_each& operator=(wait_each&& rhs)
            {
                if (this != &rhs)
                {
                    lazy_values_ = HPX_MOVE(rhs.lazy_values_);
                    ready_count_.store(rhs.ready_count_.load());
                    rhs.ready_count_ = 0;
                    f_ = HPX_MOVE(rhs.f_);
                    success_counter_ = rhs.success_counter_;
                    rhs.success_counter_ = nullptr;
                    goal_reached_on_calling_thread_ =
                        rhs.goal_reached_on_calling_thread_;
                    rhs.goal_reached_on_calling_thread_ = false;
                }
                return *this;
            }

            std::vector<Future> operator()()
            {
                ready_count_.store(0);
                goal_reached_on_calling_thread_ = false;

                // set callback functions to executed when future is ready
                std::size_t size = lazy_values_.size();
                auto ctx = hpx::execution_base::this_thread::agent();
                for (std::size_t i = 0; i != size; ++i)
                {
                    typedef typename traits::detail::shared_state_ptr_for<
                        Future>::type shared_state_ptr;
                    shared_state_ptr current =
                        traits::detail::get_shared_state(lazy_values_[i]);

                    current->execute_deferred();
                    current->set_on_completed(
                        [HPX_CXX20_CAPTURE_THIS(=)]() -> void {
                            constexpr bool is_void =
                                std::is_void_v<traits::future_traits_t<Future>>;

                            if (lazy_values_[i].has_value())
                            {
                                if (success_counter_)
                                    ++*success_counter_;

                                // invoke callback function
                                if constexpr (is_void)
                                {
                                    f_(i);
                                }
                                else
                                {
                                    f_(i, lazy_values_[i].get());
                                }
                            }

                            // keep track of ready futures
                            on_future_ready_(ctx);
                        });
                }

                // If all of the requested futures are already set then our
                // callback above has already been called, otherwise we suspend
                // ourselves.
                if (!goal_reached_on_calling_thread_)
                {
                    // wait for all of the futures to return to become ready
                    hpx::execution_base::this_thread::suspend(
                        "hpx::lcos::detail::wait_each::operator()");
                }

                // all futures should be ready
                HPX_ASSERT(ready_count_ == size);

                return HPX_MOVE(lazy_values_);
            }

            std::vector<Future> lazy_values_;
            std::atomic<std::size_t> ready_count_;
            typename std::remove_reference<F>::type f_;
            std::atomic<std::size_t>* success_counter_;
            bool goal_reached_on_calling_thread_;
        };
    }    // namespace detail

    ///////////////////////////////////////////////////////////////////////////
    // Asynchronous versions.

    /// The one argument version is special in the sense that it returns the
    /// expected value directly (without wrapping it into a tuple).
    template <typename Future, typename F>
    HPX_DEPRECATED_V(1, 8,
        "hpx::lcos::wait is deprecated and will be removed. Use hpx::wait_each "
        "instead.")
    std::enable_if_t<!std::is_void_v<traits::future_traits_t<Future>>,
        std::size_t> wait(Future&& f1, F&& f)
    {
        f(0, f1.get());
        return 1;
    }

    template <typename Future, typename F>
    HPX_DEPRECATED_V(1, 8,
        "hpx::lcos::wait is deprecated and will be removed. Use hpx::wait_each "
        "instead.")
    std::enable_if_t<std::is_void_v<traits::future_traits_t<Future>>,
        std::size_t> wait(Future&& f1, F&& f)
    {
        f1.get();
        f(0);
        return 1;
    }

    //////////////////////////////////////////////////////////////////////////
    // This overload of wait() will make sure that the passed function will be
    // invoked as soon as a value becomes available, it will not wait for all
    // results to be there.
    template <typename Future, typename F>
    HPX_DEPRECATED_V(1, 8,
        "hpx::lcos::wait is deprecated and will be removed. Use hpx::wait_each "
        "instead.")
    inline std::size_t wait(std::vector<Future>& lazy_values, F&& f,
        std::int32_t /* suspend_for */ = 10)
    {
        typedef std::vector<Future> return_type;

        if (lazy_values.empty())
            return 0;

        return_type lazy_values_;
        lazy_values_.reserve(lazy_values.size());
        std::transform(lazy_values.begin(), lazy_values.end(),
            std::back_inserter(lazy_values_),
            detail::wait_acquire_future<Future>());

        std::atomic<std::size_t> success_counter(0);
        lcos::local::futures_factory<return_type()> p(
            detail::wait_each<Future, F>(
                HPX_MOVE(lazy_values_), HPX_FORWARD(F, f), &success_counter));

        p.post();
        p.get_future().get();

        return success_counter.load();
    }

    template <typename Future, typename F>
    HPX_DEPRECATED_V(1, 8,
        "hpx::lcos::wait is deprecated and will be removed. Use hpx::wait_each "
        "instead.")
    inline std::size_t wait(
        std::vector<Future>&& lazy_values, F&& f, std::int32_t suspend_for = 10)
    {
        return wait(lazy_values, HPX_FORWARD(F, f), suspend_for);
    }

    template <typename Future, typename F>
    HPX_DEPRECATED_V(1, 8,
        "hpx::lcos::wait is deprecated and will be removed. Use hpx::wait_each "
        "instead.")
    inline std::size_t wait(std::vector<Future> const& lazy_values, F&& f,
        std::int32_t /* suspend_for */ = 10)
    {
        typedef std::vector<Future> return_type;

        if (lazy_values.empty())
            return 0;

        return_type lazy_values_;
        lazy_values_.reserve(lazy_values.size());
        std::transform(lazy_values.begin(), lazy_values.end(),
            std::back_inserter(lazy_values_),
            detail::wait_acquire_future<Future>());

        std::atomic<std::size_t> success_counter(0);
        lcos::local::futures_factory<return_type()> p(
            detail::wait_each<Future, F>(
                HPX_MOVE(lazy_values_), HPX_FORWARD(F, f), &success_counter));

        p.post();
        p.get_future().get();

        return success_counter.load();
    }
}}    // namespace hpx::lcos
