/*
 * Copyright 2019 Bloomberg Finance LP
 *
 * 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_BUILDBOXCASD_LOCALCASPROXYINSTANCE_H
#define INCLUDED_BUILDBOXCASD_LOCALCASPROXYINSTANCE_H

#include <buildboxcasd_digestcache.h>
#include <buildboxcasd_digestoperationqueue.h>
#include <buildboxcasd_expiringunorderedset.h>
#include <buildboxcasd_findmissingblobsclient.h>
#include <buildboxcasd_localcasinstance.h>
#include <buildboxcommon_casclient.h>
#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommonmetrics_countingmetricvalue.h>

#include <memory>
#include <ThreadPool.h>

using namespace build::bazel::remote::execution::v2;

using grpc::Status;

namespace buildboxcasd {

class
    LocalCasProxyInstance // NOLINT(cppcoreguidelines-special-member-functions)
        final : public LocalCasInstance {
    /* This class implements the logic for methods that service a CAS server
     * that also acts as a proxy pointing to a remote one.
     *
     * Inheriting from `LocalContentAddresableStorageInstance` avoids having
     * to reimplement behavior that is shared with that type of server.
     *
     * If `cache_findmissingblobs` is true, a caching mechanism will
     * attempt to reduce the number of `FindMissingBlobs()` requests issued for
     * a same digest after it was reported present by the remote.
     */
  public:
    typedef buildboxcommon::buildboxcommonmetrics::CountingMetricValue::Count
        CounterType;

    explicit LocalCasProxyInstance(
        const std::shared_ptr<buildboxcommon::LocalCas> &storage,
        const std::unordered_map<buildboxcommon::StageTreeRequest::StagingMode,
                                 std::shared_ptr<buildboxcommon::FileStager>>
            &fileStagers,
        const std::shared_ptr<ThreadPool> &digestThreadPool,
        const std::shared_ptr<ThreadPool> &ioThreadPool,
        const std::shared_ptr<std::unordered_set<std::string>>
            &allowedPreUnstageCommands,
        const buildboxcommon::ConnectionOptions &cas_endpoint,
        const std::string &instance_name, const bool read_only_remote = false,
        const int findmissingblobs_cache_ttl = 5 * 60);

    std::shared_ptr<CasInstance> clone(const std::string &instanceName,
                                       const std::string &root) override;
    LocalCasProxyInstance(const LocalCasProxyInstance &other,
                          const std::string &instanceName,
                          const std::string &root);

    grpc::Status FindMissingBlobs(const FindMissingBlobsRequest &request,
                                  FindMissingBlobsResponse *response) override;

    grpc::Status BatchUpdateBlobs(const BatchUpdateBlobsRequest &request,
                                  BatchUpdateBlobsResponse *response) override;

    grpc::Status BatchReadBlobs(const BatchReadBlobsRequest &request,
                                BatchReadBlobsResponse *response) override;

    grpc::Status Write(WriteRequest *request_message,
                       ServerReader<WriteRequest> &request,
                       WriteResponse *response,
                       Digest *requested_digest) override;

    grpc::Status GetTree(const GetTreeRequest &request,
                         ServerWriter<GetTreeResponse> *writer) override;

    // LocalCAS methods
    grpc::Status
    FetchMissingBlobs(const FetchMissingBlobsRequest &request,
                      FetchMissingBlobsResponse *response) override;

    grpc::Status
    UploadMissingBlobs(const UploadMissingBlobsRequest &request,
                       UploadMissingBlobsResponse *response) override;

    grpc::Status FetchTree(const FetchTreeRequest &request,
                           FetchTreeResponse *response) override;

    grpc::Status UploadTree(const UploadTreeRequest &request,
                            UploadTreeResponse *response) override;

    grpc::Status ExportFiles(const ExportFilesRequest &request,
                             ExportFilesResponse *response) override;

  protected:
    /*
     * Store digests locally, and upload to remote CAS.
     * Returns Tree digest for constructing the response.
     */
    Digest UploadAndStore(
        int dirfd, const buildboxcommon::digest_string_map &digest_blob_map,
        const buildboxcommon::digest_string_map &digest_path_map,
        const std::optional<Tree> &path_tree, const bool bypass_local_cache,
        const bool move_files_hint) override;

    void captureFileData(const Digest &digest, int fd, const std::string &path,
                         const bool bypass_local_cache,
                         const bool move_file) override;

    // Walks a tree and makes sure that all its blobs are stored locally.
    // If a CAS client is available, it will try and fetch missing blobs from a
    // remote.
    // It returns an `OK` status if all the blobs contained in the tree are
    // available locally, `FAILED_PRECONDITION` otherwise.
    Status prepareTreeForStaging(const Digest &root_digest) const override;

    bool hasRemote() const override;

  private:
    std::shared_ptr<buildboxcommon::GrpcClient> d_grpc_client;
    std::shared_ptr<buildboxcommon::CASClient> d_cas_client;
    bool d_read_only_remote = false;

    std::unordered_map<Digest, grpc::Status>
    batchUpdateRemoteCas(std::vector<buildboxcommon::CASClient::UploadRequest>
                             &upload_requests);

    grpc::Status
    findMissingBlobsInRemoteCas(const std::vector<Digest> &digests,
                                std::vector<Digest> *missing_digests);

    // In order to reduce queries for "popular" blobs, cache the responses
    // obtained from `FindMissingBlobs()`, assuming the blobs will remain
    // present in the remote for a while.
    std::shared_ptr<buildboxcasd::DigestCache> d_find_missing_blobs_cache;
    std::shared_ptr<buildboxcasd::FindMissingBlobsClient>
        d_find_missing_blobs_client;

    // Thread pool to parallelize requests to remote CAS
    std::shared_ptr<ThreadPool> d_ioThreadPool;

    // Fetch blobs from a remote CAS and store them locally.
    // Return status for each unsuccessfully downloaded blob.
    std::unordered_map<Digest, google::rpc::Status>
    fetchBlobs(const std::vector<Digest> &digests) const;

    // Fetches from a remote CAS the blobs in the given tree that are not
    // present locally.
    grpc::Status fetchTreeMissingBlobs(const Digest &root_digest,
                                       bool file_blobs = true) const;

    // This queue will make sure that `FetchTree()` requests do not run
    // simultaneously for the same digest in order to avoid redundant transfers
    // from the remote.
    std::shared_ptr<DigestOperationQueue> d_fetch_tree_operation_queue;

    Directory fetchDirectory(const Digest &root_digest) const;

    /* Returns whether the given digest is available to be read from the
     * local storage.
     */
    bool hasBlob(const Digest &digest) override;

    google::rpc::Status readBlob(const Digest &digest, std::string *data,
                                 size_t read_offset,
                                 size_t read_limit) override;

    google::rpc::Status writeBlob(const Digest &digest,
                                  const std::string &data) override;

    grpc::Status uploadBlob(const Digest &digest,
                            const std::string &path) const;

    void addToFindMissingBlobsCache(const Digest &digest);
};

} // namespace buildboxcasd

#endif // INCLUDED_BUILDBOXCASD_LOCALCASPROXYINSTANCE_H
