#include "gtest/gtest.h"

#include "mozilla/InputStreamLengthHelper.h"
#include "mozilla/SpinEventLoopUntil.h"
#include "nsCOMPtr.h"
#include "nsIInputStream.h"
#include "nsStreamUtils.h"
#include "nsString.h"
#include "nsStringStream.h"
#include "nsThreadUtils.h"
#include "nsXPCOM.h"
#include "Helpers.h"

using namespace mozilla;

TEST(TestInputStreamLengthHelper, NonLengthStream)
{
  nsCString buf;
  buf.AssignLiteral("Hello world");

  nsCOMPtr<nsIInputStream> stream;
  NS_NewCStringInputStream(getter_AddRefs(stream), buf);

  bool called = false;
  InputStreamLengthHelper::GetAsyncLength(stream, [&](int64_t aLength) {
    ASSERT_EQ(int64_t(buf.Length()), aLength);
    called = true;
  });

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil(
      "xpcom:TEST(TestInputStreamLengthHelper, NonLengthStream)"_ns,
      [&]() { return called; }));
}

class LengthStream final : public nsIInputStreamLength,
                           public nsIAsyncInputStreamLength,
                           public nsIInputStream {
 public:
  NS_DECL_ISUPPORTS

  LengthStream(int64_t aLength, nsresult aLengthRv, uint64_t aAvailable,
               bool aIsAsyncLength)
      : mLength(aLength),
        mLengthRv(aLengthRv),
        mAvailable(aAvailable),
        mIsAsyncLength(aIsAsyncLength) {}

  NS_IMETHOD Close(void) override { MOZ_CRASH("Invalid call!"); }
  NS_IMETHOD Read(char* aBuf, uint32_t aCount, uint32_t* _retval) override {
    MOZ_CRASH("Invalid call!");
  }
  NS_IMETHOD ReadSegments(nsWriteSegmentFun aWriter, void* aClosure,
                          uint32_t aCount, uint32_t* _retval) override {
    MOZ_CRASH("Invalid call!");
  }
  NS_IMETHOD IsNonBlocking(bool* _retval) override {
    MOZ_CRASH("Invalid call!");
  }

  NS_IMETHOD Length(int64_t* aLength) override {
    *aLength = mLength;
    return mLengthRv;
  }

  NS_IMETHOD AsyncLengthWait(nsIInputStreamLengthCallback* aCallback,
                             nsIEventTarget* aEventTarget) override {
    if (aCallback) {
      aCallback->OnInputStreamLengthReady(this, mLength);
    }
    return NS_OK;
  }

  NS_IMETHOD Available(uint64_t* aAvailable) override {
    *aAvailable = mAvailable;
    return NS_OK;
  }

 private:
  ~LengthStream() = default;

  int64_t mLength;
  nsresult mLengthRv;
  uint64_t mAvailable;

  bool mIsAsyncLength;
};

NS_IMPL_ADDREF(LengthStream);
NS_IMPL_RELEASE(LengthStream);

NS_INTERFACE_MAP_BEGIN(LengthStream)
  NS_INTERFACE_MAP_ENTRY(nsIInputStream)
  NS_INTERFACE_MAP_ENTRY(nsIInputStreamLength)
  NS_INTERFACE_MAP_ENTRY_CONDITIONAL(nsIAsyncInputStreamLength, mIsAsyncLength)
  NS_INTERFACE_MAP_ENTRY_AMBIGUOUS(nsISupports, nsIInputStream)
NS_INTERFACE_MAP_END

TEST(TestInputStreamLengthHelper, LengthStream)
{
  nsCOMPtr<nsIInputStream> stream = new LengthStream(42, NS_OK, 0, false);

  bool called = false;
  InputStreamLengthHelper::GetAsyncLength(stream, [&](int64_t aLength) {
    ASSERT_EQ(42, aLength);
    called = true;
  });

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil(
      "xpcom:TEST(TestInputStreamLengthHelper, LengthStream)"_ns,
      [&]() { return called; }));
}

TEST(TestInputStreamLengthHelper, InvalidLengthStream)
{
  nsCOMPtr<nsIInputStream> stream =
      new LengthStream(42, NS_ERROR_NOT_AVAILABLE, 0, false);

  bool called = false;
  InputStreamLengthHelper::GetAsyncLength(stream, [&](int64_t aLength) {
    ASSERT_EQ(-1, aLength);
    called = true;
  });

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil(
      "xpcom:TEST(TestInputStreamLengthHelper, InvalidLengthStream)"_ns,
      [&]() { return called; }));
}

TEST(TestInputStreamLengthHelper, AsyncLengthStream)
{
  nsCOMPtr<nsIInputStream> stream =
      new LengthStream(22, NS_BASE_STREAM_WOULD_BLOCK, 123, true);

  bool called = false;
  InputStreamLengthHelper::GetAsyncLength(stream, [&](int64_t aLength) {
    ASSERT_EQ(22, aLength);
    called = true;
  });

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil(
      "xpcom:TEST(TestInputStreamLengthHelper, AsyncLengthStream)"_ns,
      [&]() { return called; }));
}

TEST(TestInputStreamLengthHelper, FallbackLengthStream)
{
  nsCOMPtr<nsIInputStream> stream =
      new LengthStream(-1, NS_BASE_STREAM_WOULD_BLOCK, 123, false);

  bool called = false;
  InputStreamLengthHelper::GetAsyncLength(stream, [&](int64_t aLength) {
    ASSERT_EQ(123, aLength);
    called = true;
  });

  MOZ_ALWAYS_TRUE(SpinEventLoopUntil(
      "xpcom:TEST(TestInputStreamLengthHelper, FallbackLengthStream)"_ns,
      [&]() { return called; }));
}