#pragma once

#include <memory>

#include "envoy/server/instance.h"

#include "source/common/buffer/buffer_impl.h"

#include "source/common/http/header_map_impl.h"

#include "absl/container/flat_hash_set.h"

#include "absl/synchronization/mutex.h"

namespace Envoy {

class AdminResponse;

// Holds context for a streaming response from the admin system, enabling

// flow-control into another system. This is particularly important when the

// generated response is very large, such that holding it in memory may cause

// fragmentation or out-of-memory failures. It is possible to interleave xDS

// response handling, overload management, and other admin requests during the

// streaming of a long admin response.

//

// There can be be multiple AdminResponses at a time; each are separately

// managed. However they will obtain their data from Envoy functions that

// run on the main thread.

//

// Responses may still be active after the server has shut down, and is no

// longer running its main thread dispatcher. In this state, the callbacks

// will be called with appropriate error codes.

//

// Requests can also be cancelled explicitly by calling cancel(). After

// cancel() is called, no further callbacks will be called by the response.

//

// The lifecycle of an AdminResponse is rendered as a finite state machine

// bubble diagram:

// https://docs.google.com/drawings/d/1njUl1twApEMoxmjaG4b7optTh5fcb_YNcfSnkHbdfq0/view

class AdminResponse : public std::enable_shared_from_this<AdminResponse> {

public:

  // AdminResponse can outlive MainCommonBase. But AdminResponse needs a

  // reliable way of knowing whether MainCommonBase is alive, so we do this with

  // PtrSet, which is held by MainCommonBase and all the active AdminResponses.

  // via shared_ptr. This gives MainCommonBase a reliable way of notifying all

  // active responses that it is being shut down, and thus all responses need to

  // be terminated. And it gives a reliable way for AdminResponse to detach

  // itself, whether or not MainCommonBase is already deleted.

  //

  // In summary:

  //  * MainCommonBase can outlive AdminResponse so we need detachResponse.

  //  * AdminResponse can outlive MainCommonBase, so we need shared_ptr.

  class PtrSet {

  public:

    /**

     * Called when an AdminResponse is created. When terminateAdminRequests is

     * called, all outstanding response objects have their terminate() methods

     * called.

     *

     * @param response the response pointer to be added to the set.

     */

    void attachResponse(AdminResponse* response);

    /**

     * Called when an AdminResponse is terminated, either by completing normally

     * or having the caller call cancel on it. Either way it needs to be removed

     * from the set that will be used by terminateAdminRequests below.

     *

     * @param response the response pointer to be removed from the set.

     */

    void detachResponse(AdminResponse* response);

    /**

     * Called after the server run-loop finishes; any outstanding streaming

     * admin requests will otherwise hang as the main-thread dispatcher loop

     * will no longer run.

     */

    void terminateAdminRequests();

    mutable absl::Mutex mutex_;

    absl::flat_hash_set<AdminResponse*> response_set_ ABSL_GUARDED_BY(mutex_);

    bool accepting_admin_requests_ ABSL_GUARDED_BY(mutex_) = true;

  };

  using SharedPtrSet = std::shared_ptr<PtrSet>;

  AdminResponse(Server::Instance& server, absl::string_view path, absl::string_view method,

                SharedPtrSet response_set);

  ~AdminResponse();

  /**

   * Requests the headers for the response. This can be called from any

   * thread, and HeaderFn may also be called from any thread.

   *

   * HeadersFn will not be called after cancel(). It is invalid to

   * to call nextChunk from within HeadersFn -- the caller must trigger

   * such a call on another thread, after HeadersFn returns. Calling

   * nextChunk from HeadersFn may deadlock.

   *

   * If the server is shut down during the operation, headersFn may

   * be called with a 503, if it has not already been called.

   *

   * @param fn The function to be called with the headers and status code.

   */

  using HeadersFn = std::function<void(Http::Code, Http::ResponseHeaderMap& map)>;

  void getHeaders(HeadersFn fn);

  /**

   * Requests a new chunk. This can be called from any thread, and the BodyFn

   * callback may also be called from any thread. BodyFn will be called in a

   * loop until the Buffer passed to it is fully drained. When 'false' is

   * passed as the second arg to BodyFn, that signifies the end of the

   * response, and nextChunk must not be called again.

   *

   * BodyFn will not be called after cancel(). It is invalid to

   * to call nextChunk from within BodyFn -- the caller must trigger

   * such a call on another thread, after BodyFn returns. Calling

   * nextChunk from BodyFn may deadlock.

   *

   * If the server is shut down during the operation, bodyFn will

   * be called with an empty body and 'false' for more_data, if

   * this has not already occurred.

   *

   * @param fn A function to be called on each chunk.

   */

  using BodyFn = std::function<void(Buffer::Instance&, bool)>;

  void nextChunk(BodyFn fn);

  /**

   * Requests that any outstanding callbacks be dropped. This can be called

   * when the context in which the request is made is destroyed. This enables

   * an application to implement a. The Response itself is held as a

   * shared_ptr as that makes it much easier to manage cancellation across

   * multiple threads.

   */

  void cancel();

  /**

   * @return whether the request was cancelled.

   */

  bool cancelled() const;

private:

  /**

   * Called when the server is terminated. This calls any outstanding

   * callbacks to be called. If nextChunk is called after termination,

   * its callback is called false for the second arg, indicating

   * end of stream.

   */

  void terminate();

  void requestHeaders();

  void requestNextChunk();

  void sendAbortChunkLockHeld() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  void sendErrorLockHeld() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

  Server::Instance& server_;

  OptRef<Server::Admin> opt_admin_;

  Buffer::OwnedImpl response_;

  Http::Code code_;

  Server::Admin::RequestPtr request_;

  Http::RequestHeaderMapPtr request_headers_{Http::RequestHeaderMapImpl::create()};

  Http::ResponseHeaderMapPtr response_headers_{Http::ResponseHeaderMapImpl::create()};

  bool more_data_ = true;

  // True if cancel() was explicitly called by the user; headers and body

  // callbacks are never called after cancel().

  bool cancelled_ ABSL_GUARDED_BY(mutex_) = false;

  // True if the Envoy server has stopped running its main loop. Headers and

  // body requests can be initiated and called back are called after terminate,

  // so callers do not have to special case this -- the request will simply fail

  // with an empty response.

  bool terminated_ ABSL_GUARDED_BY(mutex_) = false;

  // Used to indicate whether the body function has been called with false

  // as its second argument. That must always happen at most once, even

  // if terminate races with the normal end-of-stream marker. more=false

  // may never be sent if the request is cancelled, nor deleted prior to

  // it being requested.

  bool sent_end_stream_ ABSL_GUARDED_BY(mutex_) = false;

  HeadersFn headers_fn_ ABSL_GUARDED_BY(mutex_);

  BodyFn body_fn_ ABSL_GUARDED_BY(mutex_);

  mutable absl::Mutex mutex_;

  SharedPtrSet shared_response_set_;

};

using AdminResponseSharedPtr = std::shared_ptr<AdminResponse>;

} // namespace Envoy