/*
 * 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 <folly/coro/BoundedQueue.h>

#include <chrono>
#include <optional>
#include <string>
#include <thread>

#include <folly/CancellationToken.h>
#include <folly/Portability.h>
#include <folly/coro/BlockingWait.h>
#include <folly/coro/Collect.h>
#include <folly/coro/GtestHelpers.h>
#include <folly/coro/Sleep.h>
#include <folly/portability/GTest.h>
#if FOLLY_HAS_COROUTINES

namespace {
struct SlowMover {
  explicit SlowMover(bool slow = false) : slow(slow) {}
  SlowMover(SlowMover&& other) noexcept { *this = std::move(other); }
  SlowMover& operator=(SlowMover&& other) noexcept {
    slow = other.slow;
    if (slow) {
      /* sleep override */ std::this_thread::sleep_for(
          std::chrono::milliseconds(50));
    }
    return *this;
  }

  bool slow;
};
} // namespace

CO_TEST(BoundedQueueTest, EnqueueDeque) {
  folly::coro::BoundedQueue<std::string, true, true> queue(100);
  constexpr auto val = "a string";
  std::string val1 = val;
  EXPECT_TRUE(queue.empty());
  EXPECT_EQ(queue.size(), 0);

  co_await queue.enqueue(val1);
  EXPECT_FALSE(queue.empty());
  co_await queue.enqueue(std::move(val1));
  EXPECT_EQ(queue.size(), 2);

  for (int i = 0; i < 2; ++i) {
    auto val2 = co_await queue.dequeue();
    EXPECT_EQ(val2, val);
  }
  EXPECT_TRUE(queue.empty());
}

CO_TEST(BoundedQueueTest, DequeueWhileBlocking) {
  folly::coro::BoundedQueue<int> queue(5);
  folly::ManualExecutor ex;

  auto fut = queue.dequeue().scheduleOn(&ex).start();
  ex.drain();
  EXPECT_FALSE(fut.isReady());
  co_await queue.enqueue(0);

  ex.drain();
  EXPECT_TRUE(fut.isReady());
  EXPECT_EQ(std::move(fut).get(), 0);
}

CO_TEST(BoundedQueueTest, EnqueueDequeMultiProducer) {
  folly::coro::BoundedQueue<int, false, true> queue(5);
  std::atomic<int> i = 0;

  std::vector<std::thread> enqueuers;
  for (int n = 0; n < 5; ++n) {
    enqueuers.emplace_back([&] {
      while (true) {
        int next = i++;
        if (next >= 100) {
          break;
        }
        folly::coro::blockingWait(
            [&, next]() mutable -> folly::coro::Task<void> {
              co_await queue.enqueue(std::move(next));
            }());
      }
    });
  }
  for (int n = 0; n < 100; ++n) {
    co_await queue.dequeue();
  }

  EXPECT_TRUE(queue.empty());

  for (int n = 0; n < 5; ++n) {
    enqueuers[n].join();
  }
}

CO_TEST(BoundedQueueTest, EnqueueDequeMultiConsumer) {
  folly::coro::BoundedQueue<int, true, false> queue(10);
  std::atomic<int> seen = 0;

  std::vector<std::thread> dequeuers;
  for (int n = 0; n < 5; ++n) {
    dequeuers.emplace_back([&] {
      folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
        while (++seen <= 100) {
          co_await queue.dequeue();
        }
      }());
    });
  }

  for (int n = 0; n < 100; ++n) {
    co_await queue.enqueue(std::move(n));
  }
  for (int n = 0; n < 5; ++n) {
    dequeuers[n].join();
  }
  EXPECT_TRUE(queue.empty());
}

CO_TEST(BoundedQueueTest, EnqueueDequeMPMCWithSingleSlot) {
  folly::coro::BoundedQueue<int, false, false> queue(1);
  std::atomic<int> seen = 0, i = 0;

  std::vector<std::thread> enqueuers;
  for (int n = 0; n < 5; ++n) {
    enqueuers.emplace_back([&] {
      folly::coro::blockingWait([&]() mutable -> folly::coro::Task<void> {
        while (true) {
          int next = i++;
          if (next >= 100) {
            break;
          }
          co_await queue.enqueue(std::move(next));
        }
      }());
    });
  }

  std::vector<std::thread> dequeuers;
  for (int n = 0; n < 5; ++n) {
    dequeuers.emplace_back([&] {
      folly::coro::blockingWait([&]() -> folly::coro::Task<void> {
        while (++seen <= 100) {
          co_await queue.dequeue();
        }
      }());
    });
  }

  for (int n = 0; n < 5; ++n) {
    enqueuers[n].join();
  }
  for (int n = 0; n < 5; ++n) {
    dequeuers[n].join();
  }
  EXPECT_TRUE(queue.empty());

  co_return;
}

CO_TEST(
    BoundedQueueTest, CancelledDequeueCompletesNormallyIfAnItemIsAvailable) {
  folly::coro::BoundedQueue<int> queue(10);
  folly::CancellationSource cancelSource;
  cancelSource.requestCancellation();

  co_await queue.enqueue(123);

  int result = co_await folly::coro::co_withCancellation(
      cancelSource.getToken(), queue.dequeue());
  EXPECT_EQ(123, result);
}

CO_TEST(BoundedQueueTest, EnqueueWait) {
  folly::coro::BoundedQueue<int> queue(2);
  co_await folly::coro::collectAll(
      [&]() -> folly::coro::Task<void> {
        for (int i = 0; i < 100; i++) {
          auto val = i;
          co_await queue.enqueue(std::move(val));
        }
      }(),
      [&]() -> folly::coro::Task<void> {
        for (int i = 0; i < 100; i++) {
          int val = co_await queue.dequeue();
          EXPECT_EQ(val, i);
        }
      }());
}

CO_TEST(BoundedQueueTest, DequeueWait) {
  folly::coro::BoundedQueue<int> queue(2);
  co_await folly::coro::collectAll(
      [&]() -> folly::coro::Task<void> {
        for (int i = 0; i < 100; i++) {
          auto val = i;
          co_await queue.enqueue(std::move(val));
        }
      }(),
      [&]() -> folly::coro::Task<void> {
        for (int i = 0; i < 100; i++) {
          int val = co_await queue.dequeue();
          EXPECT_EQ(val, i);
        }
      }());
}

CO_TEST(BoundedQueueTest, TryEnqueue) {
  folly::coro::BoundedQueue<int> queue(2);

  EXPECT_TRUE(queue.try_enqueue(1));
  EXPECT_TRUE(queue.try_enqueue(1));
  EXPECT_FALSE(queue.try_enqueue(1));
  co_await queue.dequeue();
  EXPECT_TRUE(queue.try_enqueue(1));
}

CO_TEST(BoundedQueueTest, TryDequeue) {
  folly::coro::BoundedQueue<int> queue(2);

  EXPECT_FALSE(queue.try_dequeue().has_value());
  co_await queue.enqueue(1);
  EXPECT_TRUE(queue.try_dequeue().has_value());
}

TEST(BoundedQueueTest, UnorderedEnqueueCompletion) {
  // Use optional to verify we're not accidentally dequeueing
  // default-constructed values.
  folly::coro::BoundedQueue<std::optional<SlowMover>> queue(1024);
  std::atomic<int> turn = 0;

  std::thread consumer([&] {
    ++turn;
    for (size_t i = 0; i < 2; ++i) {
      ASSERT_TRUE(folly::coro::blockingWait(queue.dequeue()).has_value());
    }
  });

  // producer2 will frequently initiate the enqueue after producer1 (thus
  // acquiring a larger ticket) but complete the move after it. The consumer
  // thus needs to block until the head-of-line item is available.
  std::thread producer1([&] {
    ++turn;
    while (turn < 3) {
    }
    ++turn;
    ASSERT_TRUE(queue.try_enqueue(SlowMover(true)));
  });
  std::thread producer2([&] {
    ++turn;
    while (turn < 4) {
    }
    /* sleep override */ std::this_thread::sleep_for(
        std::chrono::milliseconds(1));
    ASSERT_TRUE(queue.try_enqueue(SlowMover(false)));
  });

  producer1.join();
  producer2.join();
  consumer.join();
}

TEST(BoundedQueueTest, UnorderedDequeueCompletion) {
  // Use optional to verify we're not accidentally dequeueing
  // default-constructed values.
  folly::coro::BoundedQueue<std::optional<SlowMover>> queue(2);

  ASSERT_TRUE(queue.try_enqueue(SlowMover(true)));
  ASSERT_TRUE(queue.try_enqueue(SlowMover(false)));

  std::vector<std::thread> consumers;
  for (size_t i = 0; i < 3; ++i) {
    consumers.emplace_back([&]() {
      ASSERT_TRUE(folly::coro::blockingWait(queue.dequeue()).has_value());
    });
  }

  // The producer will get the ticket for the slow moving slot which will still
  // be in the process of dequeuing, so the producer needs to block until it
  // finishes and the slot becomes available.
  std::thread producer(
      [&] { folly::coro::blockingWait(queue.enqueue(SlowMover(false))); });

  producer.join();
  for (auto& consumer : consumers) {
    consumer.join();
  }
}

#endif
