/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

/**
 * @file task_set.h
 * A file providing TaskSet and ConcurrentTaskSet.  These interfaces allow the user to
 * submit/schedule multiple closures and then wait on them.
 **/

#pragma once

#include <dispenso/thread_pool.h>

namespace dispenso {

class TaskSetBase;

namespace detail {
template <typename Result>
class FutureBase;

class LimitGatedScheduler;

DISPENSO_DLL_ACCESS void pushThreadTaskSet(TaskSetBase* tasks);
DISPENSO_DLL_ACCESS void popThreadTaskSet();

} // namespace detail

DISPENSO_DLL_ACCESS TaskSetBase* parentTaskSet();

class TaskSetBase {
 public:
  TaskSetBase(
      ThreadPool& p,
      ParentCascadeCancel registerForParentCancel = ParentCascadeCancel::kOff,
      ssize_t stealingLoadMultiplier = 4)
      : pool_(p), taskSetLoadFactor_(stealingLoadMultiplier * p.numThreads()) {
#if defined DISPENSO_DEBUG
    assert(stealingLoadMultiplier > 0);
    pool_.outstandingTaskSets_.fetch_add(1, std::memory_order_acquire);
#endif

    parent_ = (registerForParentCancel == ParentCascadeCancel::kOn) ? parentTaskSet() : nullptr;

    if (parent_) {
      parent_->registerChild(this);
      if (parent_->canceled()) {
        canceled_.store(true, std::memory_order_release);
      }
    }
  }

  TaskSetBase(TaskSetBase&& other) = delete;
  TaskSetBase& operator=(TaskSetBase&& other) = delete;

  ssize_t numPoolThreads() const {
    return pool_.numThreads();
  }

  ThreadPool& pool() {
    return pool_;
  }

  void cancel() {
    canceled_.store(true, std::memory_order_release);
    cancelChildren();
  }

  bool canceled() const {
    return canceled_.load(std::memory_order_acquire);
  }

  /**
   * Check whether an exception has been captured by this task set.
   * When exceptions are disabled at compile time, this always returns false,
   * allowing the compiler to eliminate exception-related branches entirely.
   **/
#if defined(__cpp_exceptions)
  bool hasException() const {
    return guardException_.load(std::memory_order_acquire) != kUnset;
  }
#else
  constexpr bool hasException() const {
    return false;
  }
#endif

  ~TaskSetBase() {
#if defined DISPENSO_DEBUG
    pool_.outstandingTaskSets_.fetch_sub(1, std::memory_order_release);
#endif

    if (parent_) {
      parent_->unregisterChild(this);
    }
  }

 protected:
  template <typename F>
  auto packageTask(F&& f) {
    outstandingTaskCount_.fetch_add(1, std::memory_order_acquire);
    return [this, f = std::move(f)]() mutable {
      // Skip push/pop if this TaskSet is already the current parent on this
      // thread. This happens with ConcurrentTaskSet self-recursion (scheduling
      // tasks from within tasks on the same set), which is normal and expected.
      // Only actual nesting of *different* TaskSets needs stack tracking.
      bool pushed = (parentTaskSet() != this);
      if (pushed) {
        detail::pushThreadTaskSet(this);
      }
      if (!canceled_.load(std::memory_order_acquire)) {
#if defined(__cpp_exceptions)
        try {
          f();
        } catch (...) {
          trySetCurrentException();
        }
#else
        f();
#endif // __cpp_exceptions
      }
      if (pushed) {
        detail::popThreadTaskSet();
      }
      outstandingTaskCount_.fetch_sub(1, std::memory_order_release);
    };
  }

  // Package a task without incrementing the counter (for bulk scheduling where
  // the counter is incremented once for the entire batch).
  template <typename F>
  auto packageTaskNoIncrement(F&& f) {
    return [this, f = std::move(f)]() mutable {
      detail::pushThreadTaskSet(this);
      if (!canceled_.load(std::memory_order_acquire)) {
#if defined(__cpp_exceptions)
        try {
          f();
        } catch (...) {
          trySetCurrentException();
        }
#else
        f();
#endif // __cpp_exceptions
      }
      detail::popThreadTaskSet();
      outstandingTaskCount_.fetch_sub(1, std::memory_order_release);
    };
  }

  template <typename Generator>
  void scheduleBulkImpl(size_t count, Generator&& gen, moodycamel::ProducerToken* token) {
    if (count == 0) {
      return;
    }

    ssize_t numPool = pool_.numThreads();
    size_t chunkSize = static_cast<size_t>(numPool) + static_cast<size_t>(numPool) / 2;
    if (chunkSize < 1) {
      chunkSize = 1;
    }

    size_t i = 0;
    while (i < count) {
      if (canceled()) {
        break;
      }
      ssize_t outstanding = outstandingTaskCount_.load(std::memory_order_relaxed);
      ssize_t curWork = pool_.workRemaining_.load(std::memory_order_relaxed);
      // Mirror the two-tier inline decision from ThreadPool::schedule():
      // 1. TaskSet-level: outstandingTaskCount_ exceeds our own load factor
      // 2. Pool recursive: tight quickLoadFactor (numThreads*1.5) when called from a pool
      //    thread. Critical for nested patterns where inner scheduleBulk must throttle to
      //    avoid flooding the pool queue and starving outer tasks.
      // 3. Pool global: loose poolLoadFactor_ (numThreads*32) for non-recursive callers
      if (outstanding > taskSetLoadFactor_ ||
          (detail::PerPoolPerThreadInfo::isPoolRecursive(&pool_) &&
           curWork > numPool + numPool / 2) ||
          curWork > pool_.poolLoadFactor_.load(std::memory_order_relaxed)) {
#if defined(__cpp_exceptions)
        try {
          gen(i)();
        } catch (...) {
          trySetCurrentException();
        }
#else
        gen(i)();
#endif // __cpp_exceptions
        ++i;
      } else {
        ssize_t room = taskSetLoadFactor_ - outstanding;
        size_t toEnqueue = std::min({count - i, chunkSize, static_cast<size_t>(room)});
        if (toEnqueue == 0) {
          toEnqueue = 1;
        }
        outstandingTaskCount_.fetch_add(static_cast<ssize_t>(toEnqueue), std::memory_order_acquire);
        size_t base = i;
        pool_.scheduleBulkEnqueue(
            toEnqueue,
            [this, &gen, base](size_t j) { return packageTaskNoIncrement(gen(base + j)); },
            token);
        i += toEnqueue;
      }
    }
  }

  DISPENSO_DLL_ACCESS void trySetCurrentException();
  bool testAndResetException();

  void registerChild(TaskSetBase* child) {
    std::lock_guard<std::mutex> lk(mtx_);

    child->prev_ = tail_;
    child->next_ = nullptr;
    if (tail_) {
      tail_->next_ = child;
      tail_ = child;
    } else {
      head_ = tail_ = child;
    }
  }

  void unregisterChild(TaskSetBase* child) {
    std::lock_guard<std::mutex> lk(mtx_);

    if (child->prev_) {
      child->prev_->next_ = child->next_;
    } else {
      // We're head
      assert(child == head_);
      head_ = child->next_;
    }
    if (child->next_) {
      child->next_->prev_ = child->prev_;
    } else {
      // We're tail
      assert(child == tail_);
      tail_ = child->prev_;
    }
  }

  void cancelChildren() {
    std::lock_guard<std::mutex> lk(mtx_);

    auto* node = head_;
    while (node) {
      node->cancel();
      node = node->next_;
    }
  }

  alignas(kCacheLineSize) std::atomic<ssize_t> outstandingTaskCount_{0};
  alignas(kCacheLineSize) ThreadPool& pool_;
  alignas(kCacheLineSize) std::atomic<bool> canceled_{false};
  const ssize_t taskSetLoadFactor_;
  // Always present to ensure stable ABI layout regardless of __cpp_exceptions.
  enum ExceptionState { kUnset, kSetting, kSet };
  std::atomic<ExceptionState> guardException_{kUnset};
  std::exception_ptr exception_;

  TaskSetBase* parent_;

  // This mutex guards modifications/use of the intusive linked list between head_ and tail_
  std::mutex mtx_;
  TaskSetBase* head_{nullptr};
  TaskSetBase* tail_{nullptr};

  // prev_ and next_ are links in our *parent's* intrusive linked list.
  TaskSetBase* prev_;
  TaskSetBase* next_;
};

} // namespace dispenso
