#include "gtest/gtest.h"

#include "mozilla/NonBlockingAsyncInputStream.h"

#include "nsIAsyncInputStream.h"

#include "nsStreamUtils.h"

#include "nsString.h"

#include "nsStringStream.h"

#include "Helpers.h"

TEST(TestNonBlockingAsyncInputStream, Simple) {

  nsCString data;

  data.Assign("Hello world!");

  // It should not be async.

  bool nonBlocking = false;

  nsCOMPtr<nsIAsyncInputStream> async;

  {

    // Let's create a test string inputStream

    nsCOMPtr<nsIInputStream> stream;

    ASSERT_EQ(NS_OK, NS_NewCStringInputStream(getter_AddRefs(stream), data));

    async = do_QueryInterface(stream);

    ASSERT_EQ(nullptr, async);

    // It must be non-blocking

    ASSERT_EQ(NS_OK, stream->IsNonBlocking(&nonBlocking));

    ASSERT_TRUE(nonBlocking);

    // Here the non-blocking stream.

    ASSERT_EQ(NS_OK,

              NonBlockingAsyncInputStream::Create(stream.forget(),

                                                  getter_AddRefs(async)));

  }

  ASSERT_TRUE(!!async);

  // Still non-blocking

  ASSERT_EQ(NS_OK, async->IsNonBlocking(&nonBlocking));

  ASSERT_TRUE(nonBlocking);

  // Testing ::Available()

  uint64_t length;

  ASSERT_EQ(NS_OK, async->Available(&length));

  ASSERT_EQ(data.Length(), length);

  // Read works fine.

  char buffer[1024];

  uint32_t read = 0;

  ASSERT_EQ(NS_OK, async->Read(buffer, sizeof(buffer), &read));

  ASSERT_EQ(data.Length(), read);

  ASSERT_TRUE(data.Equals(nsCString(buffer, read)));

}

class ReadSegmentsData

{

public:

  ReadSegmentsData(nsIInputStream* aStream, char* aBuffer)

    : mStream(aStream)

    , mBuffer(aBuffer)

  {}

  nsIInputStream* mStream;

  char* mBuffer;

};

nsresult

ReadSegmentsFunction(nsIInputStream* aInStr,

                     void* aClosure,

                     const char* aBuffer,

                     uint32_t aOffset,

                     uint32_t aCount,

                     uint32_t* aCountWritten)

{

  ReadSegmentsData* data = static_cast<ReadSegmentsData*>(aClosure);

  if (aInStr != data->mStream) return NS_ERROR_FAILURE;

  memcpy(&data->mBuffer[aOffset], aBuffer, aCount);

  *aCountWritten = aCount;

  return NS_OK;

}

TEST(TestNonBlockingAsyncInputStream, ReadSegments) {

  nsCString data;

  data.Assign("Hello world!");

  nsCOMPtr<nsIAsyncInputStream> async;

  {

    // Let's create a test string inputStream

    nsCOMPtr<nsIInputStream> stream;

    ASSERT_EQ(NS_OK, NS_NewCStringInputStream(getter_AddRefs(stream), data));

    // Here the non-blocking stream.

    ASSERT_EQ(NS_OK,

              NonBlockingAsyncInputStream::Create(stream.forget(),

                                                  getter_AddRefs(async)));

  }

  // Read works fine.

  char buffer[1024];

  uint32_t read = 0;

  ReadSegmentsData closure(async, buffer);

  ASSERT_EQ(NS_OK, async->ReadSegments(ReadSegmentsFunction, &closure,

                                       sizeof(buffer), &read));

  ASSERT_EQ(data.Length(), read);

  ASSERT_TRUE(data.Equals(nsCString(buffer, read)));

}

TEST(TestNonBlockingAsyncInputStream, AsyncWait_Simple) {

  nsCString data;

  data.Assign("Hello world!");

  nsCOMPtr<nsIAsyncInputStream> async;

  {

    // Let's create a test string inputStream

    nsCOMPtr<nsIInputStream> stream;

    ASSERT_EQ(NS_OK, NS_NewCStringInputStream(getter_AddRefs(stream), data));

    // Here the non-blocking stream.

    ASSERT_EQ(NS_OK,

              NonBlockingAsyncInputStream::Create(stream.forget(),

                                                  getter_AddRefs(async)));

  }

  ASSERT_TRUE(!!async);

  // Testing ::Available()

  uint64_t length;

  ASSERT_EQ(NS_OK, async->Available(&length));

  ASSERT_EQ(data.Length(), length);

  // Testing ::AsyncWait - without EventTarget

  RefPtr<testing::InputStreamCallback> cb =

    new testing::InputStreamCallback();

  ASSERT_EQ(NS_OK, async->AsyncWait(cb, 0, 0, nullptr));

  ASSERT_TRUE(cb->Called());

  // Testing ::AsyncWait - with EventTarget

  cb = new testing::InputStreamCallback();

  nsCOMPtr<nsIThread> thread = do_GetCurrentThread();

  ASSERT_EQ(NS_OK, async->AsyncWait(cb, 0, 0, thread));

  ASSERT_FALSE(cb->Called());

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return cb->Called(); }));

  ASSERT_TRUE(cb->Called());

  // Read works fine.

  char buffer[1024];

  uint32_t read = 0;

  ASSERT_EQ(NS_OK, async->Read(buffer, sizeof(buffer), &read));

  ASSERT_EQ(data.Length(), read);

  ASSERT_TRUE(data.Equals(nsCString(buffer, read)));

}

TEST(TestNonBlockingAsyncInputStream, AsyncWait_ClosureOnly_withoutEventTarget) {

  nsCString data;

  data.Assign("Hello world!");

  nsCOMPtr<nsIAsyncInputStream> async;

  {

    // Let's create a test string inputStream

    nsCOMPtr<nsIInputStream> stream;

    ASSERT_EQ(NS_OK, NS_NewCStringInputStream(getter_AddRefs(stream), data));

    // Here the non-blocking stream.

    ASSERT_EQ(NS_OK,

              NonBlockingAsyncInputStream::Create(stream.forget(),

                                                  getter_AddRefs(async)));

  }

  ASSERT_TRUE(!!async);

  // Testing ::AsyncWait - no eventTarget

  RefPtr<testing::InputStreamCallback> cb =

    new testing::InputStreamCallback();

  ASSERT_EQ(NS_OK, async->AsyncWait(cb, nsIAsyncInputStream::WAIT_CLOSURE_ONLY, 0, nullptr));

  ASSERT_FALSE(cb->Called());

  ASSERT_EQ(NS_OK, async->Close());

  ASSERT_TRUE(cb->Called());

}

TEST(TestNonBlockingAsyncInputStream, AsyncWait_ClosureOnly_withEventTarget) {

  nsCString data;

  data.Assign("Hello world!");

  nsCOMPtr<nsIAsyncInputStream> async;

  {

    // Let's create a test string inputStream

    nsCOMPtr<nsIInputStream> stream;

    ASSERT_EQ(NS_OK, NS_NewCStringInputStream(getter_AddRefs(stream), data));

    // Here the non-blocking stream.

    ASSERT_EQ(NS_OK,

              NonBlockingAsyncInputStream::Create(stream.forget(),

                                                  getter_AddRefs(async)));

  }

  ASSERT_TRUE(!!async);

  // Testing ::AsyncWait - with EventTarget

  RefPtr<testing::InputStreamCallback> cb =

    new testing::InputStreamCallback();

  nsCOMPtr<nsIThread> thread = do_GetCurrentThread();

  ASSERT_EQ(NS_OK, async->AsyncWait(cb, nsIAsyncInputStream::WAIT_CLOSURE_ONLY, 0, thread));

  ASSERT_FALSE(cb->Called());

  ASSERT_EQ(NS_OK, async->Close());

  ASSERT_FALSE(cb->Called());

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil([&]() { return cb->Called(); }));

  ASSERT_TRUE(cb->Called());

}

TEST(TestNonBlockingAsyncInputStream, Helper) {

  nsCString data;

  data.Assign("Hello world!");

  nsCOMPtr<nsIAsyncInputStream> async;

  {

    // Let's create a test string inputStream

    nsCOMPtr<nsIInputStream> stream;

    ASSERT_EQ(NS_OK, NS_NewCStringInputStream(getter_AddRefs(stream), data));

    // Here the non-blocking stream.

    ASSERT_EQ(NS_OK,

              NonBlockingAsyncInputStream::Create(stream.forget(),

                                                  getter_AddRefs(async)));

  }

  ASSERT_TRUE(!!async);

  // This should return the same object because async is already non-blocking

  // and async.

  nsCOMPtr<nsIAsyncInputStream> result;

  nsCOMPtr<nsIAsyncInputStream> asyncTmp = async;

  ASSERT_EQ(NS_OK,

            NS_MakeAsyncNonBlockingInputStream(asyncTmp.forget(),

                                               getter_AddRefs(result)));

  ASSERT_EQ(async, result);

  // This will use NonBlockingAsyncInputStream wrapper.

  {

    nsCOMPtr<nsIInputStream> stream;

    ASSERT_EQ(NS_OK, NS_NewCStringInputStream(getter_AddRefs(stream), data));

    ASSERT_EQ(NS_OK,

              NS_MakeAsyncNonBlockingInputStream(stream.forget(),

                                                 getter_AddRefs(result)));

  }

  ASSERT_TRUE(async != result);

  ASSERT_TRUE(async);

}

class QIInputStream final : public nsIInputStream

                          , public nsICloneableInputStream

                          , public nsIIPCSerializableInputStream

                          , public nsISeekableStream

{

public:

  NS_DECL_ISUPPORTS

  QIInputStream(bool aNonBlockingError, bool aCloneable,

                bool aIPCSerializable, bool aSeekable)

    : mNonBlockingError(aNonBlockingError)

    , mCloneable(aCloneable)

    , mIPCSerializable(aIPCSerializable)

    , mSeekable(aSeekable)

  {}

  // nsIInputStream

  NS_IMETHOD Close() override { return NS_ERROR_NOT_IMPLEMENTED; }

  NS_IMETHOD Available(uint64_t*) override { return NS_ERROR_NOT_IMPLEMENTED; }

  NS_IMETHOD Read(char*, uint32_t, uint32_t*) override { return NS_ERROR_NOT_IMPLEMENTED; }

  NS_IMETHOD ReadSegments(nsWriteSegmentFun, void*, uint32_t, uint32_t*) override { return NS_ERROR_NOT_IMPLEMENTED; }

  NS_IMETHOD IsNonBlocking(bool* aNonBlocking) override

  {

    *aNonBlocking = true;

    return mNonBlockingError ? NS_ERROR_FAILURE : NS_OK;

  }

  // nsICloneableInputStream

  NS_IMETHOD GetCloneable(bool*) override { return NS_ERROR_NOT_IMPLEMENTED; }

  NS_IMETHOD Clone(nsIInputStream**) override { return NS_ERROR_NOT_IMPLEMENTED; }

  // nsIIPCSerializableInputStream

  void Serialize(mozilla::ipc::InputStreamParams&, FileDescriptorArray&) override {}

  bool Deserialize(const mozilla::ipc::InputStreamParams&,

                   const FileDescriptorArray&) override { return false; }

  mozilla::Maybe<uint64_t> ExpectedSerializedLength() override { return mozilla::Nothing(); }

  // nsISeekableStream

  NS_IMETHOD Seek(int32_t, int64_t) override { return NS_ERROR_NOT_IMPLEMENTED; }

  NS_IMETHOD Tell(int64_t*) override { return NS_ERROR_NOT_IMPLEMENTED; }

  NS_IMETHOD SetEOF() override { return NS_ERROR_NOT_IMPLEMENTED; }

private:

  ~QIInputStream() = default;

  bool mNonBlockingError;

  bool mCloneable;

  bool mIPCSerializable;

  bool mSeekable;

};

NS_IMPL_ADDREF(QIInputStream);

NS_IMPL_RELEASE(QIInputStream);

NS_INTERFACE_MAP_BEGIN(QIInputStream)

  NS_INTERFACE_MAP_ENTRY(nsIInputStream)

  NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsICloneableInputStream, mCloneable)

  NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsIIPCSerializableInputStream, mIPCSerializable)

  NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsISeekableStream, mSeekable)

  NS_INTERFACE_MAP_ENTRY_AMBIGUOUS(nsISupports, nsIInputStream)

NS_INTERFACE_MAP_END

TEST(TestNonBlockingAsyncInputStream, QI) {

  // Let's test ::Create() returning error.

  nsCOMPtr<nsIAsyncInputStream> async;

  {

    nsCOMPtr<nsIInputStream> stream = new QIInputStream(true, true, true, true);

    ASSERT_EQ(NS_ERROR_FAILURE,

              NonBlockingAsyncInputStream::Create(stream.forget(),

                                                  getter_AddRefs(async)));

  }

  // Let's test the QIs

  for (int i = 0; i < 8; ++i) {

    bool shouldBeCloneable = !!(i & 0x01);

    bool shouldBeSerializable = !!(i & 0x02);

    bool shouldBeSeekable = !!(i & 0x04);

    nsCOMPtr<nsICloneableInputStream> cloneable;

    nsCOMPtr<nsIIPCSerializableInputStream> ipcSerializable;

    nsCOMPtr<nsISeekableStream> seekable;

    {

      nsCOMPtr<nsIInputStream> stream =

        new QIInputStream(false, shouldBeCloneable, shouldBeSerializable, shouldBeSeekable);

      cloneable = do_QueryInterface(stream);

      ASSERT_EQ(shouldBeCloneable, !!cloneable);

      ipcSerializable = do_QueryInterface(stream);

      ASSERT_EQ(shouldBeSerializable, !!ipcSerializable);

      seekable = do_QueryInterface(stream);

      ASSERT_EQ(shouldBeSeekable, !!seekable);

      ASSERT_EQ(NS_OK, NonBlockingAsyncInputStream::Create(stream.forget(),

                                                           getter_AddRefs(async)));

    }

    // The returned async stream should be cloneable only if the underlying

    // stream is.

    cloneable = do_QueryInterface(async);

    ASSERT_EQ(shouldBeCloneable, !!cloneable);

    // The returned async stream should be serializable only if the underlying

    // stream is.

    ipcSerializable = do_QueryInterface(async);

    ASSERT_EQ(shouldBeSerializable, !!ipcSerializable);

    // The returned async stream should be seekable only if the underlying

    // stream is.

    seekable = do_QueryInterface(async);

    ASSERT_EQ(shouldBeSeekable, !!seekable);

  }

}