#pragma once

#include <chrono>

#include "envoy/buffer/buffer.h"

#include "envoy/common/pure.h"

#include "envoy/grpc/status.h"

#include "envoy/http/async_client.h"

#include "envoy/http/header_map.h"

#include "envoy/stream_info/stream_info.h"

#include "envoy/tracing/tracer.h"

#include "source/common/common/assert.h"

#include "source/common/protobuf/protobuf.h"

#include "absl/types/optional.h"

namespace Envoy {

namespace Grpc {

/**

 * An in-flight gRPC unary RPC.

 */

class AsyncRequest {

public:

  /**

   * Signals that the request should be cancelled. No further callbacks will be invoked.

   */

  virtual void cancel() PURE;

  /**

   * Returns the underlying stream info.

   */

  virtual const StreamInfo::StreamInfo& streamInfo() const PURE;

  /**

   * Detach the pending request. This is used for the case where we send a side

   * request but never cancel it even if the related downstream main request is

   * completed.

   *

   * This will will clean up all context associated with downstream request like

   * downstream stream info, parent tracing span, and so on, to avoid potential

   * dangling references.

   *

   * NOTE: the callbacks that registered to take the response will be kept to do

   * some clean up or operations when response arrives. The caller is responsible

   * for ensuring that the callbacks have enough lifetime.

   */

  virtual void detach() PURE;

};

/**

 * An in-flight gRPC stream.

 */

class RawAsyncStream {

public:

  /**

   * Send request message to the stream.

   * @param request serialized message.

   * @param end_stream close the stream locally. No further methods may be invoked on the stream

   *                   object, but callbacks may still be received until the stream is closed

   *                   remotely.

   */

  virtual void sendMessageRaw(Buffer::InstancePtr&& request, bool end_stream) PURE;

  /**

   * Close the stream locally and send an empty DATA frame to the remote. No further methods may be

   * invoked on the stream object, but callbacks may still be received until the stream is closed

   * remotely.

   */

  virtual void closeStream() PURE;

  /**

   * Close the stream locally and remotely (as needed). No further methods may be invoked on the

   * stream object and no further callbacks will be invoked.

   */

  virtual void resetStream() PURE;

  /**

   * Wait for the server to half-close its stream and then delete the RawAsyncStream object. No

   * further methods may be invoked on the stream object and no further callbacks will be invoked.

   * The server is expected to half-close within the interval specific in the StreamOptions,

   * otherwise the stream is reset.

   */

  virtual void waitForRemoteCloseAndDelete() PURE;

  /***

   * @returns if the stream has enough buffered outbound data to be over the configured buffer

   * limits

   */

  virtual bool isAboveWriteBufferHighWatermark() const PURE;

  /**

   * @returns the stream info object associated with this stream.

   */

  virtual const StreamInfo::StreamInfo& streamInfo() const PURE;

  virtual StreamInfo::StreamInfo& streamInfo() PURE;

  /***

   * Register a callback to be called when high/low write buffer watermark events occur on the

   * stream. This callback must persist beyond the lifetime of the stream or be unregistered via

   * removeWatermarkCallbacks. If there's already a watermark callback registered, this method

   * will trigger ENVOY_BUG.

   */

  virtual void setWatermarkCallbacks(Http::SidestreamWatermarkCallbacks& callbacks) PURE;

  /***

   * Remove previously set watermark callbacks. If there's no watermark callback registered, this

   * method will trigger ENVOY_BUG.

   */

  virtual void removeWatermarkCallbacks() PURE;

};

class RawAsyncRequestCallbacks {

public:

  /**

   * Called when populating the headers to send with initial metadata.

   * @param metadata initial metadata reference.

   */

  virtual void onCreateInitialMetadata(Http::RequestHeaderMap& metadata) PURE;

  /**

   * Called when the async gRPC request succeeds. No further callbacks will be invoked.

   * @param response the gRPC response bytes.

   * @param span a tracing span to fill with extra tags.

   */

  virtual void onSuccessRaw(Buffer::InstancePtr&& response, Tracing::Span& span) PURE;

  /**

   * Called when the async gRPC request fails. No further callbacks will be invoked.

   * @param status the gRPC status.

   * @param message the gRPC status message or empty string if not present.

   * @param span a tracing span to fill with extra tags.

   */

  virtual void onFailure(Status::GrpcStatus status, const std::string& message,

                         Tracing::Span& span) PURE;

};

/**

 * Notifies caller of async gRPC stream status.

 * Note the gRPC stream is full-duplex, even if the local to remote stream has been ended by

 * AsyncStream.close(), AsyncStreamCallbacks can continue to receive events until the remote

 * to local stream is closed (onRemoteClose), and vice versa. Once the stream is closed remotely, no

 * further callbacks will be invoked.

 */

class RawAsyncStreamCallbacks {

public:

  /**

   * Called when populating the headers to send with initial metadata.

   * @param metadata initial metadata reference.

   */

  virtual void onCreateInitialMetadata(Http::RequestHeaderMap& metadata) PURE;

  /**

   * Called when initial metadata is received. This will be called with empty metadata on a

   * trailers-only response, followed by onReceiveTrailingMetadata() with the trailing metadata.

   * @param metadata initial metadata reference.

   */

  virtual void onReceiveInitialMetadata(Http::ResponseHeaderMapPtr&& metadata) PURE;

  /**

   * Called when an async gRPC message is received.

   * @param response the gRPC message.

   * @return bool which is true if the message well formed and false otherwise which will cause

              the stream to shutdown with an INTERNAL error.

   */

  virtual bool onReceiveMessageRaw(Buffer::InstancePtr&& response) PURE;

  /**

   * Called when trailing metadata is received. This will also be called on non-Ok grpc-status

   * stream termination.

   * @param metadata trailing metadata reference.

   */

  virtual void onReceiveTrailingMetadata(Http::ResponseTrailerMapPtr&& metadata) PURE;

  /**

   * Called when the remote closes or an error occurs on the gRPC stream. The stream is

   * considered remotely closed after this invocation and no further callbacks will be

   * invoked. In addition, no further stream operations are permitted.

   * @param status the gRPC status.

   * @param message the gRPC status message or empty string if not present.

   */

  virtual void onRemoteClose(Status::GrpcStatus status, const std::string& message) PURE;

};

/**

 * Supports sending gRPC requests and receiving responses asynchronously. This can be used to

 * implement either plain gRPC or streaming gRPC calls.

 */

class RawAsyncClient {

public:

  /**

   * Start a gRPC unary RPC asynchronously.

   * @param service_full_name full name of the service (i.e. service_method.service()->full_name()).

   * @param method_name name of the method (i.e. service_method.name()).

   * @param request serialized message.

   * @param callbacks the callbacks to be notified of RPC status.

   * @param parent_span the current parent tracing context.

   * @param options the data struct to control the request sending.

   * @return a request handle or nullptr if no request could be started. NOTE: In this case

   *         onFailure() has already been called inline. The client owns the request and the

   *         handle should just be used to cancel.

   */

  virtual AsyncRequest* sendRaw(absl::string_view service_full_name, absl::string_view method_name,

                                Buffer::InstancePtr&& request, RawAsyncRequestCallbacks& callbacks,

                                Tracing::Span& parent_span,

                                const Http::AsyncClient::RequestOptions& options) PURE;

  /**

   * Start a gRPC stream asynchronously.

   * @param service_full_name full name of the service (i.e. service_method.service()->full_name()).

   * @param method_name name of the method (i.e. service_method.name()).

   * @param callbacks the callbacks to be notified of stream status.

   * @param options the data struct to control the stream.

   * @return a stream handle or nullptr if no stream could be started. NOTE: In this case

   *         onRemoteClose() has already been called inline. The client owns the stream and

   *         the handle can be used to send more messages or finish the stream. It is expected that

   *         closeStream() is invoked by the caller to notify the client that the stream resources

   *         may be reclaimed.

   */

  virtual RawAsyncStream* startRaw(absl::string_view service_full_name,

                                   absl::string_view method_name,

                                   RawAsyncStreamCallbacks& callbacks,

                                   const Http::AsyncClient::StreamOptions& options) PURE;

  /**

   * Returns the name of the cluster, or other destination/target, of the client.

   */

  virtual absl::string_view destination() PURE;

protected:

  // The lifetime of RawAsyncClient must be in the same thread.

  std::thread::id thread_id_{std::this_thread::get_id()};

};

using RawAsyncClientPtr = std::unique_ptr<RawAsyncClient>;

using RawAsyncClientSharedPtr = std::shared_ptr<RawAsyncClient>;

} // namespace Grpc

} // namespace Envoy