// Copyright 2018-2021 Bloomberg Finance L.P
//
// 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.

#ifndef INCLUDED_BUILDBOXCOMMON_REMOTEEXECUTIONCLIENT
#define INCLUDED_BUILDBOXCOMMON_REMOTEEXECUTIONCLIENT

#include <buildboxcommon_casclient.h>
#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_executionclient.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_logstreamreader.h>
#include <buildboxcommon_protos.h>

#include <atomic>
#include <map>
#include <mutex>
#include <optional>

namespace buildboxcommon {

typedef std::shared_ptr<grpc::ClientAsyncReaderInterface<Operation>>
    ReaderPointer;
typedef grpc::ServerWriterInterface<Operation> *WriterPointer;

typedef std::shared_ptr<google::longrunning::Operation> OperationPointer;

struct LogStreamInfo {
    bool completed = false;
    std::string stderrFileName;
    std::optional<std::string> stderrStreamName;
    std::ostream *stderrStreamTarget = nullptr;
    std::string stdoutFileName;
    std::optional<std::string> stdoutStreamName;
    std::ostream *stdoutStreamTarget = nullptr;
    std::map<std::string, bool> streamResults;
};

class RemoteExecutionClient : public ExecutionClient {
  private:
    std::shared_ptr<buildboxcommon::GrpcClient> d_executionGrpcClient;

    std::shared_ptr<Execution::StubInterface> d_executionStub;
    std::shared_ptr<Operations::StubInterface> d_operationsStub;

    std::shared_ptr<const buildboxcommon::ConnectionOptions>
        d_logstreamConnectionOptions;

    std::map<Digest, LogStreamInfo> d_logStreamInfo;
    std::mutex d_logStreamInfoMutex;

    bool d_logProgress = false;

    bool readOperation(grpc::ClientContext *context, grpc::CompletionQueue *cq,
                       ReaderPointer &reader, OperationPointer &operationPtr,
                       WriterPointer writer,
                       std::function<bool()> stopRequested,
                       const gpr_timespec &firstRequestDeadline,
                       grpc::Status *errorStatus, bool wait = true,
                       bool cancelOnStop = true,
                       const std::string &operationNamePrefix = "");

    OperationPointer submitExecution(const Digest &actionDigest,
                                     const std::atomic_bool &stopRequested,
                                     bool skipCache,
                                     const ExecutionPolicy *executionPolicy,
                                     bool wait);

    grpc::Status readOperationStream(
        grpc::ClientContext *context, ReaderPointer &readerPtr,
        WriterPointer writerPtr, OperationPointer &operationPtr,
        grpc::CompletionQueue *cq, gpr_timespec requestDeadline,
        std::function<bool()> stopRequested, bool wait, bool cancelOnStop,
        const std::string &operationNamePrefix = "");

    /**
     * Send an Execute request to the remote server, returning a pointer to
     * the last Operation message received in response.
     *
     * If `wait` is true, then this will be the final Operation response
     * assuming the request was completed, otherwise it will be the first
     * Operation received from the server.
     */
    OperationPointer performExecuteRequest(const ExecuteRequest &request,
                                           std::function<bool()> stopRequested,
                                           bool wait, bool cancelOnStop,
                                           grpc::Status *status);

    /**
     * Send an Execute request to the remote server, returning a pointer to
     * the last Operation message received in response.
     *
     * If `wait` is true, then this will be the final Operation response
     * assuming the request was completed, otherwise it will be the first
     * Operation received from the server.
     *
     * All Operation messages received using this method will be written
     * to the `grpc::ServerWriterInterface` pointed to by `writerPtr`.
     */
    OperationPointer performExecuteRequest(
        const ExecuteRequest &request, std::function<bool()> stopRequested,
        WriterPointer writerPtr, bool wait, bool cancelOnStop,
        grpc::Status *status, const std::string &operationNamePrefix = "");

    /**
     * Send a WaitExecution request to the remote server, returning a pointer
     * to the last Operation message received in response.
     *
     * If `wait` is true, then this will be the final Operation response
     * assuming the request was completed, otherwise it will be the first
     * Operation received from the server (making this method functionally
     * equivalent to `getOperation`).
     *
     * All Operation messages received using this method will be written
     * to the `grpc::ServerWriterInterface` pointed to by `writerPtr`.
     */
    OperationPointer
    performWaitExecutionRequest(const WaitExecutionRequest &request,
                                std::function<bool()> stopRequested,
                                WriterPointer writerPtr, bool wait,
                                bool cancelOnStop, grpc::Status *status,
                                const std::string &operationNamePrefix = "");

  public:
    explicit RemoteExecutionClient(
        std::shared_ptr<buildboxcommon::GrpcClient> executionGrpcClient,
        std::shared_ptr<buildboxcommon::GrpcClient> actionCacheGrpcClient,
        std::shared_ptr<const buildboxcommon::ConnectionOptions>
            logstreamConnectionOptions)
        : ExecutionClient(actionCacheGrpcClient),
          d_executionGrpcClient(executionGrpcClient),
          d_logstreamConnectionOptions(logstreamConnectionOptions)
    {
    }

    explicit RemoteExecutionClient(
        std::shared_ptr<buildboxcommon::GrpcClient> executionGrpcClient,
        std::shared_ptr<buildboxcommon::GrpcClient> actionCacheGrpcClient)
        : ExecutionClient(actionCacheGrpcClient),
          d_executionGrpcClient(executionGrpcClient),
          d_logstreamConnectionOptions(nullptr)
    {
    }

    ~RemoteExecutionClient() override {};

    RemoteExecutionClient(const RemoteExecutionClient &other) = delete;
    RemoteExecutionClient &
    operator=(const RemoteExecutionClient &other) = delete;

    RemoteExecutionClient(RemoteExecutionClient &&other) noexcept;

    RemoteExecutionClient &operator=(RemoteExecutionClient &&other) noexcept;

    virtual void
    init(std::shared_ptr<Execution::StubInterface> executionStub,
         std::shared_ptr<ActionCache::StubInterface> actionCacheStub,
         std::shared_ptr<Operations::StubInterface> operationsStub);

    void
    init(std::shared_ptr<ActionCache::StubInterface> actionCacheStub) override;

    void init() override;

    /**
     * Enable live log streaming support when executing Actions synchronously
     *
     * Callers of this method MUST also call `getLogStreamInfo` for the same
     * digest.
     *
     * Both forms take the Digest of the Action we want to enable LogStreaming
     * for.
     *
     * Plus either:
     *
     * Takes two pointers to ostreams to write the stdout and stderr streams of
     * the execution to, respectively. Or,
     *
     * Takes two strings containing filenames to stream stdout and stderr into,
     * respectively. If both are called, the filenames take precedence.
     */
    void enableLogStream(const Digest &digest, std::ostream *stdoutStream,
                         std::ostream *stderrStream);
    void enableLogStream(const Digest &digest, const std::string &stdoutFile,
                         const std::string &stderrFile);

    /**
     * Return an optional containing a `LogStreamInfo` for the given
     * ActionDigest
     */
    std::optional<LogStreamInfo> getLogStreamInfo(const Digest &digest);

    /**
     * Enables logging of partial execution metadata from read operations.
     */
    void enableProgressLog();

    /**
     * Run the action with the given digest on the given server, waiting
     * synchronously for it to complete. The Action must already be present in
     * the server's CAS.
     */
    ActionResult
    executeAction(const Digest &actionDigest,
                  const std::atomic_bool &stopRequested,
                  bool skipCache = false,
                  const ExecutionPolicy *executionPolicy = nullptr) override;

    /**
     * Run the action with the given digest on the given server asynchronously.
     * Returns the operation id. The Action must already be present in
     * the server's CAS.
     */
    google::longrunning::Operation asyncExecuteAction(
        const Digest &actionDigest, const std::atomic_bool &stopRequested,
        bool skipCache = false,
        const ExecutionPolicy *executionPolicy = nullptr) override;

    /*
     * Gets the current Operation object by operation id. Used to check its
     * execution status.
     */
    google::longrunning::Operation
    getOperation(const std::string &operationName) override;

    /**
     * Wait for the specified operation to complete. This will return the final
     * `Operation` message with `done` set to `true`, or throw an exception if
     * there was an unexpected error.
     */
    google::longrunning::Operation
    waitExecution(const std::string &operationName) override;

    /**
     * Sends the CancelOperation RPC
     */
    bool cancelOperation(const std::string &operationName) override;

    /**
     * Gets a list of Operations matching the given filter and page size.
     */
    google::longrunning::ListOperationsResponse
    listOperations(const std::string &name, const std::string &filter,
                   int page_size, const std::string &page_token) override;

    /**
     * Send an Execute request to the remote server.
     *
     * All Operation messages received using this method will be written
     * to the `grpc::ServerWriterInterface` pointed to by `writer_ptr`.
     */
    grpc::Status
    proxyExecuteRequest(const ExecuteRequest &request,
                        std::function<bool()> stopRequested,
                        WriterPointer writer_ptr,
                        const std::string &operation_name_prefix = "");

    /**
     * Send a WaitExecution request to the remote server.
     *
     * All Operation messages received using this method will be written
     * to the `grpc::ServerWriterInterface` pointed to by `writer_ptr`.
     */
    grpc::Status
    proxyWaitExecutionRequest(const WaitExecutionRequest &request,
                              std::function<bool()> stopRequested,
                              WriterPointer writer_ptr,
                              const std::string &operation_name_prefix = "");

    /**
     * Send a GetOperation request to the remote server.
     *
     * The Operation received in response to this request populates the
     * `response` parameter, and the gRPC status code received from the
     * server is returned.
     */
    grpc::Status proxyGetOperationRequest(const GetOperationRequest &request,
                                          Operation *response);

    /**
     * Send a ListOperations request to the remote server.
     *
     * The response from the server is stored in the `response` parameter,
     * and the gRPC status code received from the server is returned.
     */
    grpc::Status
    proxyListOperationsRequest(const ListOperationsRequest &request,
                               ListOperationsResponse *response);

    /**
     * Send a CancelOperation request to the remote server.
     *
     * The (empty) response from the server is stored in the `response`
     * parameter and the gRPC status code received from the server is
     * returned.
     */
    grpc::Status
    proxyCancelOperationRequest(const CancelOperationRequest &request,
                                google::protobuf::Empty *response);
};

// These utilities hide deprecation warnings: .proto files mark some operations
// as deprecated, this causes the generated code to have methods marked as
// deprecated, and as a result implementations that have to maintain backward
// compatibility are awash in build warnings.

Platform *commandMutablePlatformDeprecated(Command &cmd);
const Platform &commandPlatformDeprecated(const Command &cmd);

} // namespace buildboxcommon
#endif
