/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include "nsIThreadManager.h"
#include "nsCOMPtr.h"
#include "nsIThread.h"
#include "nsXPCOM.h"
#include "nsThreadUtils.h"
#include "nsServiceManagerUtils.h"
#include "mozilla/Atomics.h"
#include "gtest/gtest.h"
#include "mozilla/gtest/MozAssertions.h"

using mozilla::Atomic;
using mozilla::Runnable;

class WaitCondition final : public nsINestedEventLoopCondition {
 public:
  NS_DECL_THREADSAFE_ISUPPORTS

  WaitCondition(Atomic<uint32_t>& aCounter, uint32_t aMaxCount)
      : mCounter(aCounter), mMaxCount(aMaxCount) {}

  NS_IMETHODIMP IsDone(bool* aDone) override {
    *aDone = (mCounter == mMaxCount);
    return NS_OK;
  }

 private:
  ~WaitCondition() = default;

  Atomic<uint32_t>& mCounter;
  const uint32_t mMaxCount;
};

NS_IMPL_ISUPPORTS(WaitCondition, nsINestedEventLoopCondition)

class SpinRunnable final : public Runnable {
 public:
  explicit SpinRunnable(nsINestedEventLoopCondition* aCondition)
      : Runnable("SpinRunnable"), mCondition(aCondition), mResult(NS_OK) {}

  NS_IMETHODIMP Run() {
    nsCOMPtr<nsIThreadManager> threadMan =
        do_GetService("@mozilla.org/thread-manager;1");

    mResult = threadMan->SpinEventLoopUntil(
        "xpcom:TestThreadManager.cpp:SpinRunnable->Run()"_ns, mCondition);
    return NS_OK;
  }

  nsresult SpinLoopResult() { return mResult; }

 private:
  ~SpinRunnable() = default;

  nsCOMPtr<nsINestedEventLoopCondition> mCondition;
  Atomic<nsresult> mResult;
};

class CountRunnable final : public Runnable {
 public:
  explicit CountRunnable(Atomic<uint32_t>& aCounter)
      : Runnable("CountRunnable"), mCounter(aCounter) {}

  NS_IMETHODIMP Run() {
    mCounter++;
    return NS_OK;
  }

 private:
  Atomic<uint32_t>& mCounter;
};

TEST(ThreadManager, SpinEventLoopUntilSuccess)
{
  const uint32_t kRunnablesToDispatch = 100;
  nsresult rv;
  mozilla::Atomic<uint32_t> count(0);

  RefPtr condition =
      mozilla::MakeRefPtr<WaitCondition>(count, kRunnablesToDispatch);
  RefPtr spinner = mozilla::MakeRefPtr<SpinRunnable>(condition);
  nsCOMPtr<nsIThread> thread;
  rv = NS_NewNamedThread("SpinEventLoop", getter_AddRefs(thread), spinner);
  ASSERT_NS_SUCCEEDED(rv);

  RefPtr counter = mozilla::MakeRefPtr<CountRunnable>(count);
  for (uint32_t i = 0; i < kRunnablesToDispatch; ++i) {
    rv = thread->Dispatch(counter, NS_DISPATCH_NORMAL);
    ASSERT_NS_SUCCEEDED(rv);
  }

  rv = thread->Shutdown();
  ASSERT_NS_SUCCEEDED(rv);
  ASSERT_NS_SUCCEEDED(spinner->SpinLoopResult());
}

class ErrorCondition final : public nsINestedEventLoopCondition {
 public:
  NS_DECL_THREADSAFE_ISUPPORTS

  ErrorCondition(Atomic<uint32_t>& aCounter, uint32_t aMaxCount)
      : mCounter(aCounter), mMaxCount(aMaxCount) {}

  NS_IMETHODIMP IsDone(bool* aDone) override {
    if (mCounter == mMaxCount) {
      return NS_ERROR_ILLEGAL_VALUE;
    }
    return NS_OK;
  }

 private:
  ~ErrorCondition() = default;

  Atomic<uint32_t>& mCounter;
  const uint32_t mMaxCount;
};

NS_IMPL_ISUPPORTS(ErrorCondition, nsINestedEventLoopCondition)

TEST(ThreadManager, SpinEventLoopUntilError)
{
  const uint32_t kRunnablesToDispatch = 100;
  nsresult rv;
  mozilla::Atomic<uint32_t> count(0);

  RefPtr condition =
      mozilla::MakeRefPtr<ErrorCondition>(count, kRunnablesToDispatch);
  RefPtr spinner = mozilla::MakeRefPtr<SpinRunnable>(condition);
  nsCOMPtr<nsIThread> thread;
  rv = NS_NewNamedThread("SpinEventLoop", getter_AddRefs(thread), spinner);
  ASSERT_NS_SUCCEEDED(rv);

  RefPtr counter = mozilla::MakeRefPtr<CountRunnable>(count);
  for (uint32_t i = 0; i < kRunnablesToDispatch; ++i) {
    rv = thread->Dispatch(counter, NS_DISPATCH_NORMAL);
    ASSERT_NS_SUCCEEDED(rv);
  }

  rv = thread->Shutdown();
  ASSERT_NS_SUCCEEDED(rv);
  ASSERT_NS_FAILED(spinner->SpinLoopResult());
}