/*
 * Copyright (c) Facebook, Inc. and its 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.
 */

#pragma once

#include <atomic>

#include <folly/Portability.h>
#include <folly/detail/Futex.h>
#include <folly/experimental/coro/Coroutine.h>
#include <folly/io/async/HHWheelTimer.h>

namespace folly {
namespace fibers {

class Fiber;
class FiberManager;

/**
 * @class Baton
 *
 * Primitive which allows one to put current Fiber to sleep and wake it from
 * another Fiber/thread.
 */
class Baton {
 public:
  class TimeoutHandler;

  class Waiter {
   public:
    virtual void post() = 0;

    virtual ~Waiter() {}
  };

  Baton() noexcept;

  ~Baton() noexcept = default;

  bool ready() const {
    auto state = waiter_.load();
    return state == POSTED;
  }

  /**
   * Registers a waiter for the baton. The waiter will be notified when
   * the baton is posted.
   */
  void setWaiter(Waiter& waiter);

  /**
   * Puts active fiber to sleep. Returns when post is called.
   */
  void wait();

  /**
   * Put active fiber to sleep indefinitely. However, timeoutHandler may
   * be used elsewhere on the same thread in order to schedule a wakeup
   * for the active fiber.  Users of timeoutHandler must be on the same thread
   * as the active fiber and may only schedule one timeout, which must occur
   * after the active fiber calls wait.
   */
  void wait(TimeoutHandler& timeoutHandler);

  /**
   * Puts active fiber to sleep. Returns when post is called.
   *
   * @param mainContextFunc this function is immediately executed on the main
   *        context.
   */
  template <typename F>
  void wait(F&& mainContextFunc);

  /**
   * Checks if the baton has been posted without blocking.
   *
   * @return    true iff the baton has been posted.
   */
  bool try_wait();

  /**
   * Puts active fiber to sleep. Returns when post is called or the timeout
   * expires.
   *
   * @param timeout Baton will be automatically awaken if timeout expires
   *
   * @return true if was posted, false if timeout expired
   */
  template <typename Rep, typename Period>
  bool try_wait_for(const std::chrono::duration<Rep, Period>& timeout) {
    return try_wait_for(timeout, [] {});
  }

  /**
   * Puts active fiber to sleep. Returns when post is called or the timeout
   * expires.
   *
   * @param timeout Baton will be automatically awaken if timeout expires
   * @param mainContextFunc this function is immediately executed on the main
   *        context.
   *
   * @return true if was posted, false if timeout expired
   */
  template <typename Rep, typename Period, typename F>
  bool try_wait_for(
      const std::chrono::duration<Rep, Period>& timeout, F&& mainContextFunc);

  /**
   * Puts active fiber to sleep. Returns when post is called or the deadline
   * expires.
   *
   * @param deadline Baton will be automatically awaken if deadline expires
   *
   * @return true if was posted, false if timeout expired
   */
  template <typename Clock, typename Duration>
  bool try_wait_until(
      const std::chrono::time_point<Clock, Duration>& deadline) {
    return try_wait_until(deadline, [] {});
  }

  /**
   * Puts active fiber to sleep. Returns when post is called or the deadline
   * expires.
   *
   * @param deadline Baton will be automatically awaken if deadline expires
   * @param mainContextFunc this function is immediately executed on the main
   *        context.
   *
   * @return true if was posted, false if timeout expired
   */
  template <typename Clock, typename Duration, typename F>
  bool try_wait_until(
      const std::chrono::time_point<Clock, Duration>& deadline,
      F&& mainContextFunc);

  /**
   * Puts active fiber to sleep. Returns when post is called or the deadline
   * expires.
   *
   * @param deadline Baton will be automatically awaken if deadline expires
   * @param mainContextFunc this function is immediately executed on the main
   *        context.
   *
   * @return true if was posted, false if timeout expired
   */
  template <typename Clock, typename Duration, typename F>
  bool try_wait_for(
      const std::chrono::time_point<Clock, Duration>& deadline,
      F&& mainContextFunc);

  /// Alias to try_wait_for. Deprecated.
  template <typename Rep, typename Period>
  bool timed_wait(const std::chrono::duration<Rep, Period>& timeout) {
    return try_wait_for(timeout);
  }

  /// Alias to try_wait_for. Deprecated.
  template <typename Rep, typename Period, typename F>
  bool timed_wait(
      const std::chrono::duration<Rep, Period>& timeout, F&& mainContextFunc) {
    return try_wait_for(timeout, static_cast<F&&>(mainContextFunc));
  }

  /// Alias to try_wait_until. Deprecated.
  template <typename Clock, typename Duration>
  bool timed_wait(const std::chrono::time_point<Clock, Duration>& deadline) {
    return try_wait_until(deadline);
  }

  /// Alias to try_wait_until. Deprecated.
  template <typename Clock, typename Duration, typename F>
  bool timed_wait(
      const std::chrono::time_point<Clock, Duration>& deadline,
      F&& mainContextFunc) {
    return try_wait_until(deadline, static_cast<F&&>(mainContextFunc));
  }

  /**
   * Wakes up Fiber which was waiting on this Baton (or if no Fiber is waiting,
   * next wait() call will return immediately).
   */
  void post();

  /**
   * Reset's the baton (equivalent to destroying the object and constructing
   * another one in place).
   * Caller is responsible for making sure no one is waiting on/posting the
   * baton when reset() is called.
   */
  void reset();

  /**
   * Provides a way to schedule a wakeup for a wait()ing fiber.
   * A TimeoutHandler must be passed to Baton::wait(TimeoutHandler&)
   * before a timeout is scheduled. It is only safe to use the
   * TimeoutHandler on the same thread as the wait()ing fiber.
   * scheduleTimeout() may only be called once prior to the end of the
   * associated Baton's life.
   */
  class TimeoutHandler final : public HHWheelTimer::Callback {
   public:
    void scheduleTimeout(std::chrono::milliseconds timeout);

   private:
    friend class Baton;

    std::function<void()> timeoutFunc_{nullptr};
    FiberManager* fiberManager_{nullptr};

    void timeoutExpired() noexcept override {
      assert(timeoutFunc_ != nullptr);
      timeoutFunc_();
    }

    void callbackCanceled() noexcept override {}
  };

 private:
  class FiberWaiter;

  explicit Baton(intptr_t state) : waiter_(state) {}

  void postHelper(intptr_t new_value);
  void postThread();
  void waitThread();

  template <typename F>
  inline void waitFiber(FiberManager& fm, F&& mainContextFunc);

  template <typename Clock, typename Duration>
  bool timedWaitThread(
      const std::chrono::time_point<Clock, Duration>& deadline);

  static constexpr intptr_t NO_WAITER = 0;
  static constexpr intptr_t POSTED = -1;
  static constexpr intptr_t TIMEOUT = -2;
  static constexpr intptr_t THREAD_WAITING = -3;

  union {
    std::atomic<intptr_t> waiter_;
    struct {
      folly::detail::Futex<> futex{};
      int32_t _unused_packing;
    } futex_;
  };
};

#if FOLLY_HAS_COROUTINES
namespace detail {
class BatonAwaitableWaiter : public Baton::Waiter {
 public:
  explicit BatonAwaitableWaiter(Baton& baton) : baton_(baton) {}

  void post() override {
    assert(h_);
    h_();
  }

  bool await_ready() const { return baton_.ready(); }

  void await_resume() {}

  void await_suspend(coro::coroutine_handle<> h) {
    assert(!h_);
    h_ = std::move(h);
    baton_.setWaiter(*this);
  }

 private:
  coro::coroutine_handle<> h_;
  Baton& baton_;
};
} // namespace detail

inline detail::BatonAwaitableWaiter /* implicit */ operator co_await(
    Baton& baton) {
  return detail::BatonAwaitableWaiter(baton);
}
#endif
} // namespace fibers
} // namespace folly

#include <folly/fibers/Baton-inl.h>