/*
 * 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.
 */

#pragma once

#include <memory>

#include <folly/io/async/AsyncSocket.h>
#include <folly/io/async/test/BlockingSocket.h>
#include <folly/io/async/test/CallbackStateEnum.h>
#include <folly/io/async/test/ConnCallback.h>
#include <folly/net/NetOps.h>
#include <folly/net/NetworkSocket.h>
#include <folly/portability/Sockets.h>

namespace folly::test {

class WriteCallback : public folly::AsyncTransport::WriteCallback,
                      public folly::AsyncWriter::ReleaseIOBufCallback {
 public:
  explicit WriteCallback(bool enableReleaseIOBufCallback = false)
      : state(STATE_WAITING),
        bytesWritten(0),
        numIoBufCount(0),
        numIoBufBytes(0),
        exception(folly::AsyncSocketException::UNKNOWN, "none"),
        releaseIOBufCallback(enableReleaseIOBufCallback ? this : nullptr) {}

  void writeSuccess() noexcept override {
    state = STATE_SUCCEEDED;
    if (successCallback) {
      successCallback();
    }
  }

  void writeErr(
      size_t nBytesWritten,
      const folly::AsyncSocketException& ex) noexcept override {
    LOG(ERROR) << ex.what();
    state = STATE_FAILED;
    this->bytesWritten = nBytesWritten;
    exception = ex;
    if (errorCallback) {
      errorCallback();
    }
  }

  void writeStarting() noexcept override { writeStartingInvocations++; }

  folly::AsyncWriter::ReleaseIOBufCallback* getReleaseIOBufCallback() noexcept
      override {
    return releaseIOBufCallback;
  }

  void releaseIOBuf(std::unique_ptr<folly::IOBuf> ioBuf) noexcept override {
    numIoBufCount += ioBuf->countChainElements();
    numIoBufBytes += ioBuf->computeChainDataLength();
  }

  StateEnum state;
  std::atomic<size_t> bytesWritten;
  std::atomic<size_t> numIoBufCount;
  std::atomic<size_t> numIoBufBytes;
  folly::AsyncSocketException exception;
  VoidCallback successCallback;
  VoidCallback errorCallback;
  ReleaseIOBufCallback* releaseIOBufCallback;
  size_t writeStartingInvocations{0};
};

class ReadCallback : public folly::AsyncTransport::ReadCallback {
 public:
  explicit ReadCallback(size_t _maxBufferSz = 4096)
      : state(STATE_WAITING),
        exception(folly::AsyncSocketException::UNKNOWN, "none"),
        buffers(),
        maxBufferSz(_maxBufferSz) {}

  ~ReadCallback() override {
    for (auto& buffer : buffers) {
      buffer.free();
    }
    currentBuffer.free();
  }

  void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
    if (!currentBuffer.buffer) {
      currentBuffer.allocate(maxBufferSz);
    }
    *bufReturn = currentBuffer.buffer;
    *lenReturn = currentBuffer.length;
  }

  void readDataAvailable(size_t len) noexcept override {
    currentBuffer.length = len;
    buffers.push_back(currentBuffer);
    currentBuffer.reset();
    if (dataAvailableCallback) {
      dataAvailableCallback();
    }
  }

  void readEOF() noexcept override { state = STATE_SUCCEEDED; }

  void readErr(const folly::AsyncSocketException& ex) noexcept override {
    state = STATE_FAILED;
    exception = ex;
  }

  void verifyData(const char* expected, size_t expectedLen) const {
    verifyData((const unsigned char*)expected, expectedLen);
  }

  void verifyData(const unsigned char* expected, size_t expectedLen) const {
    size_t offset = 0;
    for (size_t idx = 0; idx < buffers.size(); ++idx) {
      const auto& buf = buffers[idx];
      size_t cmpLen = std::min(buf.length, expectedLen - offset);
      CHECK_EQ(memcmp(buf.buffer, expected + offset, cmpLen), 0);
      CHECK_EQ(cmpLen, buf.length);
      offset += cmpLen;
    }
    CHECK_EQ(offset, expectedLen);
  }

  void clearData() {
    for (auto& buffer : buffers) {
      buffer.free();
    }
    buffers.clear();
  }

  size_t dataRead() const {
    size_t ret = 0;
    for (const auto& buf : buffers) {
      ret += buf.length;
    }
    return ret;
  }

  class Buffer {
   public:
    Buffer() : buffer(nullptr), length(0) {}
    Buffer(char* buf, size_t len) : buffer(buf), length(len) {}

    void reset() {
      buffer = nullptr;
      length = 0;
    }
    void allocate(size_t len) {
      assert(buffer == nullptr);
      this->buffer = static_cast<char*>(malloc(len));
      this->length = len;
    }
    void free() {
      ::free(buffer);
      reset();
    }

    char* buffer;
    size_t length;
  };

  StateEnum state;
  folly::AsyncSocketException exception;
  std::vector<Buffer> buffers;
  Buffer currentBuffer;
  VoidCallback dataAvailableCallback;
  const size_t maxBufferSz;
};

class ReadvCallback : public folly::AsyncTransport::ReadCallback {
 public:
  ReadvCallback(size_t bufferSize, size_t len)
      : state_(STATE_WAITING),
        exception_(folly::AsyncSocketException::UNKNOWN, "none"),
        queue_(folly::IOBufIovecBuilder::Options().setBlockSize(bufferSize)),
        len_(len) {
    setReadMode(folly::AsyncTransport::ReadCallback::ReadMode::ReadVec);
  }

  ~ReadvCallback() override = default;

  void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
    std::ignore = bufReturn;
    std::ignore = lenReturn;

    CHECK(false); // this should not be called
  }

  void getReadBuffers(folly::IOBufIovecBuilder::IoVecVec& iovs) override {
    queue_.allocateBuffers(iovs, len_);
  }

  void readDataAvailable(size_t len) noexcept override {
    auto tmp = queue_.extractIOBufChain(len);
    if (!buf_) {
      buf_ = std::move(tmp);
    } else {
      buf_->prependChain(std::move(tmp));
    }
  }

  void reset() { buf_.reset(); }

  void readEOF() noexcept override { state_ = STATE_SUCCEEDED; }

  void readErr(const folly::AsyncSocketException& ex) noexcept override {
    state_ = STATE_FAILED;
    exception_ = ex;
  }

  void verifyData(const std::string& data) const {
    CHECK(buf_);
    auto r = buf_->coalesce();
    std::string tmp;
    tmp.assign(reinterpret_cast<const char*>(r.begin()), r.end() - r.begin());
    CHECK_EQ(data, tmp);
  }

  std::unique_ptr<folly::IOBuf> buf_;

 private:
  StateEnum state_;
  folly::AsyncSocketException exception_;
  folly::IOBufIovecBuilder queue_;
  const size_t len_;
};

class BufferCallback : public folly::AsyncTransport::BufferCallback {
 public:
  BufferCallback(folly::AsyncSocket* socket, size_t expectedBytes)
      : socket_(socket),
        expectedBytes_(expectedBytes),
        buffered_(false),
        bufferCleared_(false) {}

  void onEgressBuffered() override {
    size_t bytesWritten = socket_->getAppBytesWritten();
    size_t bytesBuffered = socket_->getAppBytesBuffered();
    CHECK_GT(bytesBuffered, 0);
    CHECK_EQ(expectedBytes_, bytesWritten + bytesBuffered);
    buffered_ = true;
  }

  void onEgressBufferCleared() override {
    size_t bytesWritten = socket_->getAppBytesWritten();
    size_t bytesBuffered = socket_->getAppBytesBuffered();
    CHECK_EQ(0, bytesBuffered);
    CHECK_EQ(expectedBytes_, bytesWritten);
    bufferCleared_ = true;
  }

  bool hasBuffered() const { return buffered_; }

  bool hasBufferCleared() const { return bufferCleared_; }

 private:
  folly::AsyncSocket* socket_{nullptr};
  size_t expectedBytes_{0};
  bool buffered_{false};
  bool bufferCleared_{false};
};

class ZeroCopyReadCallback : public folly::AsyncTransport::ReadCallback {
 public:
  explicit ZeroCopyReadCallback(
      folly::AsyncTransport::ReadCallback::ZeroCopyMemStore* memStore,
      size_t _maxBufferSz = 4096)
      : memStore_(memStore),
        state(STATE_WAITING),
        exception(folly::AsyncSocketException::UNKNOWN, "none"),
        maxBufferSz(_maxBufferSz) {}

  ~ZeroCopyReadCallback() override { currentBuffer.free(); }

  // zerocopy
  folly::AsyncTransport::ReadCallback::ZeroCopyMemStore*
  readZeroCopyEnabled() noexcept override {
    return memStore_;
  }

  void getZeroCopyFallbackBuffer(
      void** bufReturn, size_t* lenReturn) noexcept override {
    if (!currentZeroCopyBuffer.buffer) {
      currentZeroCopyBuffer.allocate(maxBufferSz);
    }
    *bufReturn = currentZeroCopyBuffer.buffer;
    *lenReturn = currentZeroCopyBuffer.length;
  }

  void readZeroCopyDataAvailable(
      std::unique_ptr<folly::IOBuf>&& zeroCopyData,
      size_t additionalBytes) noexcept override {
    auto ioBuf = std::move(zeroCopyData);
    if (additionalBytes) {
      auto tmp = folly::IOBuf::takeOwnership(
          currentZeroCopyBuffer.buffer,
          currentZeroCopyBuffer.length,
          0,
          additionalBytes);
      currentZeroCopyBuffer.reset();
      if (ioBuf) {
        ioBuf->prependChain(std::move(tmp));
      } else {
        ioBuf = std::move(tmp);
      }
    }

    if (!data_) {
      data_ = std::move(ioBuf);
    } else {
      data_->prependChain(std::move(ioBuf));
    }
  }

  void getReadBuffer(void** bufReturn, size_t* lenReturn) override {
    if (!currentBuffer.buffer) {
      currentBuffer.allocate(maxBufferSz);
    }
    *bufReturn = currentBuffer.buffer;
    *lenReturn = currentBuffer.length;
  }

  void readDataAvailable(size_t len) noexcept override {
    auto ioBuf = folly::IOBuf::takeOwnership(
        currentBuffer.buffer, currentBuffer.length, 0, len);
    currentBuffer.reset();

    if (!data_) {
      data_ = std::move(ioBuf);
    } else {
      data_->prependChain(std::move(ioBuf));
    }
  }

  void readEOF() noexcept override { state = STATE_SUCCEEDED; }

  void readErr(const folly::AsyncSocketException& ex) noexcept override {
    state = STATE_FAILED;
    exception = ex;
  }

  void verifyData(const std::string& expected) const {
    verifyData((const unsigned char*)expected.data(), expected.size());
  }

  void verifyData(const unsigned char* expected, size_t expectedLen) const {
    CHECK(!!data_);
    auto len = data_->computeChainDataLength();
    CHECK_EQ(len, expectedLen);

    auto* buf = data_.get();
    auto* current = buf;
    size_t offset = 0;

    do {
      size_t cmpLen = std::min(current->length(), expectedLen - offset);
      CHECK_EQ(cmpLen, current->length());
      CHECK_EQ(memcmp(current->data(), expected + offset, cmpLen), 0);
      offset += cmpLen;

      current = current->next();
    } while (current != buf);

    std::ignore = expected;
    CHECK_EQ(offset, expectedLen);
  }

  class Buffer {
   public:
    Buffer() = default;
    Buffer(char* buf, size_t len) : buffer(buf), length(len) {}
    ~Buffer() {
      if (buffer) {
        ::free(buffer);
      }
    }

    void reset() {
      buffer = nullptr;
      length = 0;
    }
    void allocate(size_t len) {
      CHECK(buffer == nullptr);
      buffer = static_cast<char*>(malloc(len));
      length = len;
    }
    void free() {
      ::free(buffer);
      reset();
    }

    char* buffer{nullptr};
    size_t length{0};
  };
  folly::AsyncTransport::ReadCallback::ZeroCopyMemStore* memStore_;
  StateEnum state;
  folly::AsyncSocketException exception;
  Buffer currentBuffer, currentZeroCopyBuffer;
  VoidCallback dataAvailableCallback;
  const size_t maxBufferSz;
  std::unique_ptr<folly::IOBuf> data_;
};

class ReadVerifier {};

class TestSendMsgParamsCallback
    : public folly::AsyncSocket::SendMsgParamsCallback {
 public:
  TestSendMsgParamsCallback(int flags, uint32_t dataSize, void* data)
      : flags_(flags),
        writeFlags_(folly::WriteFlags::NONE),
        dataSize_(dataSize),
        data_(data),
        queriedFlags_(false),
        queriedData_(false) {}

  void reset(int flags) {
    flags_ = flags;
    writeFlags_ = folly::WriteFlags::NONE;
    queriedFlags_ = false;
    queriedData_ = false;
  }

  int getFlagsImpl(
      folly::WriteFlags flags, int /*defaultFlags*/) noexcept override {
    queriedFlags_ = true;
    if (writeFlags_ == folly::WriteFlags::NONE) {
      writeFlags_ = flags;
    } else {
      assert(flags == writeFlags_);
    }
    return flags_;
  }

  void getAncillaryData(
      folly::WriteFlags flags,
      void* data,
      const folly::AsyncSocket::WriteRequestTag& tag,
      const bool /* byteEventsEnabled */) noexcept override {
    CHECK_EQ(tag, expectedTag_);
    queriedData_ = true;
    if (writeFlags_ == folly::WriteFlags::NONE) {
      writeFlags_ = flags;
    } else {
      assert(flags == writeFlags_);
    }
    assert(data != nullptr);
    memcpy(data, data_, dataSize_);
  }

  uint32_t getAncillaryDataSize(
      folly::WriteFlags flags,
      const folly::AsyncSocket::WriteRequestTag& tag,
      const bool /* byteEventsEnabled */) noexcept override {
    CHECK_EQ(tag, expectedTag_);
    if (writeFlags_ == folly::WriteFlags::NONE) {
      writeFlags_ = flags;
    } else {
      assert(flags == writeFlags_);
    }
    return dataSize_;
  }

  void wroteBytes(
      const folly::AsyncSocket::WriteRequestTag& tag) noexcept override {
    CHECK_EQ(tag, expectedTag_);
    tagLastWritten_ = tag;
  }

  int flags_;
  folly::WriteFlags writeFlags_;
  uint32_t dataSize_;
  void* data_;
  bool queriedFlags_;
  bool queriedData_;
  folly::AsyncSocket::WriteRequestTag expectedTag_{
      folly::AsyncSocket::WriteRequestTag::EmptyDummy()};
  std::optional<folly::AsyncSocket::WriteRequestTag> tagLastWritten_;
};

class TestServer {
 public:
  // Create a TestServer.
  // This immediately starts listening on an ephemeral port.
  explicit TestServer(bool enableTFO = false, int bufSize = -1) : fd_() {
    namespace fsp = folly::portability::sockets;
    fd_ = folly::netops::socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (fd_ == folly::NetworkSocket()) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "failed to create test server socket",
          errno);
    }
    if (folly::netops::set_socket_non_blocking(fd_) != 0) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "failed to put test server socket in "
          "non-blocking mode",
          errno);
    }
    if (enableTFO) {
#if FOLLY_ALLOW_TFO
      folly::detail::tfo_enable(fd_, 100);
#endif
    }

    struct addrinfo hints, *res;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_flags = AI_PASSIVE;

    if (getaddrinfo(nullptr, "0", &hints, &res)) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "Attempted to bind address to socket with "
          "bad getaddrinfo",
          errno);
    }

    SCOPE_EXIT {
      freeaddrinfo(res);
    };

    if (bufSize > 0) {
      folly::netops::setsockopt(
          fd_, SOL_SOCKET, SO_SNDBUF, &bufSize, sizeof(bufSize));
      folly::netops::setsockopt(
          fd_, SOL_SOCKET, SO_RCVBUF, &bufSize, sizeof(bufSize));
    }

    if (folly::netops::bind(fd_, res->ai_addr, res->ai_addrlen)) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "failed to bind to async server socket for port 10",
          errno);
    }

    if (folly::netops::listen(fd_, 10) != 0) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "failed to listen on test server socket",
          errno);
    }

    address_.setFromLocalAddress(fd_);
    // The local address will contain 0.0.0.0.
    // Change it to 127.0.0.1, so it can be used to connect to the server
    address_.setFromIpPort("127.0.0.1", address_.getPort());
  }

  ~TestServer() {
    if (fd_ != folly::NetworkSocket()) {
      folly::netops::close(fd_);
    }
  }

  // Get the address for connecting to the server
  const folly::SocketAddress& getAddress() const { return address_; }

  folly::NetworkSocket acceptFD(int timeout = 50) {
    folly::netops::PollDescriptor pfd;
    pfd.fd = fd_;
    pfd.events = POLLIN;
    int ret = folly::netops::poll(&pfd, 1, timeout);
    if (ret == 0) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "test server accept() timed out");
    } else if (ret < 0) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "test server accept() poll failed",
          errno);
    }

    auto acceptedFd = folly::netops::accept(fd_, nullptr, nullptr);
    if (acceptedFd == folly::NetworkSocket()) {
      throw folly::AsyncSocketException(
          folly::AsyncSocketException::INTERNAL_ERROR,
          "test server accept() failed",
          errno);
    }

    return acceptedFd;
  }

  std::shared_ptr<BlockingSocket> accept(int timeout = 50) {
    auto fd = acceptFD(timeout);
    return std::make_shared<BlockingSocket>(fd);
  }

  std::shared_ptr<folly::AsyncSocket> acceptAsync(
      folly::EventBase* evb, int timeout = 50) {
    auto fd = acceptFD(timeout);
    return folly::AsyncSocket::newSocket(evb, fd);
  }

  /**
   * Accept a connection, read data from it, and verify that it matches the
   * data in the specified buffer.
   */
  void verifyConnection(const char* buf, size_t len) {
    // accept a connection
    std::shared_ptr<BlockingSocket> acceptedSocket = accept();
    // read the data and compare it to the specified buffer
    std::unique_ptr<uint8_t[]> readbuf(new uint8_t[len]);
    acceptedSocket->readAll(readbuf.get(), len);
    CHECK_EQ(memcmp(buf, readbuf.get(), len), 0);
    // make sure we get EOF next
    uint32_t bytesRead = acceptedSocket->read(readbuf.get(), len);
    CHECK_EQ(bytesRead, 0);
  }

 private:
  folly::NetworkSocket fd_;
  folly::SocketAddress address_;
};

} // namespace folly::test
