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

#include <chrono>
#include <memory>

#include <folly/io/IOBuf.h>
#include <folly/io/async/AsyncTransport.h>
#include <folly/io/async/EventBase.h>
#include <thrift/lib/cpp/protocol/TProtocolTypes.h>
#include <thrift/lib/cpp2/async/ClientChannel.h>
#include <thrift/lib/cpp2/transport/core/ClientConnectionIf.h>

namespace apache {
namespace thrift {

/**
 * This is the client side interface for Thrift RPCs.  You create an
 * object of this type and pass it as an argument to the constructor
 * of the client code generated by the Thrift compiler.
 *
 * ThriftClient objects are lightweight and you can create a new one
 * for each RPC.  However ThriftClient objects have minimal state and
 * multiple RPCs needing the same state may use the same ThriftClient
 * object at the same time.
 *
 * ThriftClient objects are provided a ClientConnectionIf object and
 * an event base as parameters during construction.
 *
 * The ClientConnectionIf object handles the underlying connection for
 * the RPCs.
 *
 * The event base is used to perform the callbacks for asynchronous
 * RPCs.  Callbacks for synchronous RPCs are always performed on the
 * event base of the underlying connection.
 *
 * The event base can be the same as that of the underlying connection
 * (and a special constructor is provided for this purpose).  Use this
 * only for synchronous RPCs.  While asynchronous RPCs will also work
 * in this case, we do not recommend using the event base of the
 * underling connection to perform callbacks for asynchronous RPCs.
 *
 * Synchronous RPCs may be performed from any thread except the one
 * that manages the underlying connection.
 *
 * Asynchronous RPCs may be performed from any thread.
 */
class ThriftClient : public ClientChannel {
 public:
  // Use "Ptr" instead of "unique_ptr<ThriftClient>".
  using Ptr =
      std::unique_ptr<ThriftClient, folly::DelayedDestruction::Destructor>;

  // Creates a ThriftClient object that uses "connection".  Callbacks
  // for asynchronous RPCs are run on "callbackEvb".  Callbacks for
  // synchronous RPCs are run on the event base of the connection.
  ThriftClient(
      std::shared_ptr<ClientConnectionIf> connection,
      folly::EventBase* callbackEvb);

  // Creates a ThriftClient object that uses "connection".  Callbacks
  // for all RPCs are run on the event base of the connection.
  explicit ThriftClient(std::shared_ptr<ClientConnectionIf> connection);

  ThriftClient(const ThriftClient&) = delete;
  ThriftClient& operator=(const ThriftClient&) = delete;

  void setProtocolId(uint16_t protocolId);
  void setHTTPHost(const std::string& host);
  void setHTTPUrl(const std::string& url);

  // begin RequestChannel methods
  void sendRequestResponse(
      const RpcOptions& rpcOptions,
      MethodMetadata&&,
      SerializedRequest&&,
      std::shared_ptr<apache::thrift::transport::THeader> header,
      RequestClientCallback::Ptr cb) override;

  void sendRequestNoResponse(
      const RpcOptions& rpcOptions,
      MethodMetadata&&,
      SerializedRequest&&,
      std::shared_ptr<apache::thrift::transport::THeader> header,
      RequestClientCallback::Ptr cb) override;

  void sendRequestStream(
      const RpcOptions&,
      MethodMetadata&&,
      SerializedRequest&&,
      std::shared_ptr<transport::THeader>,
      StreamClientCallback* clientCallback) override {
    clientCallback->onFirstResponseError(
        folly::make_exception_wrapper<transport::TTransportException>(
            "This channel doesn't support stream RPC"));
  }

  void sendRequestSink(
      const RpcOptions&,
      MethodMetadata&&,
      SerializedRequest&&,
      std::shared_ptr<transport::THeader>,
      SinkClientCallback* clientCallback) override {
    clientCallback->onFirstResponseError(
        folly::make_exception_wrapper<transport::TTransportException>(
            "This channel doesn't support sink RPC"));
  }

  folly::EventBase* getEventBase() const override;

  uint16_t getProtocolId() override;

  void setCloseCallback(CloseCallback* cb) override;

  // end RequestChannel methods

  // begin ClientChannel methods

  // These methods are delegated to the connection object.  Given that
  // connection objects may be shared by multiple ThriftClient
  // objects, calls to these methods will affect all these
  // ThriftClient objects.  Therefore, these methods should only be
  // called by frameworks that manage all the ThriftClient objects.
  //
  // TODO: Refactor this to be cleaner.

  folly::AsyncTransport* getTransport() override;
  bool good() override;
  SaturationStatus getSaturationStatus() override;
  void attachEventBase(folly::EventBase* eventBase) override;
  void detachEventBase() override;
  bool isDetachable() override;
  uint32_t getTimeout() override;
  void setTimeout(uint32_t ms) override;
  void closeNow() override;
  CLIENT_TYPE getClientType() override;

  // end ClientChannel methods

 protected:
  std::shared_ptr<ClientConnectionIf> connection_;
  folly::EventBase* callbackEvb_;
  std::string httpHost_;
  std::string httpUrl_;
  uint16_t protocolId_{apache::thrift::protocol::T_BINARY_PROTOCOL};

  // The default timeout for a Thrift RPC.
  static const std::chrono::milliseconds kDefaultRpcTimeout;

  // Destructor is private because this class inherits from
  // folly:DelayedDestruction.
  virtual ~ThriftClient();

  std::unique_ptr<ThriftChannelIf::RequestMetadata> createRequestMetadata(
      const RpcOptions& rpcOptions,
      RpcKind kind,
      apache::thrift::ProtocolId protocolId,
      transport::THeader* header);

  void sendRequestHelper(
      const RpcOptions& rpcOptions,
      RpcKind kind,
      std::unique_ptr<folly::IOBuf> buf,
      std::shared_ptr<apache::thrift::transport::THeader> header,
      RequestClientCallback::Ptr cb) noexcept;

  // Scheduled by sendRequestHelper in the connection thread.  Both
  // operations getChannel() and sendThriftRequest() can be scheduled
  // together to avoid scheduling two operations.  Also keeping these
  // two operations together is useful in HTTP2 because the stream id
  // is generated when sendThriftRequest is called, and we would like
  // the stream id to be increasing for more recently created
  // channels.
  static void getChannelAndSendThriftRequest(
      ClientConnectionIf* connection,
      ThriftChannelIf::RequestMetadata&& metadata,
      std::unique_ptr<folly::IOBuf> payload,
      std::unique_ptr<ThriftClientCallback> callback) noexcept;
};

} // namespace thrift
} // namespace apache