/*
 * 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.
 */

#include <buildboxcommon_casclient.h>
#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_futuregroup.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_mergeutil.h>
#include <buildboxcommon_merklize.h>
#include <buildboxcommon_permissions.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommon_tempconstants.h>
#include <buildboxcommon_timeutils.h>

#include <algorithm>
#include <cerrno>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <fcntl.h>
#include <future>
#include <grpc/grpc.h>
#include <grpcpp/support/status.h>
#include <iterator>
#include <mutex>
#include <optional>
#include <sstream>
#include <stdexcept>
#include <string>
#include <sys/stat.h>
#include <ThreadPool.h>
#include <tuple>
#include <unistd.h>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

namespace buildboxcommon {

const size_t CASClient::s_bytestreamChunkSizeBytes = 1024UL * 1024UL;
// The default limit for gRPC messages is 4 MiB.
// Limit payload to 1 MiB to leave sufficient headroom for metadata.

#ifdef __SANITIZE_ADDRESS__
// Use protobuf reflection to work around ASAN container-overflow false
// positive with protobuf repeated numeric/enum fields.
std::unordered_set<int>
getDigestFunctions(const CacheCapabilities &cacheCapabilities)
{
    const auto nDigestFunctions = cacheCapabilities.digest_functions_size();
    std::unordered_set<int> digestFunctions;

    const auto reflection = cacheCapabilities.GetReflection();
    const auto fieldDescriptor =
        cacheCapabilities.GetDescriptor()->FindFieldByName("digest_functions");
    for (int i = 0; i < nDigestFunctions; i++) {
        digestFunctions.insert(reflection->GetRepeatedEnumValue(
            cacheCapabilities, fieldDescriptor, i));
    }

    return digestFunctions;
}
#else
std::unordered_set<int>
getDigestFunctions(const CacheCapabilities &cacheCapabilities)
{
    const auto &digestFunctions = cacheCapabilities.digest_functions();
    return std::unordered_set<int>(digestFunctions.cbegin(),
                                   digestFunctions.cend());
}
#endif

void CASClient::init(bool checkCapabilities)
{
    auto channel = d_grpcClient->channel();

    std::shared_ptr<ByteStream::Stub> bytestreamClient =
        ByteStream::NewStub(channel);
    std::shared_ptr<ContentAddressableStorage::Stub> casClient =
        ContentAddressableStorage::NewStub(channel);
    std::shared_ptr<Capabilities::Stub> capabilitiesClient;
    if (checkCapabilities) {
        capabilitiesClient = Capabilities::NewStub(channel);
    }
    std::shared_ptr<LocalContentAddressableStorage::StubInterface>
        localCasClient = LocalContentAddressableStorage::NewStub(channel);
    init(bytestreamClient, casClient, localCasClient, capabilitiesClient);
}

void CASClient::init(
    std::shared_ptr<ByteStream::StubInterface> bytestreamClient,
    std::shared_ptr<ContentAddressableStorage::StubInterface> casClient,
    std::shared_ptr<LocalContentAddressableStorage::StubInterface>
        localCasClient,
    std::shared_ptr<Capabilities::StubInterface> capabilitiesClient,
    size_t maxBatchTotalSizeBytes)
{
    this->d_bytestreamClient = std::move(bytestreamClient);
    this->d_casClient = std::move(casClient);
    this->d_localCasClient = std::move(localCasClient);

    if (maxBatchTotalSizeBytes == 0) {
        this->d_maxBatchTotalSizeBytes = d_grpcClient->maxMessageSizeBytes();
    }
    else {
        this->d_maxBatchTotalSizeBytes = maxBatchTotalSizeBytes;
    }

    BUILDBOX_LOG_DEBUG("Setting d_maxBatchTotalSizeBytes = "
                       << d_maxBatchTotalSizeBytes << " bytes by default");

    // Request server capabilities and adjust our defaults according to the
    // server response
    auto getCapabilitiesLambda = [&](grpc::ClientContext &context) {
        GetCapabilitiesRequest request;
        request.set_instance_name(d_grpcClient->instanceName());

        ServerCapabilities response;
        auto status =
            capabilitiesClient->GetCapabilities(&context, request, &response);
        if (status.ok()) {
            const size_t serverMaxBatchTotalSizeBytes =
                response.cache_capabilities().max_batch_total_size_bytes();
            // 0 means no server limit
            if (serverMaxBatchTotalSizeBytes > 0 &&
                serverMaxBatchTotalSizeBytes <
                    this->d_maxBatchTotalSizeBytes) {
                BUILDBOX_LOG_INFO(
                    "Reconfiguring d_maxBatchTotalSizeBytes down from "
                    << d_maxBatchTotalSizeBytes << " to "
                    << serverMaxBatchTotalSizeBytes
                    << " due to server max_batch_total_size_bytes of "
                    << serverMaxBatchTotalSizeBytes);
                this->d_maxBatchTotalSizeBytes = serverMaxBatchTotalSizeBytes;
            }

            // Verify server supports configured digest function.
            const auto configured_digest_function =
                DigestGenerator::digestFunction();
            const auto supportedFunctions =
                getDigestFunctions(response.cache_capabilities());
            if (!supportedFunctions.contains(configured_digest_function)) {
                const std::string error_message =
                    "CAS server does not support the configured digest "
                    "function: " +
                    DigestFunction_Value_Name(configured_digest_function);

                BUILDBOX_LOG_ERROR(error_message);
                throw std::runtime_error(error_message);
            }
        }
        return status;
    };

    if (capabilitiesClient) {
        auto retrier = d_grpcClient->makeRetrier(getCapabilitiesLambda,
                                                 "GetCapabilities()");
        if (!retrier.issueRequest() ||
            (!retrier.status().ok() && retrier.status().error_code() !=
                                           grpc::StatusCode::UNIMPLEMENTED)) {
            GrpcError::throwGrpcError(retrier.status());
        }
        else if (retrier.status().error_code() ==
                 grpc::StatusCode::UNIMPLEMENTED) {
            BUILDBOX_LOG_DEBUG(
                "Get capabilities request failed. Using default. "
                << retrier.status().error_message());
        }
    }

    // Generate UUID to use for uploads
    this->d_uuid = StringUtils::getUUIDString();
}

std::shared_ptr<GrpcClient> CASClient::getGrpcClient() const
{
    return d_grpcClient;
}

size_t CASClient::bytestreamChunkSizeBytes()
{
    return s_bytestreamChunkSizeBytes;
}

std::optional<LocalServerDetails> CASClient::localServerDetails()
{
    std::optional<LocalServerDetails> result{};
    auto caller = [&](grpc::ClientContext &context) {
        GetLocalServerDetailsRequest request;
        request.set_instance_name(d_grpcClient->instanceName());

        LocalServerDetails response;
        auto status = d_localCasClient->GetLocalServerDetails(
            &context, request, &response);

        if (status.ok()) {
            result.emplace(response);
        }
        else if (status.error_code() == grpc::StatusCode::UNIMPLEMENTED) {
            // Server is not local, we can stop.
            return grpc::Status::OK;
        }
        // Other errors, lets get the retrier to help us.
        return status;
    };

    auto retrier =
        d_grpcClient->makeRetrier(caller, "GetLocalServerDetails()");
    if (!retrier.issueRequest() || !retrier.status().ok()) {
        GrpcError::throwGrpcError(retrier.status());
    }
    return result;
}

std::string CASClient::makeResourceName(const Digest &digest, bool isUpload)
{
    std::string resourceName;

    if (!d_grpcClient->instanceName().empty()) {
        resourceName.append(d_grpcClient->instanceName());
        resourceName.append("/");
    }

    if (isUpload) {
        resourceName.append("uploads/");
        resourceName.append(this->d_uuid);
        resourceName.append("/");
    }

    resourceName.append("blobs/");
    resourceName.append(digest.hash());
    resourceName.append("/");
    resourceName.append(std::to_string(digest.size_bytes()));
    return resourceName;
}

std::string CASClient::fetchString(const Digest &digest,
                                   GrpcClient::RequestStats *requestStats)
{
    BUILDBOX_LOG_TRACE("Downloading " << digest.hash() << " to string");
    const std::string resourceName = this->makeResourceName(digest, false);

    std::string downloaded_data;
    auto streamConsumer = [&](const ReadResponse &readResponse,
                              size_t startOffset) {
        // Retry safety
        if (startOffset == 0) {
            downloaded_data.clear();
            downloaded_data.reserve(digest.size_bytes());
        }
        downloaded_data += readResponse.data();
        return grpc::Status::OK;
    };
    this->fetchStream(digest, streamConsumer, 0, std::nullopt, requestStats);

    // Post-download check
    const auto bytes_downloaded =
        static_cast<google::protobuf::int64>(downloaded_data.size());
    if (bytes_downloaded != digest.size_bytes()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Expected "
                                    << digest.size_bytes()
                                    << " bytes, but downloaded blob was "
                                    << bytes_downloaded << " bytes");
    }
    const auto downloaded_digest = DigestGenerator::hash(downloaded_data);
    if (downloaded_digest != digest) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Expected blob with digest "
                                    << digest
                                    << ", but downloaded blob has digest "
                                    << downloaded_digest);
    }

    BUILDBOX_LOG_TRACE(resourceName << ": " << bytes_downloaded
                                    << " bytes retrieved");
    return downloaded_data;
}

void CASClient::fetchStream(const Digest &digest,
                            const StreamConsumer &streamConsumer,
                            int64_t offset, std::optional<int64_t> limit,
                            GrpcClient::RequestStats *requestStats)
{
    BUILDBOX_LOG_TRACE("Downloading " << digest.hash() << " as a stream");
    const std::string resourceName = this->makeResourceName(digest, false);

    auto fetchLambda = [&](grpc::ClientContext &context) {
        ReadRequest request;
        request.set_resource_name(resourceName);
        request.set_read_offset(static_cast<int64_t>(offset));
        if (limit.has_value()) {
            request.set_read_limit(static_cast<int64_t>(limit.value()));
        }

        auto reader = this->d_bytestreamClient->Read(&context, request);
        size_t currentStartOffset = offset;

        ReadResponse response;
        while (reader->Read(&response)) {
            auto consumeStatus = streamConsumer(response, currentStartOffset);
            if (!consumeStatus.ok()) {
                return consumeStatus;
            }
            currentStartOffset += response.data().size();
        }

        grpc::Status read_status = reader->Finish();
        return read_status;
    };

    d_grpcClient->issueRequest(fetchLambda, "ByteStream.Read()",
                               digest.size_bytes(), requestStats);
}

void CASClient::download(int fd, const Digest &digest,
                         GrpcClient::RequestStats *requestStats)
{
    BUILDBOX_LOG_TRACE("Downloading " << digest.hash() << " to file");
    const std::string resourceName = this->makeResourceName(digest, false);

    int64_t bytesDownloaded = 0;
    auto digestContext = DigestGenerator::createDigestContext();

    auto downloadLambda = [&](grpc::ClientContext &context) {
        ReadRequest request;
        request.set_resource_name(resourceName);
        request.set_read_offset(bytesDownloaded);

        auto reader = this->d_bytestreamClient->Read(&context, request);

        ReadResponse response;
        while (reader->Read(&response)) {
            const auto &data = response.data();
            const char *data_ptr = data.data();
            size_t remaining = data.size();
            while (remaining > 0) {
                auto num_written = write(fd, data_ptr, remaining);
                if (num_written < 0) {
                    BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                        std::system_error, errno, std::generic_category,
                        "Error in write to descriptor "
                            << fd << ": " << std::strerror(errno));
                }
                data_ptr += // NOLINT
                            // (cppcoreguidelines-pro-bounds-pointer-arithmetic)
                    num_written;
                remaining -= num_written;
            }
            digestContext.update(data.c_str(), data.size());
            bytesDownloaded += static_cast<int64_t>(data.size());
        }

        auto read_status = reader->Finish();
        if (read_status.ok()) {
            struct stat st{};
            auto rc = fstat(fd, &st);
            if (rc < 0) {
                BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                    std::system_error, errno, std::generic_category,
                    "Error calling stat on file descriptor "
                        << fd << ": " << std::strerror(errno));
            }
            if (st.st_size != digest.size_bytes()) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "Expected " << digest.size_bytes()
                                << " bytes, but downloaded blob was "
                                << st.st_size << " bytes");
            }

            const auto downloaded_digest = digestContext.finalizeDigest();
            if (downloaded_digest != digest) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "Expected blob with digest "
                        << digest << ", but downloaded blob has digest "
                        << downloaded_digest);
            }

            BUILDBOX_LOG_TRACE(resourceName << ": " << st.st_size
                                            << " bytes retrieved");
        }

        return read_status;
    };

    d_grpcClient->issueRequest(downloadLambda, "ByteStream.Read()",
                               digest.size_bytes(), requestStats);
}

void CASClient::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)
{
    // The "path" must already exist in the filesystem this function
    // will populate it with the contents of "directory".

    const Directory directory = return_directory_callback(digest);

    // Downloading the files in this directory:
    OutputMap outputs;
    std::vector<Digest> file_digests;
    file_digests.reserve(static_cast<size_t>(directory.files_size()));

    for (const FileNode &fileNode : directory.files()) {
        file_digests.push_back(fileNode.digest());

        const std::string file_path = path + "/" + fileNode.name();
        outputs.emplace(fileNode.digest().hash(),
                        std::make_pair(file_path, fileNode));
    }
    download_callback(file_digests, outputs);

    // Creating the subdirectories in this level and recursively fetching
    // their contents:
    for (const DirectoryNode &directory_node : directory.directories()) {
        const std::string directory_path = path + "/" + directory_node.name();
        if (mkdir(directory_path.c_str(), PERMISSION_RWXALL) == -1) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Error in mkdir for directory \"" << directory_path << "\"");
        }

        downloadDirectory(directory_node.digest(), directory_path,
                          download_callback, return_directory_callback,
                          requestStats);
    }

    // Create symlinks, note: we just create the symlink. It's not the
    // responsibility of the worker/casd, to ensure the target is valid and
    // has contents.
    for (const SymlinkNode &symlink_node : directory.symlinks()) {
        if (symlink_node.target().empty() || symlink_node.name().empty()) {
            BUILDBOX_LOG_WARNING(
                "Symlink Node name or target empty skipping.");
            continue;
        }
        // Prepend the path to the symlink_node name.
        const std::string symlink_path = path + "/" + symlink_node.name();

        if (symlink(symlink_node.target().c_str(), symlink_path.c_str()) !=
            0) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Unable to create symlink: \""
                    << symlink_path + "\" to target: \""
                    << symlink_node.target() << "\"");
        }
    }

    const auto &properties = directory.node_properties();

    if (properties.has_unix_mode()) {
        const mode_t mode =
            static_cast<mode_t>(properties.unix_mode().value());
        if (chmod(path.c_str(), mode) != 0) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Could not chmod directory " << path << " to " << mode);
        }
    }
    if (properties.has_mtime()) {
        FileUtils::setFileTimes(
            path.c_str(), std::nullopt,
            TimeUtils::parse_timestamp(properties.mtime()));
    }
}

void CASClient::downloadDirectory(const Digest &digest,
                                  const std::string &path,
                                  GrpcClient::RequestStats *requestStats)
{
    download_callback_t download_blobs =
        [this, requestStats](const std::vector<Digest> &file_digests,
                             const OutputMap &outputs) {
            this->downloadBlobs(file_digests, outputs, requestStats);
        };

    return_directory_callback_t download_directory =
        [this, requestStats](const Digest &message_digest) {
            return this->fetchMessage<Directory>(message_digest, requestStats);
        };

    this->downloadDirectory(digest, path, download_blobs, download_directory,
                            requestStats);
}

void CASClient::downloadTree(const Digest &treeDigest,
                             const std::string &downloadPath,
                             GrpcClient::RequestStats *requestStats)
{
    // Get the blob containing the tree and format as a tree object
    buildboxcommon::Tree tree =
        fetchMessage<buildboxcommon::Tree>(treeDigest, requestStats);

    // The tree is flatend. To get information about a sub directory the
    // digest must be calculated and searched for in the tree's children.
    // Create a map to make lookup easier.
    std::unordered_map<Digest, Directory> treeNodes;
    // Populate map from tree
    for (const auto &child : tree.children()) {
        treeNodes[DigestGenerator::hash(child)] = child;
    }

    // Add the root directory to the treeNodes
    const Digest treeRootDigest = DigestGenerator::hash(tree.root());
    treeNodes[treeRootDigest] = tree.root();

    // Create the function for downloading blobs
    download_callback_t download_blobs =
        [this, requestStats](const std::vector<Digest> &file_digests,
                             const OutputMap &outputs) {
            this->downloadBlobs(file_digests, outputs, requestStats);
        };

    // Create the function to find the directory in the treeNodes from the
    // digest
    return_directory_callback_t find_directory = [treeNodes](
                                                     const Digest &digest) {
        // As the tree is flat information about the directory must be
        // looked up from the treeNodes.
        auto it = treeNodes.find(digest);
        if (it != treeNodes.end()) {
            return it->second;
        }
        else {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error,
                "Expected to download "
                    << digest.hash() << ", but couldn't find in treeNodes.");
        }
    };

    // Download the tree at the given downloadPath
    downloadDirectory(treeRootDigest, downloadPath, download_blobs,
                      find_directory, requestStats);
}

void CASClient::upload(const std::string &data, const Digest &digest,
                       GrpcClient::RequestStats *requestStats)
{
    BUILDBOX_LOG_DEBUG("Uploading " << digest.hash() << " from string");
    const auto data_size = static_cast<google::protobuf::int64>(data.size());

    if (data_size != digest.size_bytes()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::logic_error, "Digest length of "
                                  << digest.size_bytes() << " bytes for "
                                  << digest.hash()
                                  << " does not match string length of "
                                  << data_size << " bytes");
    }

    const std::string resourceName = this->makeResourceName(digest, true);

    WriteResponse response;

    auto uploadLambda = [&](grpc::ClientContext &context) {
        auto writer = this->d_bytestreamClient->Write(&context, &response);

        size_t offset = 0;
        bool lastChunk = false;
        while (!lastChunk) {
            WriteRequest request;
            request.set_resource_name(resourceName);
            request.set_write_offset(
                static_cast<google::protobuf::int64>(offset));

            const size_t uploadLength =
                std::min(bytestreamChunkSizeBytes(), data.size() - offset);

            request.set_data(&data[offset], uploadLength);
            offset += uploadLength;

            lastChunk = (offset == data.size());
            if (lastChunk) {
                request.set_finish_write(lastChunk);
            }

            if (!writer->Write(request)) {
                break;
            }
        }

        writer->WritesDone();
        auto status = writer->Finish();
        if (status.ok()) {
            if (response.committed_size() != digest.size_bytes()) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "Expected to upload "
                        << digest.size_bytes() << " bytes for "
                        << digest.hash() << ", but server reports "
                        << response.committed_size() << " bytes committed");
            }

            BUILDBOX_LOG_DEBUG(resourceName << ": " << offset
                                            << " bytes uploaded");
        }

        return status;
    };

    d_grpcClient->issueRequest(uploadLambda, "ByteStream.Write()",
                               digest.size_bytes(), requestStats);
}

void CASClient::upload(int fd, const Digest &digest,
                       GrpcClient::RequestStats *requestStats)
{
    std::vector<char> buffer(bytestreamChunkSizeBytes());
    BUILDBOX_LOG_DEBUG("Uploading " << digest.hash() << " from file");

    const std::string resourceName = this->makeResourceName(digest, true);

    WriteResponse response;
    auto uploadLambda = [&](grpc::ClientContext &context) {
        auto writer = this->d_bytestreamClient->Write(&context, &response);

        int64_t offset = 0;
        bool lastChunk = false;
        while (!lastChunk) {
            const ssize_t bytesRead =
                pread(fd, &buffer[0], bytestreamChunkSizeBytes(), offset);
            if (bytesRead < 0) {
                BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                    std::system_error, errno, std::generic_category,
                    "Error in read on descriptor " << fd);
            }

            WriteRequest request;
            request.set_resource_name(resourceName);
            request.set_write_offset(offset);
            request.set_data(&buffer[0], static_cast<size_t>(bytesRead));

            if (offset + bytesRead < digest.size_bytes()) {
                if (bytesRead == 0) {
                    BUILDBOXCOMMON_THROW_EXCEPTION(
                        std::runtime_error,
                        "Upload of " << digest.hash()
                                     << " failed: unexpected end of file");
                }
            }
            else {
                lastChunk = true;
                request.set_finish_write(true);
            }

            if (!writer->Write(request)) {
                break;
            }

            offset += bytesRead;
        }

        writer->WritesDone();
        auto status = writer->Finish();
        if (status.ok()) {
            if (response.committed_size() != digest.size_bytes()) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "Expected to upload "
                        << digest.size_bytes() << " bytes for "
                        << digest.hash() << ", but server reports "
                        << response.committed_size() << " bytes committed");
            }

            BUILDBOX_LOG_DEBUG(resourceName << ": " << offset
                                            << " bytes uploaded");
        }
        return status;
    };

    d_grpcClient->issueRequest(uploadLambda, "ByteStream.Write()",
                               digest.size_bytes(), requestStats);
}

void CASClient::uploadRequest(const UploadRequest &request,
                              GrpcClient::RequestStats *requestStats)
{
    if (!request.path.empty()) {
        const FileDescriptor fd(
            openat(request.dirfd, request.path.c_str(), O_RDONLY | O_CLOEXEC));
        if (fd.get() < 0) {
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Error in open for file \"" << request.path << "\"");
        }
        upload(fd.get(), request.digest, requestStats);
    }
    else if (request.fd >= 0) {
        upload(request.fd, request.digest, requestStats);
    }
    else {
        upload(request.data, request.digest);
    }
}

std::vector<CASClient::UploadResult>
CASClient::uploadBlobs(const std::vector<UploadRequest> &requests,
                       GrpcClient::RequestStats *requestStats)
{
    return CASClient::uploadBlobs(requests, nullptr, requestStats);
}

std::vector<CASClient::UploadResult>
CASClient::uploadBlobs(const std::vector<UploadRequest> &requests,
                       ThreadPool *threadPool,
                       GrpcClient::RequestStats *requestStats)
{
    std::vector<CASClient::UploadResult> results;

    // We first sort the requests by their sizes in ascending order, so
    // that we can then iterate through that result greedily trying to add
    // as many digests as possible to each request.
    std::vector<UploadRequest> request_list(requests);
    std::sort(request_list.begin(), request_list.end(),
              [](const UploadRequest &r1, const UploadRequest &r2) {
                  return r1.digest.size_bytes() < r2.digest.size_bytes();
              });

    // Grouping the requests into batches (we only need to look at the
    // Digests for their sizes):
    std::vector<Digest> digests;
    digests.reserve(request_list.size());
    for (const auto &r : request_list) {
        digests.push_back(r.digest);
    }

    const auto batches = makeBatches(digests);

    {
        FutureGroup<std::vector<CASClient::UploadResult>> futureGroup(
            threadPool);
        std::vector<std::tuple<
            size_t, size_t,
            std::shared_future<std::vector<CASClient::UploadResult>>>>
            batchResults;
        for (const auto &batch_range : batches) {
            const size_t batch_start = batch_range.first;
            const size_t batch_end = batch_range.second;

            const auto batchUploadLambda = [batch_start, batch_end,
                                            &request_list, &requestStats,
                                            this]() {
                std::vector<CASClient::UploadResult> digests_not_uploaded =
                    batchUpload(request_list, batch_start, batch_end,
                                requestStats);
                return digests_not_uploaded;
            };

            batchResults.emplace_back(batch_start, batch_end,
                                      futureGroup.add(batchUploadLambda));
        }

        // Check results of BatchUpdateBlobs
        for (auto &result : batchResults) {
            auto &[batchStart, batchEnd, resultFuture] = result;
            try {
                const auto digestsNotUploaded = resultFuture.get();
                std::move(digestsNotUploaded.cbegin(),
                          digestsNotUploaded.cend(),
                          std::back_inserter(results));
            }
            catch (const GrpcError &e) {
                for (auto d = batchStart; d < batchEnd; d++) {
                    results.emplace_back(request_list.at(d).digest, e.status);
                }
            }
            catch (const std::runtime_error &e) {
                BUILDBOX_LOG_ERROR("Batch upload failed: " +
                                   std::string(e.what()));

                // Report an INTERNAL error code for failed uploads in case
                // there is no throw.
                const auto failed_status = grpc::Status(
                    grpc::StatusCode::INTERNAL, grpc::string(e.what()));
                for (auto d = batchStart; d < batchEnd; d++) {
                    results.emplace_back(request_list.at(d).digest,
                                         failed_status);
                }
            }
        }
    }

    {
        // Fetching all those digests that might need to be uploaded using the
        // Bytestream API. Those will be in the range [batch_end,
        // batches.size()).
        const size_t batch_end =
            batches.empty() ? 0 : batches.rbegin()->second;
        FutureGroup<void> futureGroup(threadPool);
        std::vector<std::pair<size_t, std::shared_future<void>>> streamResults;

        for (auto d = batch_end; d < request_list.size(); d++) {
            const auto uploadLambda = [d, &request_list, &requestStats,
                                       this]() {
                uploadRequest(request_list[d], requestStats);
            };

            streamResults.emplace_back(d, futureGroup.add(uploadLambda));
        }

        // Check results of streaming uploads
        for (auto &result : streamResults) {
            auto &[d, resultFuture] = result;
            try {
                resultFuture.get();
            }
            catch (const GrpcError &e) {
                results.emplace_back(request_list[d].digest, e.status);
            }
            catch (const std::runtime_error &e) {
                BUILDBOX_LOG_ERROR("Failed to upload blob: " +
                                   std::string(e.what()));
                results.emplace_back(
                    request_list[d].digest,
                    grpc::Status(grpc::StatusCode::INTERNAL, e.what()));
            }
        }
    }

    return results;
}

CASClient::DownloadBlobsResult
CASClient::downloadBlobs(const std::vector<Digest> &digests, int temp_dirfd,
                         ThreadPool *threadPool,
                         GrpcClient::RequestStats *requestStats)
{
    CASClient::DownloadBlobsResult downloaded_data;
    // Used if multi-threaded
    std::optional<std::mutex> downloadDataMutex =
        (threadPool == nullptr) ? std::nullopt
                                : std::make_optional<std::mutex>();

    // Writing the data directly into the result. (We know that the status
    // code will be OK for each of these blobs.)
    auto write_blob = [&](const std::string &hash, const std::string &data) {
        google::rpc::Status status;
        status.set_code(grpc::StatusCode::OK);

        {
            std::optional<std::unique_lock<std::mutex>> lock =
                downloadDataMutex.has_value()
                    ? std::make_optional<std::unique_lock<std::mutex>>(
                          *downloadDataMutex)
                    : std::nullopt;
            downloaded_data.emplace(hash, std::make_pair(status, data));
        }
    };

    const CASClient::DownloadResults download_results = downloadBlobs(
        digests, write_blob, temp_dirfd, threadPool, requestStats);

    // And adding the codes of the hashes that failed into the result:
    for (const auto &entry : download_results) {
        const Digest &digest = entry.first;
        const google::rpc::Status &status = entry.second;

        if (status.code() != grpc::StatusCode::OK) {
            downloaded_data.emplace(digest.hash(), std::make_pair(status, ""));
        }
    }

    return downloaded_data;
}

CASClient::DownloadBlobsResult
CASClient::downloadBlobs(const std::vector<Digest> &digests,
                         ThreadPool *threadPool,
                         GrpcClient::RequestStats *requestStats)
{
    return downloadBlobs(digests, -1, threadPool, requestStats);
}

CASClient::DownloadBlobsResult
CASClient::downloadBlobsToDirectory(const std::vector<Digest> &digests,
                                    int temp_dirfd, ThreadPool *threadPool,
                                    GrpcClient::RequestStats *requestStats)
{
    return downloadBlobs(digests, temp_dirfd, threadPool, requestStats);
}

void CASClient::downloadBlobs(const std::vector<Digest> &digests,
                              const OutputMap &outputs,
                              GrpcClient::RequestStats *requestStats)
{

    auto write_blob = [&](const std::string &hash, const std::string &data) {
        const std::pair<OutputMap::const_iterator, OutputMap::const_iterator>
            range = outputs.equal_range(hash);

        for (auto it = range.first; it != range.second; it++) {
            const auto &[path, fileNode] = it->second;

            mode_t file_permissions = 0;
            if (fileNode.node_properties().has_unix_mode()) {
                file_permissions = static_cast<mode_t>(
                    fileNode.node_properties().unix_mode().value());
            }
            else {
                // use default if none available
                file_permissions =
                    S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH; // 0644
                if (fileNode.is_executable()) {
                    file_permissions |= S_IXUSR | S_IXGRP | S_IXOTH; // 0755
                }
            }

            std::optional<std::chrono::time_point<std::chrono::system_clock>>
                mtime = std::nullopt;
            if (fileNode.node_properties().has_mtime()) {
                mtime = TimeUtils::parse_timestamp(
                    fileNode.node_properties().mtime());
            }

            FileUtils::writeFileAtomically(path, data, file_permissions, "",
                                           TempDefaults::DEFAULT_TMP_PREFIX,
                                           mtime);
        }
    };

    const auto results =
        downloadBlobs(digests, write_blob, -1, nullptr, requestStats);
    // check error and throw
    for (const auto &[digest, status] : results) {
        if (status.code() != grpc::StatusCode::OK) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error, "Failed to download blob, digest="
                                        << digest << " code=" << status.code()
                                        << " message=" << status.message());
        }
    }
}

CASClient::DownloadResults
CASClient::downloadBlobs(const std::vector<Digest> &digests,
                         const WriteBlobCallback &write_blob, int temp_dirfd,
                         ThreadPool *threadPool,
                         GrpcClient::RequestStats *requestStats)
{
    DownloadResults download_results;
    download_results.reserve(digests.size());

    // We first sort the digests by their sizes in ascending order, so that
    // we can then iterate through that result greedily trying to add as
    // many digests as possible to each request.
    auto request_list(digests);
    std::sort(request_list.begin(), request_list.end(),
              [](const Digest &d1, const Digest &d2) {
                  return d1.size_bytes() < d2.size_bytes();
              });
    const auto batches = makeBatches(request_list);

    {
        FutureGroup<DownloadResults> futureGroup(threadPool);
        std::vector<std::tuple<size_t, // start
                               size_t, // end
                               std::shared_future<DownloadResults>>>
            batchResultFutures;
        for (const auto &batch_range : batches) {
            const size_t batch_start = batch_range.first;
            const size_t batch_end = batch_range.second;

            // For each batch, we make the request and then store whether it
            // was successful in the result:
            auto batchDownloadLambda = [batch_start, batch_end, &request_list,
                                        &write_blob, &requestStats, temp_dirfd,
                                        this]() {
                auto batchResults =
                    batchDownload(request_list, batch_start, batch_end,
                                  write_blob, requestStats, temp_dirfd);

                return batchResults;
            };
            batchResultFutures.emplace_back(
                batch_start, batch_end, futureGroup.add(batchDownloadLambda));
        }
        // Handle results of batch requests
        for (auto &[batchStart, batchEnd, resultFuture] : batchResultFutures) {
            try {
                const auto batchResults = resultFuture.get();

                std::move(batchResults.cbegin(), batchResults.cend(),
                          std::back_inserter(download_results));
            }
            catch (const GrpcError &e) {
                google::rpc::Status failed_status;
                failed_status.set_code(e.status.error_code());
                failed_status.set_message(e.status.error_message());
                for (auto d = batchStart; d < batchEnd; d++) {
                    const Digest &digest = request_list.at(d);
                    download_results.emplace_back(digest, failed_status);
                }
            }
            catch (const std::runtime_error &e) {
                BUILDBOX_LOG_ERROR("Batch download failed: " +
                                   std::string(e.what()));

                // If we don't throw, we need to report in the result that the
                // requested digests failed:
                google::rpc::Status failed_status;
                failed_status.set_code(grpc::StatusCode::INTERNAL);

                for (auto d = batchStart; d < batchEnd; d++) {
                    const Digest &digest = request_list.at(d);
                    download_results.emplace_back(digest, failed_status);
                }
            }
        }
    }

    {
        // Fetching all those digests that might need to be downloaded using
        // the Bytestream API. Those will be in the range [batch_end,
        // batches.size()).
        size_t batch_end = 0;
        if (batches.empty()) {
            batch_end = 0;
        }
        else {
            batch_end = batches.rbegin()->second;
        }

        FutureGroup<void> futureGroup(threadPool);
        std::vector<std::pair<size_t, std::shared_future<void>>>
            streamResultFutures;
        for (auto d = batch_end; d < request_list.size(); d++) {
            const Digest &digest = request_list[d];

            auto downloadLambda = [digest, &write_blob, &requestStats,
                                   temp_dirfd, this]() {
                if (temp_dirfd < 0) {
                    const auto data = fetchString(digest, requestStats);
                    write_blob(digest.hash(), data);
                }
                else {
                    // Download blob directly into a file to avoid excessive
                    // memory usage for large files.
                    const auto &path = digest.hash();
                    FileDescriptor fd(
                        openat(temp_dirfd, path.c_str(),
                               O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC,
                               PERMISSION_RWUSR));
                    if (fd.get() < 0) {
                        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                            std::system_error, errno, std::generic_category,
                            "CASClient::downloadBlobs: Failed to create file "
                            "\"" << path
                                 << "\" in temporary directory");
                    }
                    download(fd.get(), digest, requestStats);

                    write_blob(digest.hash(), path);
                }
            };

            streamResultFutures.emplace_back(d,
                                             futureGroup.add(downloadLambda));
        }

        // Collect results of downloads
        for (auto &[d, resultFuture] : streamResultFutures) {
            google::rpc::Status downloadStatus;
            try {
                resultFuture.get();
            }
            catch (const GrpcError &e) {
                downloadStatus.set_code(e.status.error_code());
                downloadStatus.set_message(e.status.error_message());
            }
            catch (const std::runtime_error &e) {
                BUILDBOX_LOG_ERROR("Error: fetchString(): " +
                                   std::string(e.what()));
                downloadStatus.set_code(grpc::StatusCode::INTERNAL);
            }
            download_results.emplace_back(request_list[d], downloadStatus);
        }
    }

    return download_results;
}

std::unique_ptr<CASClient::StagedDirectory> CASClient::stage(
    const Digest &root_digest, const StageTreeRequest_StagingMode stagingMode,
    const std::string &path,
    const ProcessCredentialsGetter &processCredentialsGetter,
    const std::vector<std::string> &postActionCommands,
    const std::optional<RemoteApisSocketConfig> &remoteApisSocket) const
{
    auto context = std::make_shared<grpc::ClientContext>();

    std::shared_ptr<
        grpc::ClientReaderWriterInterface<StageTreeRequest, StageTreeResponse>>
        reader_writer(d_localCasClient->StageTree(context.get()));

    StageTreeRequest request;
    request.set_instance_name(d_grpcClient->instanceName());
    request.mutable_root_digest()->CopyFrom(root_digest);
    request.set_path(path);
    request.set_staging_mode(stagingMode);

    // Specify the credentials used to access the staged tree
    if (processCredentialsGetter) {
        const auto creds = processCredentialsGetter();
        request.mutable_access_credentials()->set_uid(creds.uid);
        request.mutable_access_credentials()->set_gid(creds.gid);
    }

    for (const auto &command : postActionCommands) {
        *request.add_pre_unstage_commands() = command;
    }

    if (remoteApisSocket) {
        request.mutable_remote_apis_socket()->CopyFrom(*remoteApisSocket);
    }

    const bool successful_write = reader_writer->Write(request);
    if (!successful_write) {
        const grpc::Status write_error_status = reader_writer->Finish();
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Error staging \"" << toString(root_digest) << "\" into \"" << path
                               << "\": \""
                               << write_error_status.error_message() << "\"");
    }

    StageTreeResponse response;
    const bool successful_read = reader_writer->Read(&response);
    if (!successful_read) {
        reader_writer->WritesDone();
        const grpc::Status read_error_status = reader_writer->Finish();
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Error staging \"" << toString(root_digest) << "\" into \"" << path
                               << "\": \"" << read_error_status.error_message()
                               << "\"");
    }

    return std::make_unique<StagedDirectory>(context, reader_writer,
                                             response.path());
}

DirectoryTree CASClient::getTree(const Digest &root_digest)
{
    grpc::ClientContext context;
    GetTreeRequest request;
    request.set_instance_name(d_grpcClient->instanceName());
    request.mutable_root_digest()->CopyFrom(root_digest);

    std::unique_ptr<grpc::ClientReaderInterface<GetTreeResponse>> reader(
        d_casClient->GetTree(&context, request));

    DirectoryTree tree;
    tree.root = root_digest;

    GetTreeResponse response;
    while (reader->Read(&response)) {
        BUILDBOX_LOG_TRACE("\n" << toString(response.directories()));
        for (const auto &directory : response.directories()) {
            tree.add(directory);
        }
    }

    grpc::Status status = reader->Finish();
    if (!status.ok()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Error getting tree for digest \""
                                           << toString(root_digest)
                                           << "\", status = ["
                                           << status.error_code() << ": \""
                                           << status.error_message() + "\"]");
    }

    if (!tree.map.contains(tree.root)) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Incomplete GetTree response, missing "
                                           << tree.root);
    }

    for (const auto &[digest, directory] : tree.map) {
        for (const auto &node : directory.directories()) {
            if (!tree.map.contains(node.digest())) {
                BUILDBOXCOMMON_THROW_EXCEPTION(
                    std::runtime_error,
                    "Incomplete GetTree response, missing " << node.digest());
            }
        }
    }

    return tree;
}

DirectoryTree
CASClient::getTreeFromTreeMessage(const Digest &tree_digest,
                                  GrpcClient::RequestStats *requestStats)
{
    Tree treeMessage = this->fetchMessage<Tree>(tree_digest, requestStats);

    DirectoryTree tree;

    tree.root = DigestGenerator::hash(treeMessage.root());
    tree.map.emplace(tree.root, treeMessage.root());
    for (const auto &child : treeMessage.children()) {
        tree.add(child);
    };

    return tree;
}

FetchTreeResponse CASClient::fetchTree(const Digest &digest,
                                       const bool fetch_file_blobs,
                                       GrpcClient::RequestStats *requestStats)
{
    FetchTreeRequest request;
    request.set_instance_name(d_grpcClient->instanceName());
    request.set_fetch_file_blobs(fetch_file_blobs);
    *request.mutable_root_digest() = digest;

    FetchTreeResponse response;
    const auto fetchTreeLambda = [this, &request,
                                  &response](grpc::ClientContext &context) {
        return d_localCasClient->FetchTree(&context, request, &response);
    };

    d_grpcClient->issueRequest(fetchTreeLambda, "LocalCAS.FetchTree()",
                               requestStats);
    return response;
}

CaptureTreeResponse
CASClient::captureTree(const std::vector<std::string> &paths,
                       const std::vector<std::string> &properties,
                       bool bypass_local_cache, mode_t unixModeMask,
                       GrpcClient::RequestStats *requestStats,
                       const Command_OutputDirectoryFormat outputDirFormat,
                       bool allowChmodToRead) const
{
    return captureTree("", paths, properties, bypass_local_cache, unixModeMask,
                       requestStats, outputDirFormat, allowChmodToRead);
}

CaptureTreeResponse CASClient::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,
    const Command_OutputDirectoryFormat outputDirFormat,
    bool allowChmodToRead) const
{
    CaptureTreeRequest request;
    request.set_instance_name(d_grpcClient->instanceName());
    request.set_bypass_local_cache(bypass_local_cache);
    request.set_root(root);
    request.set_output_directory_format(outputDirFormat);
    request.mutable_unix_mode_mask()->set_value(
        static_cast<uint32_t>(unixModeMask));
    request.set_allow_chmod_to_read(allowChmodToRead);

    for (const std::string &path : paths) {
        request.add_path(path);
    }

    for (const std::string &property : properties) {
        request.add_node_properties(property);
    }

    CaptureTreeResponse response;

    const auto captureLambda = [&](grpc::ClientContext &context) {
        return d_localCasClient->CaptureTree(&context, request, &response);
    };

    d_grpcClient->issueRequest(captureLambda, "LocalCAS.CaptureTree()",
                               requestStats);

    return response;
}

std::vector<CaptureFilesResponse_Response>
CASClient::captureFiles(const std::vector<std::string> &paths,
                        const std::vector<std::string> &properties,
                        bool bypass_local_cache, bool skipUpload,
                        GrpcClient::RequestStats *requestStats) const
{
    return captureFiles("", paths, properties, bypass_local_cache, skipUpload,
                        requestStats);
}

std::vector<HashFilesResponse_Response>
CASClient::hashFiles(const std::vector<std::string> &paths,
                     const std::vector<std::string> &properties,
                     GrpcClient::RequestStats *requestStats) const
{
    std::vector<HashFilesResponse_Response> responses;

    // reuse capture batches since the request structure is similar
    const auto batches = makeCaptureBatches(paths);
    for (const auto &batch_range : batches) {
        const size_t batch_start = batch_range.first;
        const size_t batch_end = batch_range.second;

        HashFilesRequest request;
        request.set_instance_name(d_grpcClient->instanceName());

        for (auto i = batch_start; i < batch_end; i++) {
            request.add_path(paths[i]);
        }

        for (const std::string &property : properties) {
            request.add_node_properties(property);
        }

        HashFilesResponse hashFilesResponse;

        const auto hashLambda = [&](grpc::ClientContext &context) {
            return d_localCasClient->HashFiles(&context, request,
                                               &hashFilesResponse);
        };

        d_grpcClient->issueRequest(hashLambda, "LocalCAS.HashFiles()",
                                   requestStats);

        for (const auto &response : hashFilesResponse.responses()) {
            responses.push_back(response);
        }
    }

    return responses;
}

std::vector<CaptureFilesResponse_Response> CASClient::captureFiles(
    const std::string &root, const std::vector<std::string> &paths,
    const std::vector<std::string> &properties, bool bypass_local_cache,
    bool skipUpload, GrpcClient::RequestStats *requestStats) const
{
    std::vector<CaptureFilesResponse_Response> responses;

    const auto batches = makeCaptureBatches(paths);
    for (const auto &batch_range : batches) {
        const size_t batch_start = batch_range.first;
        const size_t batch_end = batch_range.second;

        CaptureFilesRequest request;
        request.set_instance_name(d_grpcClient->instanceName());
        request.set_bypass_local_cache(bypass_local_cache);
        request.set_root(root);
        request.set_skip_upload(skipUpload);

        for (auto i = batch_start; i < batch_end; i++) {
            request.add_path(paths[i]);
        }

        for (const std::string &property : properties) {
            request.add_node_properties(property);
        }

        CaptureFilesResponse captureFilesResponse;

        const auto captureLambda = [&](grpc::ClientContext &context) {
            return d_localCasClient->CaptureFiles(&context, request,
                                                  &captureFilesResponse);
        };

        d_grpcClient->issueRequest(captureLambda, "LocalCAS.CaptureFiles()",
                                   requestStats);

        for (const auto &response : captureFilesResponse.responses()) {
            responses.push_back(response);
        }
    }

    return responses;
}

std::vector<CASClient::UploadResult>
CASClient::batchUpload(const std::vector<UploadRequest> &requests,
                       const size_t start_index, const size_t end_index,
                       GrpcClient::RequestStats *requestStats)
{
    assert(start_index <= end_index);
    assert(end_index <= requests.size());

    BatchUpdateBlobsRequest request;
    request.set_instance_name(d_grpcClient->instanceName());

    for (auto d = start_index; d < end_index; d++) {
        const UploadRequest &upload_request = requests[d];
        auto entry = request.add_requests();
        entry->mutable_digest()->CopyFrom(upload_request.digest);
        if (!upload_request.path.empty()) {
            entry->set_data(FileUtils::getFileContents(
                upload_request.dirfd, upload_request.path.c_str()));
        }
        else if (upload_request.fd >= 0) {
            entry->set_data(FileUtils::getFileContents(upload_request.fd));
        }
        else {
            entry->set_data(upload_request.data);
        }
    }

    BUILDBOX_LOG_TRACE("BatchUpdateBlobs Request serialized message size = "
                       << request.ByteSizeLong());

    BatchUpdateBlobsResponse response;
    auto batchUploadLambda = [&](grpc::ClientContext &context) {
        auto status =
            this->d_casClient->BatchUpdateBlobs(&context, request, &response);
        return status;
    };

    d_grpcClient->issueRequest(batchUploadLambda, "BatchUpdateBlobs()",
                               requestStats);

    BUILDBOX_LOG_TRACE("BatchUpdateBlobs Response serialized message size = "
                       << response.ByteSizeLong());

    std::vector<CASClient::UploadResult> results;
    for (const auto &uploadResponse : response.responses()) {
        if (uploadResponse.status().code() != GRPC_STATUS_OK) {
            results.emplace_back(
                uploadResponse.digest(),
                grpc::Status(grpc::StatusCode(uploadResponse.status().code()),
                             uploadResponse.status().message()));
        }
    }
    return results;
}

void writeFile(int dirfd, const std::string &path, const char *buf, size_t n)
{
    FileDescriptor fd(openat(dirfd, path.c_str(),
                             O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC,
                             PERMISSION_RWALL));
    if (fd.get() < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to create file \"" << path << "\" in temporary directory");
    }

    while (n > 0) {
        ssize_t ret = write(fd.get(), buf, n);
        if (ret < 0) {
            if (errno == EINTR) {
                // Interrupted by signal, retry
                continue;
            }
            BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
                std::system_error, errno, std::system_category,
                "Failed to write file \"" << path
                                          << "\" in temporary directory");
        }
        // write() should never return 0
        assert(ret != 0);
        buf += // NOLINT (cppcoreguidelines-pro-bounds-pointer-arithmetic)
            ret;
        n -= ret;
    }
}

CASClient::DownloadResults CASClient::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)
{
    assert(start_index <= end_index);
    assert(end_index <= digests.size());

    BatchReadBlobsRequest request;
    request.set_instance_name(d_grpcClient->instanceName());

    for (auto d = start_index; d < end_index; d++) {
        auto digest = request.add_digests();
        digest->CopyFrom(digests[d]);
    }
    BUILDBOX_LOG_TRACE("BatchReadBlobs Request serialized message size = "
                       << request.ByteSizeLong());

    BatchReadBlobsResponse response;
    auto batchDownloadLambda = [&](grpc::ClientContext &context) {
        auto status =
            this->d_casClient->BatchReadBlobs(&context, request, &response);
        return status;
    };

    d_grpcClient->issueRequest(batchDownloadLambda, "BatchReadBlobs()",
                               requestStats);

    BUILDBOX_LOG_TRACE("BatchReadBlobs Response serialized message size = "
                       << response.ByteSizeLong());

    DownloadResults download_results;
    download_results.reserve(static_cast<size_t>(response.responses_size()));

    for (const auto &downloadResponse : response.responses()) {
        if (downloadResponse.status().code() == GRPC_STATUS_OK) {
            const auto downloaded_digest =
                DigestGenerator::hash(downloadResponse.data());
            if (downloaded_digest != downloadResponse.digest()) {
                google::rpc::Status status;
                status.set_code(grpc::StatusCode::INTERNAL);
                std::ostringstream error;
                error << "Expected blob with digest "
                      << downloadResponse.digest()
                      << ", but downloaded blob has digest "
                      << downloaded_digest;
                status.set_message(error.str());
                download_results.emplace_back(downloadResponse.digest(),
                                              status);
                continue;
            }

            if (temp_dirfd < 0) {
                write_blob_function(downloadResponse.digest().hash(),
                                    downloadResponse.data());
            }
            else {
                const auto path = downloadResponse.digest().hash();
                const auto &data = downloadResponse.data();
                writeFile(temp_dirfd, path, data.c_str(), data.size());
                write_blob_function(downloadResponse.digest().hash(), path);
            }
        }

        download_results.emplace_back(downloadResponse.digest(),
                                      downloadResponse.status());
    }

    return download_results;
}

std::vector<std::pair<size_t, size_t>>
CASClient::makeBatches(const std::vector<Digest> &digests)
{
    // The below value is set based on rounding up of static data sizes
    // in the gRPC classes BatchUpdateBlobsRequest(upload) and
    // BatchReadBlobsRequest(download). This is an attempt to factor in
    // this size into the overall batch size equation
    static const size_t SIZEOF_ESTIMATED_TOP_LEVEL_GRPC_CONTAINER = 256;

    // The bulk of the per-blob overhead are the Digest and Status
    // data structures. The 'digest' hash needs 128 bytes for the largest
    // hash in the protocol (SHA-512) and the 'status' can contain an
    // optional string message
    static const size_t PER_BLOB_METADATA_SIZE = 256;

    // A batch is a pair of indexes into the vector of `digests` that
    // can all fit in one `d_maxBatchTotalSizeBytes` sized buffer;
    // The indices are semantically represented by [batch_start, batch_end)
    std::vector<std::pair<size_t, size_t>> batches;
    const size_t max_batch_size =
        d_maxBatchTotalSizeBytes - SIZEOF_ESTIMATED_TOP_LEVEL_GRPC_CONTAINER;
    size_t batch_start = 0;
    size_t batch_end = 0;
    while (batch_end < digests.size()) {
        if (digests[batch_end].size_bytes() >
            static_cast<int64_t>(max_batch_size - PER_BLOB_METADATA_SIZE)) {
            // All digests from `batch_end` to the end of the list are
            // larger than what we can request; stop.
            return batches;
        }

        // Adding all the digests that we can until we run out or exceed
        // the batch request limit...
        size_t bytes_in_batch = 0;
        while (batch_end < digests.size() &&
               digests[batch_end].size_bytes() <=
                   static_cast<int64_t>(max_batch_size - bytes_in_batch -
                                        PER_BLOB_METADATA_SIZE)) {
            bytes_in_batch +=
                (digests[batch_end].size_bytes() + PER_BLOB_METADATA_SIZE);
            batch_end++;
        }

        batches.emplace_back(batch_start, batch_end);
        batch_start = batch_end;
    }

    return batches;
}

std::vector<std::pair<size_t, size_t>>
CASClient::makeCaptureBatches(const std::vector<std::string> &paths) const
{
    // The below value is set based on rounding up of static data sizes
    // in the gRPC message type CaptureFilesResponse. This is an attempt to
    // factor in this size into the overall batch size equation
    static const size_t SIZEOF_ESTIMATED_TOP_LEVEL_GRPC_CONTAINER = 256;

    // The bulk of the per-blob overhead are the Digest and Status
    // data structures. The 'digest' hash needs 128 bytes for the largest
    // hash in the protocol (SHA-512) and the 'status' can contain an
    // optional string message
    static const size_t PER_BLOB_METADATA_SIZE = 256;

    // A batch is a pair of indexes into the vector of `paths` that
    // can all fit in one `d_maxBatchTotalSizeBytes` sized buffer;
    // The indices are semantically represented by [batch_start, batch_end)
    std::vector<std::pair<size_t, size_t>> batches;
    const size_t max_batch_size =
        d_maxBatchTotalSizeBytes - SIZEOF_ESTIMATED_TOP_LEVEL_GRPC_CONTAINER;
    size_t batch_start = 0;
    size_t batch_end = 0;
    while (batch_end < paths.size()) {
        // Adding all the digests that we can until we run out or exceed
        // the batch request limit...
        size_t bytes_in_batch = 0;
        while (batch_end < paths.size() &&
               paths[batch_end].size() <=
                   max_batch_size - bytes_in_batch - PER_BLOB_METADATA_SIZE) {
            bytes_in_batch +=
                (paths[batch_end].size() + PER_BLOB_METADATA_SIZE);
            batch_end++;
        }

        batches.emplace_back(batch_start, batch_end);
        batch_start = batch_end;
    }

    return batches;
}

std::vector<Digest>
CASClient::findMissingBlobs(const std::vector<Digest> &digests,
                            GrpcClient::RequestStats *requestStats)
{
    FindMissingBlobsRequest request;
    request.set_instance_name(d_grpcClient->instanceName());

    // We take the given digests and split them across requests to not
    // exceed the maximum size of a gRPC message:
    std::vector<FindMissingBlobsRequest> requests_to_issue;
    size_t batch_size = 0;
    for (const Digest &digest : digests) {
        const size_t digest_size = digest.ByteSizeLong();
        if (batch_size + digest_size > bytestreamChunkSizeBytes()) {
            requests_to_issue.push_back(request);
            request.clear_blob_digests();
            batch_size = 0;
        }
        else {
            batch_size += digest_size;
        }

        auto entry = request.add_blob_digests();
        entry->CopyFrom(digest);
    }
    requests_to_issue.push_back(request);

    std::vector<Digest> missing_blobs;
    for (const auto &request_to_issue : requests_to_issue) {
        FindMissingBlobsResponse response;
        auto findMissingBlobsLambda = [&](grpc::ClientContext &context) {
            auto status = this->d_casClient->FindMissingBlobs(
                &context, request_to_issue, &response);
            return status;
        };

        d_grpcClient->issueRequest(findMissingBlobsLambda,
                                   "FindMissingBlobs()", requestStats);

        missing_blobs.insert(missing_blobs.end(),
                             response.missing_blob_digests().cbegin(),
                             response.missing_blob_digests().cend());
    }

    return missing_blobs;
}

std::vector<CASClient::UploadResult>
CASClient::uploadDirectory(const std::string &path,
                           Digest *root_directory_digest, Tree *tree,
                           GrpcClient::RequestStats *requestStats)
{
    const FileDescriptor dirfd(
        open(path.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC));
    if (dirfd.get() < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::system_category,
            "Failed to open path \"" << path << "\"");
    }
    return uploadDirectory(dirfd.get(), root_directory_digest, tree,
                           requestStats);
}

std::vector<CASClient::UploadResult>
CASClient::uploadDirectory(int dirfd, Digest *root_directory_digest,
                           Tree *tree, GrpcClient::RequestStats *requestStats)
{
    // Recursing through the directory and building a map:

    const Merklizer merklizer(false /*no follow*/);
    auto merklizeResult = merklizer.merklize(dirfd);

    if (root_directory_digest != nullptr) {
        root_directory_digest->CopyFrom(merklizeResult.d_rootDigest);
    }
    if (tree != nullptr) {
        tree->CopyFrom(merklizeResult.tree());
    }

    // FindMissingBlobs(), this modifes the maps in-place
    missingDigests(&merklizeResult.d_digestToDirectoryBlob,
                   &merklizeResult.d_digestToPath, requestStats);

    // Batch upload the blobs missing in the remote:
    std::vector<UploadRequest> upload_requests;
    upload_requests.reserve(merklizeResult.d_digestToDirectoryBlob.size() +
                            merklizeResult.d_digestToPath.size());
    for (const auto &entry : merklizeResult.d_digestToDirectoryBlob) {
        const Digest digest = entry.first;
        upload_requests.emplace_back(digest, entry.second);
    }
    for (const auto &entry : merklizeResult.d_digestToPath) {
        const Digest digest = entry.first;
        upload_requests.emplace_back(
            UploadRequest::from_path(digest, dirfd, entry.second));
    }

    return uploadBlobs(upload_requests, nullptr, requestStats);
}

void CASClient::missingDigests(digest_string_map *digest_blob_map,
                               digest_string_map *digest_path_map,
                               GrpcClient::RequestStats *requestStats)
{
    std::vector<Digest> digests_in_directory;
    digests_in_directory.reserve(digest_path_map->size() +
                                 digest_blob_map->size());
    for (const auto &entry : *digest_blob_map) {
        digests_in_directory.push_back(entry.first);
    }
    for (const auto &entry : *digest_path_map) {
        digests_in_directory.push_back(entry.first);
    }

    const std::vector<Digest> missing_blobs =
        findMissingBlobs(digests_in_directory, requestStats);

    digest_string_map missing_blob_map, missing_path_map;
    for (const Digest &digest : missing_blobs) {
        // Digest must be in at least one of the two maps.
        // It's not necessary to return it in both maps.
        // Prefer blob over path to avoid unnecessary file read.
        const auto &digest_blob_it = digest_blob_map->find(digest);
        if (digest_blob_it != digest_blob_map->end()) {
            missing_blob_map[digest] = digest_blob_it->second;
        }
        else {
            missing_path_map[digest] = digest_path_map->at(digest);
        }
    }
    digest_blob_map->swap(missing_blob_map);
    digest_path_map->swap(missing_path_map);
}

ExportFilesResponse
CASClient::exportFiles(std::span<OutputFile> files,
                       const std::filesystem::path &path,
                       GrpcClient::RequestStats *requestStats)
{
    FileDescriptor rootDirfd(
        open(path.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC));
    if (rootDirfd.get() < 0) {
        BUILDBOXCOMMON_THROW_SYSTEM_EXCEPTION(
            std::system_error, errno, std::generic_category,
            "Error opening directory '" << path
                                        << "': " << std::strerror(errno));
    }

    ExportFilesRequest request;
    request.set_instance_name(d_grpcClient->instanceName());
    request.set_path(std::filesystem::absolute(path));
    for (const auto &file : files) {
        const auto filePath = std::filesystem::path(file.path());
        if (!filePath.is_relative()) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::logic_error,
                "Path of OutputFile is not relative: '" << filePath << "'");
        }
        // Ensure parent path exists
        const auto parentPath = filePath.parent_path();
        if (!parentPath.empty()) {
            FileUtils::createDirectory(rootDirfd.get(), parentPath.c_str());
        }
        request.add_output_files()->CopyFrom(file);
    }

    ExportFilesResponse response;
    const auto exportFilesLambda = [this, &request,
                                    &response](grpc::ClientContext &context) {
        return d_localCasClient->ExportFiles(&context, request, &response);
    };

    d_grpcClient->issueRequest(exportFilesLambda, "LocalCAS.ExportFiles()",
                               requestStats);
    return response;
}

CASClient::StagedDirectory::StagedDirectory(
    std::shared_ptr<grpc::ClientContext> context,
    std::shared_ptr<
        grpc::ClientReaderWriterInterface<StageTreeRequest, StageTreeResponse>>
        reader_writer,
    const std::string &path)
    : d_context(std::move(context)), d_reader_writer(std::move(reader_writer)),
      d_path(path)
{
}

CASClient::StagedDirectory::~StagedDirectory()
{
    // According to the LocalCAS spec., an empty request tells the
    // server to clean up.
    d_reader_writer->Write(StageTreeRequest());
    d_reader_writer->WritesDone();
}

} // namespace buildboxcommon
