// Copyright 2018-2023 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_EXECUTIONCLIENT
#define INCLUDED_BUILDBOXCOMMON_EXECUTIONCLIENT

#include <buildboxcommon_casclient.h>
#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_protos.h>
#include <ThreadPool.h>

#include <atomic>
#include <map>
#include <set>

namespace buildboxcommon {

class ExecutionClient {
  private:
    std::shared_ptr<buildboxcommon::GrpcClient> d_actionCacheGrpcClient;

    std::shared_ptr<ActionCache::StubInterface> d_actionCacheStub;

    bool d_actionCacheUpdatesAllowed = true;

  public:
    explicit ExecutionClient(
        std::shared_ptr<buildboxcommon::GrpcClient> actionCacheGrpcClient)
        : d_actionCacheGrpcClient(actionCacheGrpcClient)
    {
    }

    virtual ~ExecutionClient() {};

    // Rule of Five: Delete copy constructor and copy assignment operator
    ExecutionClient(const ExecutionClient &) = delete;
    ExecutionClient &operator=(const ExecutionClient &) = delete;

    // Rule of Five: Define move constructor and move assignment operator
    ExecutionClient(ExecutionClient &&other) noexcept = default;
    ExecutionClient &operator=(ExecutionClient &&other) noexcept = default;

    /**
     * Disable action cache updates by the execution client. A remote execution
     * server may still cache the action result. To disable caching completely,
     * set the `Action.do_not_cache` field instead.
     */
    void disableActionCacheUpdates();

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

    virtual void init();

    /**
     * Attempts to fetch the ActionResult with the given digest from the action
     * cache and store it in the `result` parameter. The return value
     * indicates whether the ActionResult was found in the action cache.
     * If it wasn't, `result` is not modified.
     *
     */
    virtual bool fetchFromActionCache(const Digest &actionDigest,
                                      const std::set<std::string> &outputs,
                                      ActionResult *result);

    /**
     * Upload a new execution result. The Action message must be uploaded to
     * CAS before calling this method.
     */
    virtual void updateActionCache(const Digest &actionDigest,
                                   const ActionResult &result);

    /**
     * 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.
     */
    virtual ActionResult
    executeAction(const Digest &actionDigest,
                  const std::atomic_bool &stopRequested,
                  bool skipCache = false,
                  const ExecutionPolicy *executionPolicy = nullptr) = 0;

    /**
     * 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.
     */
    virtual google::longrunning::Operation
    asyncExecuteAction(const Digest &actionDigest,
                       const std::atomic_bool &stop_requested,
                       bool skipCache = false,
                       const ExecutionPolicy *executionPolicy = nullptr) = 0;

    /**
     * Download all output directories and files and store them in the
     * specified directory.
     *
     * Will throw an exception if one or multiple blobs cannot be downloaded.
     */
    virtual void downloadOutputs(buildboxcommon::CASClient *casClient,
                                 const ActionResult &actionResult, int dirfd,
                                 ThreadPool *threadPool);

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

    /**
     * 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.
     */
    virtual google::longrunning::Operation
    waitExecution(const std::string &operationName) = 0;

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

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

    std::shared_ptr<buildboxcommon::GrpcClient>
    getActionCacheGrpcClient() const;

    std::shared_ptr<ActionCache::StubInterface> getActionCacheStub() const;

    bool isActionCacheUpdatesAllowed() const;

  protected:
    /**
     * Return the ActionResult for the given Operation. Throws an exception
     * if the Operation finished with an error, or if the Operation hasn't
     * finished yet.
     */
    ActionResult getActionResult(const Operation &operation);
};
} // namespace buildboxcommon
#endif
