/*
 * Copyright 2020 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 <buildboxcasd_findmissingblobsclient.h>

#include <buildboxcommon_grpcerror.h>

using namespace buildboxcasd;

FindMissingBlobsClient::FindMissingBlobsClient(
    buildboxcommon::CASClient &casClient, DigestCache *cache)
    : FindMissingBlobsClient(
          [&casClient](const std::vector<Digest> &digests) {
              return casClient.findMissingBlobs(digests);
          },
          cache)
{
}

FindMissingBlobsClient::FindMissingBlobsClient(
    const CasClientFindMissingBlobsFunction &findMissingBlobsFunction,
    DigestCache *cache)
    : d_cas_client_findmissingblobs_function(findMissingBlobsFunction),
      d_digests_in_remote_cache(cache)
{
}

grpc::Status FindMissingBlobsClient::findMissingBlobs(
    const std::vector<Digest> &digests,
    std::vector<Digest> *digestsFoundMissing)
{
    assert(digestsFoundMissing != nullptr);

    std::vector<Digest> request;
    if (!cacheEnabled()) {
        request = digests;
    }
    else {
        request.reserve(digests.size());
        // If caching is enabled, we first check our cache to avoid asking
        // about blobs that we assume are still present in the remote:
        for (const Digest &digest : digests) {
            if (!d_digests_in_remote_cache->hasDigest(digest)) {
                request.push_back(digest);
            }
        }
    }

    if (!request.empty()) {
        try {
            *digestsFoundMissing =
                d_cas_client_findmissingblobs_function(request);
        }
        catch (const buildboxcommon::GrpcError &e) {
            if (e.status.error_code() == grpc::StatusCode::CANCELLED) {
                BUILDBOX_LOG_INFO("FindMissingBlobs() cancelled: "
                                  << e.status.error_code() << ": "
                                  << e.status.error_message());
            }
            else {
                BUILDBOX_LOG_ERROR(
                    "FindMissingBlobs() failed. Remote returned "
                    << e.status.error_code() << ": "
                    << e.status.error_message());
            }
            return e.status;
        }

        if (cacheEnabled()) {
            addBlobsPresentInRemote(&request, digestsFoundMissing);
        }
    }

    return grpc::Status(grpc::StatusCode::OK, "");
}

void FindMissingBlobsClient::addBlobsPresentInRemote(
    std::vector<Digest> *queriedDigests, std::vector<Digest> *missingDigests)
{
    if (!cacheEnabled()) {
        return;
    }
    // Add entries in `{queried_digests} - {missing_digests}` to the cache.

    std::sort(queriedDigests->begin(), queriedDigests->end());
    std::sort(missingDigests->begin(), missingDigests->end());

    auto queriedIt = queriedDigests->cbegin();
    auto missingIt = missingDigests->cbegin();
    while (queriedIt != queriedDigests->cend() &&
           missingIt != missingDigests->cend()) {
        if (*queriedIt < *missingIt) {
            d_digests_in_remote_cache->addDigest(*queriedIt);
            queriedIt++;
        }
        else if (*missingIt < *queriedIt)
            missingIt++;
        else {
            queriedIt++;
            missingIt++;
        }
    }

    for (auto it = queriedIt; it != queriedDigests->cend(); it++) {
        d_digests_in_remote_cache->addDigest(*it);
    }
}
