/*
 * Copyright 2018-2021 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_BUILDBOXCOMMON_CASCLIENT
#define INCLUDED_BUILDBOXCOMMON_CASCLIENT

#include <cstddef>
#include <cstdint>
#include <filesystem>
#include <functional>
#include <map>
#include <memory>
#include <span>
#include <ThreadPool.h>
#include <unistd.h>
#include <unordered_map>

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_merklize.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_requestmetadata.h>
#include <buildboxcommon_systemutils.h>

namespace buildboxcommon {

typedef std::function<grpc::Status(const ReadResponse &, size_t)>
    StreamConsumer;

/**
 * Tree of Directory objects supporting efficient lookup by Digest
 */
struct DirectoryTree {
    // Root Directory Digest
    Digest root;
    // Digest -> Directory map (ordered for deterministic Tree generation)
    std::map<Digest, Directory> map;

    void add(const Directory &directory)
    {
        this->map.emplace(DigestGenerator::hash(directory), directory);
    }
};

/**
 * Implements a mechanism to communicate with remote CAS servers, and includes
 * data members to keep track of an ongoing batch upload or batch download
 * request.
 */
class CASClient {
  public:
    CASClient(std::shared_ptr<GrpcClient> grpcClient)
        : d_grpcClient(grpcClient)
    {
    }

    virtual ~CASClient() = default;

    /**
     * Connect to the CAS server with the given GrpcClient.
     */
    virtual void init(bool checkCapabilities = true);

    /**
     * Connect to the CAS server with the given clients.
     */
    virtual void
    init(std::shared_ptr<ByteStream::StubInterface> bytestreamClient,
         std::shared_ptr<ContentAddressableStorage::StubInterface> casClient,
         std::shared_ptr<LocalContentAddressableStorage::StubInterface>
             d_localCasClient,
         std::shared_ptr<Capabilities::StubInterface> capabilitiesClient,
         size_t maxBatchTotalSizeBytes = 0);

    /**
     * Check if the client is connected to a local cas server, and if so,
     * get the server details.
     */
    std::optional<LocalServerDetails> localServerDetails();

    /**
     * Download the blob with the given digest and return it.
     *
     * If the server returned an error, or the size of the received blob does
     * not match the digest, throw an `std::runtime_error` exception.
     */
    virtual std::string
    fetchString(const Digest &digest,
                GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Download a stream with the given digest and consume each chunk.
     *
     * Different from `fetchString`, this function does not store or return
     * the whole string but a callback to consume each chunk should be
     * provided.
     */
    virtual void fetchStream(const Digest &digest,
                             const StreamConsumer &streamConsumer,
                             int64_t offset, std::optional<int64_t> limit,
                             GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Download the blob with the given digest to the given file descriptor.
     *
     * If the file descriptor cannot be written to, the size of the
     * received blob does not match the digest, or the server
     * returned an error, throw an `std::runtime_error` exception.
     */
    virtual void download(int fd, const Digest &digest,
                          GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Download a directory recursively from a root digest.
     */
    virtual void
    downloadDirectory(const Digest &digest, const std::string &path,
                      GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Download a directory from a single blob using a tree message digest.
     */
    virtual void
    downloadTree(const Digest &treeDigest, const std::string &downloadPath,
                 GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Upload the given string. If it can't be uploaded successfully, throw
     * an exception.
     */
    virtual void upload(const std::string &data, const Digest &digest,
                        GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Upload a blob from the given file descriptor. If it can't be uploaded
     * successfully, throw an exception.
     */
    virtual void upload(int fd, const Digest &digest,
                        GrpcClient::RequestStats *requestStats = nullptr);

    struct UploadRequest {
        Digest digest;
        std::string data;
        int dirfd;
        std::string path;
        int fd;

        UploadRequest(const Digest &_digest, const std::string _data)
            : digest(_digest), data(_data), dirfd(AT_FDCWD), fd(-1) {};

        static UploadRequest from_path(const Digest &_digest,
                                       const std::string _path)
        {
            auto request = UploadRequest(_digest);
            request.path = _path;
            return request;
        }

        static UploadRequest from_path(const Digest &_digest, int dirfd,
                                       const std::string _path)
        {
            auto request = UploadRequest(_digest);
            request.dirfd = dirfd;
            request.path = _path;
            return request;
        }

        static UploadRequest from_fd(const Digest &_digest, int fd)
        {
            auto request = UploadRequest(_digest);
            request.fd = fd;
            return request;
        }

      private:
        UploadRequest(const Digest &_digest)
            : digest(_digest), dirfd(AT_FDCWD), fd(-1) {};
    };

    struct UploadResult {
        Digest digest;
        grpc::Status status;

        UploadResult(const Digest &_digest, const grpc::Status &_status)
            : digest(_digest), status(_status) {};
    };

    /* Upload multiple digests in an efficient way, allowing each digest to
     * potentially fail separately.
     *
     * If `threadPool` is not nullptr, upload requests are parallelized with
     * it.
     *
     * Return a list containing the Digests that failed to be uploaded and the
     * errors they received. (An empty result indicates that all digests were
     * uploaded.)
     */
    virtual std::vector<UploadResult>
    uploadBlobs(const std::vector<UploadRequest> &requests,
                GrpcClient::RequestStats *requestStats = nullptr);
    virtual std::vector<UploadResult>
    uploadBlobs(const std::vector<UploadRequest> &requests,
                ThreadPool *threadPool,
                GrpcClient::RequestStats *requestStats = nullptr);

    std::shared_ptr<GrpcClient> getGrpcClient() const;
    typedef std::unordered_map<std::string,
                               std::pair<google::rpc::Status, std::string>>
        DownloadBlobsResult;

    /* Given a list of digests, download the data and return it in a map
     * indexed by hash. Allow each digest to potentially fail separately.
     *
     * The returned map's values are pairs of (status, data) where the second
     * component will be empty if the status contains an non-OK code.
     *
     * When downloading many and/or large blobs it is strongly recommended to
     * use `downloadBlobsToDirectory()` instead to avoid excessive memory
     * usage.
     *
     * If `threadPool` is not nullptr, the downloads are parallelized.
     */
    virtual DownloadBlobsResult
    downloadBlobs(const std::vector<Digest> &digests, ThreadPool *threadPool,
                  GrpcClient::RequestStats *requestStats = nullptr);

    /* Given a list of digests, download the data to a temporary directory
     * and return the file paths in a map indexed by hash. Allow each digest to
     * potentially fail separately. The specified directory must be empty.
     *
     * The returned map's values are pairs of (status, path) where the second
     * component will be empty if the status contains an non-OK code.
     *
     * This is suitable for downloading many and/or large blobs.
     */
    virtual DownloadBlobsResult
    downloadBlobsToDirectory(const std::vector<Digest> &digests,
                             int temp_dirfd, ThreadPool *threadPool = nullptr,
                             GrpcClient::RequestStats *requestStats = nullptr);

    /* Given a list of digests, download the data and store each blob in the
     * path specified by the entry's first member in the `outputs` map. The
     * resulting file will use the second member of the pair to set the file
     * permissions.
     *
     * If any errors are encountered in the process of fetching the blobs, it
     * aborts and throws an `std::runtime_error` exception. (It might leave
     * directories in an inconsistent state, i.e. with missing files.)
     */
    typedef std::unordered_multimap<std::string /*Digest*/,
                                    std::pair<std::string /*Path*/, FileNode>>
        OutputMap;
    virtual void
    downloadBlobs(const std::vector<Digest> &digests, const OutputMap &outputs,
                  GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Given a list of digests, creates and sends a `FindMissingBlobsRequest`
     * to the server.
     *
     * Returns a list of Digests that the remote server reports not having,
     * or throws a runtime_exception if the request failed.
     */
    virtual std::vector<Digest>
    findMissingBlobs(const std::vector<Digest> &digests,
                     GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Uploads the contents of the given path or directory file descriptor.
     *
     * Returns a list of Digests that the remote server reports not having,
     * or throws a runtime_exception if the request failed.
     *
     * If the optional `directory_digest` pointer is provided, the Digest of
     * the uploaded directory is copied to it. Likewise for the `tree`
     * pointer and the `Tree` object that is generated for the directory.
     */
    virtual std::vector<UploadResult>
    uploadDirectory(int dirfd, Digest *directory_digest = nullptr,
                    Tree *tree = nullptr,
                    GrpcClient::RequestStats *requestStats = nullptr);
    virtual std::vector<UploadResult>
    uploadDirectory(const std::string &path,
                    Digest *directory_digest = nullptr, Tree *tree = nullptr,
                    GrpcClient::RequestStats *requestStats = nullptr);

    /*
     * Send a LocalCas protocol `CaptureTree()` request containing the given
     * paths and the given property names. If successful, returns a
     * `CaptureTreeResponse` object (it contains a Status for each path).
     *
     * If the request fails, throws an `std::runtime_exception`.
     */
    virtual CaptureTreeResponse
    captureTree(const std::string &root, const std::vector<std::string> &paths,
                const std::vector<std::string> &properties,
                bool bypass_local_cache, mode_t unixModeMask,
                GrpcClient::RequestStats *requestStats = nullptr,
                const Command_OutputDirectoryFormat outputDirFormat =
                    Command_OutputDirectoryFormat_TREE_ONLY,
                bool allowChmodToRead = false) const;
    virtual CaptureTreeResponse
    captureTree(const std::vector<std::string> &paths,
                const std::vector<std::string> &properties,
                bool bypass_local_cache, mode_t unixModeMask,
                GrpcClient::RequestStats *requestStats = nullptr,
                const Command_OutputDirectoryFormat outputDirFormat =
                    Command_OutputDirectoryFormat_TREE_ONLY,
                bool allowChmodToRead = false) const;

    /*
     * Send a LocalCas protocol `CaptureFiles()` request containing the given
     * paths and the given property names. If successful, returns a
     * `CaptureFilesResponse` object (it contains a Status for each path).
     *
     * If the request fails, throws an `std::runtime_exception`.
     */
    virtual std::vector<CaptureFilesResponse_Response>
    captureFiles(const std::string &root,
                 const std::vector<std::string> &paths,
                 const std::vector<std::string> &properties,
                 bool bypass_local_cache, bool skipUpload = false,
                 GrpcClient::RequestStats *requestStats = nullptr) const;
    virtual std::vector<CaptureFilesResponse_Response>
    captureFiles(const std::vector<std::string> &paths,
                 const std::vector<std::string> &properties,
                 bool bypass_local_cache, bool skipUpload = false,
                 GrpcClient::RequestStats *requestStats = nullptr) const;

    virtual std::vector<HashFilesResponse_Response>
    hashFiles(const std::vector<std::string> &paths,
              const std::vector<std::string> &properties,
              GrpcClient::RequestStats *requestStats = nullptr) const;

    class StagedDirectory {
        /*
         * Represents a staged directory. It encapsulates the gRPC stream's
         * status, keeping it open (and preventing the server from cleaning
         * up).
         *
         * On destruction it sends an empty message to the server to clean up
         * the directory.
         */

      public:
        explicit StagedDirectory(
            std::shared_ptr<grpc::ClientContext> context,
            std::shared_ptr<grpc::ClientReaderWriterInterface<
                StageTreeRequest, StageTreeResponse>>
                reader_writer,
            const std::string &path);

        ~StagedDirectory();

        // Do now allow making copies:
        StagedDirectory(const StagedDirectory &) = delete;
        StagedDirectory &operator=(const StagedDirectory &) = delete;
        StagedDirectory(StagedDirectory &&) = delete;
        StagedDirectory &operator=(StagedDirectory &&) = delete;

        inline std::string path() const { return d_path; }

      private:
        const std::shared_ptr<grpc::ClientContext> d_context;
        const std::shared_ptr<grpc::ClientReaderWriterInterface<
            StageTreeRequest, StageTreeResponse>>
            d_reader_writer;
        const std::string d_path;
    };

    /**
     * Stage a directory using the LocalCAS `Stage()` call.
     *
     * The `path` parameter is optional. If not provided, the server will
     * assign a temporary directory.
     *
     * The `uid` and `gid` can be specified as credentials to LocalCAS.
     * If unspecified, the effective `uid` and `gid` are used.
     *
     * On success return a `unique_ptr` to a `StagedDirectory` object that when
     * destructed will request the server to clean up.
     *
     * On error throw an `std::runtime_error` exception.
     */
    virtual std::unique_ptr<StagedDirectory>
    stage(const Digest &root_digest,
          const StageTreeRequest_StagingMode stagingMode,
          const std::string &path = "",
          const ProcessCredentialsGetter &processCredentialsGetter =
              SystemUtils::getProcessCredentials,
          const std::vector<std::string> &postActionCommands = {},
          const std::optional<RemoteApisSocketConfig> &remoteApisSocket =
              std::nullopt) const;

    /**
     * Get the Directory tree whose root digest is 'digest', using the CAS
     * GetTree() call.
     *
     * On success, return a `DirectoryTree` object
     *
     * On error throw an 'std::runtime_error' exception
     */
    virtual DirectoryTree getTree(const Digest &digest);

    /**
     * Get the Directory tree from the digest of a tree message.
     *
     * If the directory tree is available as a tree message calling this is
     * faster than using the root digest with getTree().
     *
     * On error throw an 'std::runtime_error' exception
     */
    DirectoryTree
    getTreeFromTreeMessage(const Digest &tree_digest,
                           GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Get the Directory tree whose root digest is 'digest', using the CAS
     * GetTree() call, streaming the response.
     *
     * This will not raise an error if the returned tree is incomplete due to
     * missing Directory objects.
     */
    std::unique_ptr<grpc::ClientReaderInterface<GetTreeResponse>>
    getTree(grpc::ClientContext *context, const Digest &rootDigest,
            int pageSize = 0, const std::string &pageToken = "");

    /**
     * Issue a LocalCAS `FetchTree()` call and return the response.
     *
     * That call will fetch the entire directory tree rooted at a node from a
     * remote CAS to a local cache.
     *
     * The request is equivalent to `GetTree`, storing the `Directory`
     * objects in the local cache. If `fetch_file_blobs` is set, it will also
     * fetch all blobs referenced by the `Directory` objects (equivalent to
     * `FetchMissingBlobs`).
     *
     * If no remote CAS is configured in the LocalCAS server, it will check
     * presence of the entire directory tree (and optionally also file blobs)
     * in its local cache.
     *
     * On errors throws `GrpcError`.
     */
    virtual FetchTreeResponse
    fetchTree(const Digest &digest, const bool fetch_file_blobs,
              GrpcClient::RequestStats *requestStats = nullptr);

    /**
     * Fetch the Protocol Buffer message of the given type and digest and
     * deserialize it.
     */
    template <typename Msg>
    inline Msg fetchMessage(const Digest &digest,
                            GrpcClient::RequestStats *requestStats = nullptr)
    {
        Msg result;
        if (!result.ParseFromString(this->fetchString(digest, requestStats))) {
            throw std::runtime_error("Could not deserialize fetched message");
        }
        return result;
    }

    /**
     * Upload the given Protocol Buffer message to CAS and return its
     * Digest.
     */
    template <typename Msg>
    inline Digest
    uploadMessage(const Msg &msg,
                  GrpcClient::RequestStats *requestStats = nullptr)
    {
        const std::string str = msg.SerializeAsString();
        const Digest digest = DigestGenerator::hash(str);
        this->upload(str, digest, requestStats);
        return digest;
    }

    static size_t bytestreamChunkSizeBytes();

    /**
     * Export files to the specified directory.
     */
    ExportFilesResponse
    exportFiles(std::span<OutputFile> files, const std::filesystem::path &path,
                GrpcClient::RequestStats *requestStats = nullptr);

  protected:
    typedef std::function<void(const std::string &hash,
                               const std::string &data)>
        WriteBlobCallback;

    typedef std::pair<Digest, google::rpc::Status> DownloadResult;
    typedef std::vector<DownloadResult> DownloadResults;

    /* Download the digests in the specified list and invoke the
     * `write_blob_callback` function after each blob is downloaded.
     *
     * If `temp_dirfd` is specified, each blob is written into
     * a file in that directory and the filename is passed to
     * `write_blob_callback` instead of the blob data. The specified
     * directory must be empty.
     *
     * If `threadPool` is not nullptr, the downloads are parallelized.
     *
     * Note: marked as `protected` to unit-test.
     */
    DownloadResults
    downloadBlobs(const std::vector<Digest> &digests,
                  const WriteBlobCallback &write_blob_callback, int temp_dirfd,
                  ThreadPool *threadPool,
                  GrpcClient::RequestStats *requestStats = nullptr);

    typedef std::function<void(const std::vector<Digest> &digest,
                               const OutputMap &outputs)>
        download_callback_t;
    typedef std::function<Directory(const Digest &digest)>
        return_directory_callback_t;

    void
    downloadTree(const Digest &digest, const std::string &path,
                 const download_callback_t &download_callback,
                 const return_directory_callback_t &return_directory_callback,
                 GrpcClient::RequestStats *requestStats);

    void downloadDirectory(
        const Digest &digest, const std::string &path,
        const download_callback_t &download_callback,
        const return_directory_callback_t &return_directory_callback,
        GrpcClient::RequestStats *requestStats);

  private:
    std::shared_ptr<GrpcClient> d_grpcClient;
    std::shared_ptr<ByteStream::StubInterface> d_bytestreamClient;
    std::shared_ptr<ContentAddressableStorage::StubInterface> d_casClient;
    std::shared_ptr<LocalContentAddressableStorage::StubInterface>
        d_localCasClient;

    size_t d_maxBatchTotalSizeBytes = 0;

    std::string d_uuid;

    // Maximum number of bytes that can be sent in a single gRPC message.
    static const size_t s_bytestreamChunkSizeBytes;

    std::string makeResourceName(const Digest &digest, bool is_upload);

    /* Given a list of digests, download the data and return it in a map
     * indexed by hash. Allow each digest to potentially fail separately.
     *
     * The returned map's values are pairs of (status, data) where the second
     * component will be empty if the status contains an non-OK code.
     *
     * If `temp_dirfd` is specified, blobs are written to files in that
     * directory and the second component of the returned map's values will be
     * the path to the file. The specified directory must be empty.
     */
    DownloadBlobsResult
    downloadBlobs(const std::vector<Digest> &digests, int temp_dirfd,
                  ThreadPool *threadPool = nullptr,
                  GrpcClient::RequestStats *requestStats = nullptr);

    /* Uploads the requests contained in the range [start_index,
     * end_index).
     *
     * The sum of bytes of the data in this range MUST NOT exceed the
     * maximum batch size request allowed.
     */
    std::vector<UploadResult>
    batchUpload(const std::vector<UploadRequest> &requests,
                const size_t start_index, const size_t end_index,
                GrpcClient::RequestStats *requestStats);

    /* Downloads the data for the Digests stored in the range
     * [start_index, end_index) of the given vector.
     *
     * The sum of sizes inside the range MUST NOT exceed the maximum batch
     * size request allowed.
     *
     * The `write_blob_function` callback will be invoked for each blob that is
     * successfully fetched.
     *
     * Returns a map with the status code received for each Digest.
     */
    DownloadResults batchDownload(const std::vector<Digest> &digests,
                                  const size_t start_index,
                                  const size_t end_index,
                                  const WriteBlobCallback &write_blob_function,
                                  GrpcClient::RequestStats *requestStats,
                                  int temp_dirfd = -1);

    /* Given a list of digests sorted by increasing size, forms batches
     * according to the value of `d_maxBatchTotalSizeBytes`.
     */
    std::vector<std::pair<size_t, size_t>>
    makeBatches(const std::vector<Digest> &digests);

    /* Given a list of paths, forms CaptureFiles batches according to the
     * value of`d_maxBatchTotalSizeBytes`.
     */
    std::vector<std::pair<size_t, size_t>>
    makeCaptureBatches(const std::vector<std::string> &paths) const;

    /* Given a directory map, invoke `findMissingBlobs()` and return a map
     * with the subset of protos that need to be uploaded.
     */
    void missingDigests(digest_string_map *digest_path_map,
                        digest_string_map *digest_blob_map,
                        GrpcClient::RequestStats *requestStats);

    /* Upload a single request. */
    void uploadRequest(const UploadRequest &request,
                       GrpcClient::RequestStats *requestStats);
};

} // namespace buildboxcommon

#endif
