/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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.
 */

#include <exception>
#include <memory>

#include <folly/CancellationToken.h>
#include <folly/Executor.h>
#include <folly/ScopeGuard.h>
#include <folly/coro/Baton.h>
#include <folly/coro/Mutex.h>
#include <folly/coro/Task.h>
#include <folly/coro/ViaIfAsync.h>
#include <folly/coro/WithCancellation.h>
#include <folly/coro/detail/Barrier.h>
#include <folly/coro/detail/BarrierTask.h>
#include <folly/coro/detail/CurrentAsyncFrame.h>
#include <folly/coro/detail/Helpers.h>

#if FOLLY_HAS_COROUTINES

namespace folly {
namespace coro {
namespace detail {

enum class CallbackRecordSelector { Invalid, Value, None, Error };

constexpr inline std::in_place_index_t<0> const callback_record_value{};
constexpr inline std::in_place_index_t<1> const callback_record_none{};
constexpr inline std::in_place_index_t<2> const callback_record_error{};

//
// CallbackRecord records the result of a single invocation of a callback.
//
// This is very related to Try and expected, but this also records None in
// addition to Value and Error results.
//
// When the callback supports multiple overloads of Value then T would be
// something like a variant<tuple<..>, ..>
//
// When the callback supports multiple overloads of Error then all the errors
// are coerced to folly::exception_wrapper
//
template <class T>
class CallbackRecord {
  static void clear(CallbackRecord* that) {
    auto selector =
        std::exchange(that->selector_, CallbackRecordSelector::Invalid);
    if (selector == CallbackRecordSelector::Value) {
      detail::deactivate(that->value_);
    } else if (selector == CallbackRecordSelector::Error) {
      detail::deactivate(that->error_);
    }
  }
  template <class OtherReference>
  static void convert_variant(
      CallbackRecord* that, const CallbackRecord<OtherReference>& other) {
    if (other.hasValue()) {
      detail::activate(that->value_, other.value_.get());
    } else if (other.hasError()) {
      detail::activate(that->error_, other.error_.get());
    }
    that->selector_ = other.selector_;
  }
  template <class OtherReference>
  static void convert_variant(
      CallbackRecord* that, CallbackRecord<OtherReference>&& other) {
    if (other.hasValue()) {
      detail::activate(that->value_, std::move(other.value_).get());
    } else if (other.hasError()) {
      detail::activate(that->error_, std::move(other.error_).get());
    }
    that->selector_ = other.selector_;
  }

 public:
  ~CallbackRecord() { clear(this); }

  CallbackRecord() noexcept : selector_(CallbackRecordSelector::Invalid) {}

  template <class V>
  CallbackRecord(const std::in_place_index_t<0>&, V&& v) noexcept(
      std::is_nothrow_constructible_v<T, V>)
      : CallbackRecord() {
    detail::activate(value_, std::forward<V>(v));
    selector_ = CallbackRecordSelector::Value;
  }
  explicit CallbackRecord(const std::in_place_index_t<1>&) noexcept
      : selector_(CallbackRecordSelector::None) {}
  CallbackRecord(
      const std::in_place_index_t<2>&, folly::exception_wrapper e) noexcept
      : CallbackRecord() {
    detail::activate(error_, std::move(e));
    selector_ = CallbackRecordSelector::Error;
  }

  CallbackRecord(CallbackRecord&& other) noexcept(
      std::is_nothrow_move_constructible_v<T>)
      : CallbackRecord() {
    convert_variant(this, std::move(other));
  }

  CallbackRecord& operator=(CallbackRecord&& other) noexcept(
      std::is_nothrow_move_constructible_v<T>) {
    if (&other != this) {
      clear(this);
      convert_variant(this, std::move(other));
    }
    return *this;
  }

  template <class U>
  CallbackRecord(CallbackRecord<U>&& other) noexcept(
      std::is_nothrow_constructible_v<T, U>)
      : CallbackRecord() {
    convert_variant(this, std::move(other));
  }

  bool hasNone() const noexcept {
    return selector_ == CallbackRecordSelector::None;
  }

  bool hasError() const noexcept {
    return selector_ == CallbackRecordSelector::Error;
  }

  decltype(auto) error() & {
    DCHECK(hasError());
    return error_.get();
  }

  decltype(auto) error() && {
    DCHECK(hasError());
    return std::move(error_).get();
  }

  decltype(auto) error() const& {
    DCHECK(hasError());
    return error_.get();
  }

  decltype(auto) error() const&& {
    DCHECK(hasError());
    return std::move(error_).get();
  }

  bool hasValue() const noexcept {
    return selector_ == CallbackRecordSelector::Value;
  }

  decltype(auto) value() & {
    DCHECK(hasValue());
    return value_.get();
  }

  decltype(auto) value() && {
    DCHECK(hasValue());
    return std::move(value_).get();
  }

  decltype(auto) value() const& {
    DCHECK(hasValue());
    return value_.get();
  }

  decltype(auto) value() const&& {
    DCHECK(hasValue());
    return std::move(value_).get();
  }

  explicit operator bool() const noexcept {
    return selector_ != CallbackRecordSelector::Invalid;
  }

 private:
  union {
    detail::ManualLifetime<T> value_;
    detail::ManualLifetime<folly::exception_wrapper> error_;
  };
  CallbackRecordSelector selector_;
};

template <typename Reference, typename Value, typename GeneratorType>
AsyncGenerator<Reference, Value> mergeImpl(
    folly::Executor::KeepAlive<> executor, GeneratorType sources) {
  struct SharedState {
    explicit SharedState(folly::Executor::KeepAlive<> executor_)
        : executor(std::move(executor_)) {}

    const folly::Executor::KeepAlive<> executor;
    const folly::CancellationSource cancelSource;
    coro::Mutex mutex;
    coro::Baton recordPublished;
    coro::Baton recordConsumed;
    coro::Baton allTasksCompleted;
    detail::CallbackRecord<Reference> record;
  };

  auto makeConsumerTask =
      [](std::shared_ptr<SharedState> state,
         GeneratorType sources_) -> Task<void> {
    auto makeWorkerTask =
        [](std::shared_ptr<SharedState> state_,
           AsyncGenerator<Reference, Value> generator)
        -> detail::DetachedBarrierTask {
      exception_wrapper ex;
      auto cancelToken = state_->cancelSource.getToken();
      try {
        while (auto item = co_await co_viaIfAsync(
                   state_->executor.get_alias(),
                   co_withCancellation(cancelToken, generator.next()))) {
          // We have a new value to emit in the merged stream.
          {
            auto lock = co_await co_viaIfAsync(
                state_->executor.get_alias(), state_->mutex.co_scoped_lock());

            if (cancelToken.isCancellationRequested()) {
              // Consumer has detached and doesn't want any more values.
              // Discard this value.
              break;
            }

            // Publish the value.
            state_->record = detail::CallbackRecord<Reference>{
                detail::callback_record_value, *std::move(item)};
            state_->recordPublished.post();

            // Wait until the consumer is finished with it.
            co_await co_viaIfAsync(
                state_->executor.get_alias(), state_->recordConsumed);
            state_->recordConsumed.reset();

            // Clear the result before releasing the lock.
            state_->record = {};
          }

          if (cancelToken.isCancellationRequested()) {
            break;
          }
        }
      } catch (...) {
        ex = exception_wrapper{current_exception()};
      }

      if (ex) {
        state_->cancelSource.requestCancellation();

        auto lock = co_await co_viaIfAsync(
            state_->executor.get_alias(), state_->mutex.co_scoped_lock());
        if (!state_->record.hasError()) {
          state_->record = detail::CallbackRecord<Reference>{
              detail::callback_record_error, std::move(ex)};
          state_->recordPublished.post();
        }
      }
    };

    detail::Barrier barrier{1};

    auto& asyncFrame = co_await detail::co_current_async_stack_frame;

    // Save the initial context and restore it after starting each task
    // as the task may have modified the context before suspending and we
    // want to make sure the next task is started with the same initial
    // context.
    const auto context = RequestContext::saveContext();

    exception_wrapper ex;
    try {
      while (auto item = co_await sources_.next()) {
        if (state->cancelSource.isCancellationRequested()) {
          break;
        }
        makeWorkerTask(state, *std::move(item)).start(&barrier, asyncFrame);
        RequestContext::setContext(context);
      }
    } catch (...) {
      ex = exception_wrapper{current_exception()};
    }

    if (ex) {
      state->cancelSource.requestCancellation();

      auto lock = co_await co_viaIfAsync(
          state->executor.get_alias(), state->mutex.co_scoped_lock());
      if (!state->record.hasError()) {
        state->record = detail::CallbackRecord<Reference>{
            detail::callback_record_error, std::move(ex)};
        state->recordPublished.post();
      }
    }

    // Wait for all worker tasks to finish consuming the entirety of their
    // input streams.
    co_await detail::UnsafeResumeInlineSemiAwaitable{barrier.arriveAndWait()};

    // Guaranteed there are no more concurrent producers trying to acquire
    // the mutex here.
    if (!state->record.hasError()) {
      // Stream not yet been terminated with an error.
      // Terminate the stream with the 'end()' signal.
      assert(!state->record.hasValue());
      state->record =
          detail::CallbackRecord<Reference>{detail::callback_record_none};
      state->recordPublished.post();
    }
  };

  auto state = std::make_shared<SharedState>(executor);

  SCOPE_EXIT {
    state->cancelSource.requestCancellation();
    // Make sure we resume the worker thread so that it has a chance to notice
    // that cancellation has been requested.
    state->recordConsumed.post();
  };

  // Start a task that consumes the stream of input streams.
  co_withExecutor(executor, makeConsumerTask(state, std::move(sources)))
      .start(
          [state](auto&&) { state->allTasksCompleted.post(); },
          state->cancelSource.getToken());

  // Consume values produced by the input streams.
  while (true) {
    if (!state->recordPublished.ready()) {
      folly::CancellationCallback cb{
          co_await co_current_cancellation_token,
          [&] { state->cancelSource.requestCancellation(); }};
      co_await state->recordPublished;
    }
    state->recordPublished.reset();

    if (state->record.hasValue()) {
      // next value
      co_yield std::move(state->record).value();
      state->recordConsumed.post();
    } else {
      // We're closing the output stream. In the spirit of structured
      // concurrency, let's make sure to not leave any background tasks behind.
      co_await state->allTasksCompleted;

      if (state->record.hasError()) {
        std::move(state->record).error().throw_exception();
      } else {
        // none
        assert(state->record.hasNone());
        break;
      }
    }
  }
}
} // namespace detail

template <typename Reference, typename Value>
AsyncGenerator<Reference, Value> merge(
    folly::Executor::KeepAlive<> executor,
    AsyncGenerator<AsyncGenerator<Reference, Value>&&> sources) {
  return detail::mergeImpl<
      Reference,
      Value,
      AsyncGenerator<AsyncGenerator<Reference, Value>&&>>(
      std::move(executor), std::move(sources));
}

template <typename Reference, typename Value>
AsyncGenerator<Reference, Value> merge(
    folly::Executor::KeepAlive<> executor,
    AsyncGenerator<AsyncGenerator<Reference, Value>> sources) {
  return detail::mergeImpl<
      Reference,
      Value,
      AsyncGenerator<AsyncGenerator<Reference, Value>>>(
      std::move(executor), std::move(sources));
}

} // namespace coro
} // namespace folly

#endif // FOLLY_HAS_COROUTINES
