/*
 * Copyright 2025 Bloomberg Finance LP
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.
 */

#ifndef INCLUDED_BUILDBOXCOMMON_FUTUREGROUP_H
#define INCLUDED_BUILDBOXCOMMON_FUTUREGROUP_H

#include <atomic>
#include <future>
#include <stdexcept>
#include <ThreadPool.h>
#include <tuple>
#include <type_traits>
#include <utility>

#include <buildboxcommon_exception.h>

/*
 * RAII class to manage a group of tasks run on a thread pool.
 *
 * `add()` schedules a task and returns a future for its result. The group's
 * destructor blocks until every scheduled task has finished, so tasks may
 * safely reference objects that outlive the group but are destroyed after it.
 */
template <typename R> class FutureGroup {
  public:
    explicit FutureGroup(ThreadPool *threadPool) : d_threadPool(threadPool) {}

    // Delete all other constructors and assignment operators
    FutureGroup() = delete;
    FutureGroup(const FutureGroup &) = delete;
    FutureGroup &operator=(const FutureGroup &) = delete;
    FutureGroup(FutureGroup &&) = delete;
    FutureGroup &operator=(FutureGroup &&) = delete;

    ~FutureGroup()
    {
        d_stopped.store(true, std::memory_order_relaxed);
        // Wait (lock-free) until every enqueued task has finished.
        for (size_t n = d_pending.load(std::memory_order_acquire); n != 0;
             n = d_pending.load(std::memory_order_acquire)) {
            d_pending.wait(n, std::memory_order_acquire);
        }
    }

    template <typename F, typename... Args>
    std::future<R> add(F &&func, Args &&...args)
    {
        using Ret = std::invoke_result_t<F, Args...>;
        static_assert(std::is_same_v<Ret, R>,
                      "All tasks must return the same type R");

        auto task =
            [func = std::forward<F>(func),
             args = std::make_tuple(std::forward<Args>(args)...)]() mutable {
                return std::apply(std::move(func), std::move(args));
            };

        if (d_threadPool == nullptr) {
            // No thread pool: run the task inline and capture its result (or
            // exception) into a ready future.
            std::promise<R> promise;
            try {
                if constexpr (std::is_void_v<R>) {
                    task();
                    promise.set_value();
                }
                else {
                    promise.set_value(task());
                }
            }
            catch (...) {
                promise.set_exception(std::current_exception());
            }
            return promise.get_future();
        }

        // The guard bumps the pending count now (before the task is enqueued)
        // and drops it when the task object is destroyed, whether the task ran
        // to completion, threw, or was discarded because enqueue() threw. This
        // keeps the count balanced without a try/catch around enqueue().
        auto wrapped = [this, guard = CompletionGuard(this),
                        task = std::move(task)]() mutable -> R {
            if (d_stopped.load(std::memory_order_relaxed)) {
                throw std::runtime_error("FutureGroup has been stopped");
            }
            return task();
        };
        return d_threadPool->enqueue(std::move(wrapped));
    }

  private:
    class CompletionGuard {
      public:
        explicit CompletionGuard(FutureGroup *group) : d_group(group)
        {
            d_group->d_pending.fetch_add(1, std::memory_order_relaxed);
        }
        CompletionGuard(CompletionGuard &&other) noexcept
            : d_group(other.d_group)
        {
            other.d_group = nullptr;
        }
        CompletionGuard(const CompletionGuard &) = delete;
        CompletionGuard &operator=(const CompletionGuard &) = delete;
        CompletionGuard &operator=(CompletionGuard &&) = delete;
        ~CompletionGuard()
        {
            if (d_group != nullptr) {
                d_group->taskFinished();
            }
        }

      private:
        FutureGroup *d_group;
    };

    void taskFinished()
    {
        if (d_pending.fetch_sub(1, std::memory_order_release) == 1) {
            d_pending.notify_one();
        }
    }

    ThreadPool *d_threadPool;
    std::atomic_bool d_stopped = false;
    std::atomic<size_t> d_pending = 0;
};

#endif
