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

#include <buildboxcommon_connectionoptions.h>
#include <buildboxcommon_grpcerror.h>

#include <buildboxcommon_protos.h>

#include <chrono>
#include <string>
#include <utility>

namespace buildboxcommon {

using namespace build::bazel::remote::asset::v1;

void AssetClient::init()
{
    const std::shared_ptr<grpc::Channel> channel = d_grpcClient->channel();

    std::shared_ptr<Fetch::Stub> fetchClient = Fetch::NewStub(channel);
    std::shared_ptr<Push::Stub> pushClient = Push::NewStub(channel);
    init(fetchClient, pushClient);
}

void AssetClient::init(std::shared_ptr<Fetch::StubInterface> fetchClient,
                       std::shared_ptr<Push::StubInterface> pushClient)
{
    this->d_fetchClient = std::move(fetchClient);
    this->d_pushClient = std::move(pushClient);
}

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

FetchBlobResponse
AssetClient::fetchBlob(const FetchBlobRequest &request,
                       GrpcClient::RequestStats *requestStats)
{
    FetchBlobResponse response;
    d_grpcClient->issueRequest(
        [&](grpc::ClientContext &context) {
            return d_fetchClient->FetchBlob(&context, request, &response);
        },
        __func__, requestStats);

    return response;
}

AssetClient::FetchResult AssetClient::fetchBlob(
    const std::vector<std::string> &uris,
    const std::vector<std::pair<std::string, std::string>> &qualifiers,
    GrpcClient::RequestStats *requestStats)
{
    FetchBlobRequest remoteRequest;

    for (const std::string &uri : uris) {
        remoteRequest.add_uris(uri);
    }

    for (const std::pair<std::string, std::string> &qualifier : qualifiers) {
        Qualifier *const ptr = remoteRequest.add_qualifiers();
        ptr->set_name(qualifier.first);
        ptr->set_value(qualifier.second);
    }

    remoteRequest.set_instance_name(d_grpcClient->instanceName());

    FetchBlobResponse response = fetchBlob(remoteRequest, requestStats);

    std::vector<std::pair<std::string, std::string>> convertedQualifiers;
    convertedQualifiers.reserve(response.qualifiers().size());
    for (const Qualifier &qualifier : response.qualifiers()) {
        convertedQualifiers.emplace_back(qualifier.name(), qualifier.value());
    }

    return FetchResult{
        response.uri(), std::move(convertedQualifiers),
        std::chrono::time_point<std::chrono::steady_clock>(
            std::chrono::duration_cast<std::chrono::steady_clock::duration>(
                std::chrono::seconds(response.expires_at().seconds()) +
                std::chrono::nanoseconds(response.expires_at().nanos()))),
        response.blob_digest()};
}

FetchDirectoryResponse
AssetClient::fetchDirectory(const FetchDirectoryRequest &request,
                            GrpcClient::RequestStats *requestStats)
{
    FetchDirectoryResponse response;
    d_grpcClient->issueRequest(
        [&](grpc::ClientContext &context) {
            return d_fetchClient->FetchDirectory(&context, request, &response);
        },
        __func__, requestStats);

    return response;
}

AssetClient::FetchResult AssetClient::fetchDirectory(
    const std::vector<std::string> &uris,
    const std::vector<std::pair<std::string, std::string>> &qualifiers,
    GrpcClient::RequestStats *requestStats)
{
    FetchDirectoryRequest remoteRequest;
    for (const std::string &uri : uris) {
        remoteRequest.add_uris(uri);
    }

    for (const std::pair<std::string, std::string> &qualifier : qualifiers) {
        Qualifier *const ptr = remoteRequest.add_qualifiers();
        ptr->set_name(qualifier.first);
        ptr->set_value(qualifier.second);
    }

    remoteRequest.set_instance_name(d_grpcClient->instanceName());

    FetchDirectoryResponse response =
        fetchDirectory(remoteRequest, requestStats);

    std::vector<std::pair<std::string, std::string>> convertedQualifiers;
    convertedQualifiers.reserve(response.qualifiers().size());
    for (const Qualifier &qualifier : response.qualifiers()) {
        convertedQualifiers.emplace_back(qualifier.name(), qualifier.value());
    }

    return FetchResult{
        response.uri(), std::move(convertedQualifiers),
        std::chrono::time_point<std::chrono::steady_clock>(
            std::chrono::duration_cast<std::chrono::steady_clock::duration>(
                std::chrono::seconds(response.expires_at().seconds()) +
                std::chrono::nanoseconds(response.expires_at().nanos()))),
        response.root_directory_digest()};
}

std::optional<FetchBlobResponse>
AssetClient::tryFetchBlob(const FetchBlobRequest &request,
                          GrpcClient::RequestStats *requestStats)
{
    FetchBlobResponse response;
    auto retrier = d_grpcClient->makeRetrier(
        [&](grpc::ClientContext &context) {
            return d_fetchClient->FetchBlob(&context, request, &response);
        },
        __func__);

    // Treat NOT_FOUND as OK (for cache misses)
    retrier.addOkStatusCode(grpc::StatusCode::NOT_FOUND);

    retrier.issueRequest();

    if (requestStats != nullptr) {
        requestStats->d_grpcRetryCount += retrier.retryAttempts();
    }

    // Return nullopt for NOT_FOUND
    if (retrier.status().error_code() == grpc::StatusCode::NOT_FOUND) {
        return std::nullopt;
    }

    // For tryFetch methods, we don't throw exceptions on NOT_FOUND
    if (!retrier.status().ok()) {
        GrpcError::throwGrpcError(retrier.status());
    }

    return response;
}

std::optional<AssetClient::FetchResult> AssetClient::tryFetchBlob(
    const std::vector<std::string> &uris,
    const std::vector<std::pair<std::string, std::string>> &qualifiers,
    GrpcClient::RequestStats *requestStats)
{
    FetchBlobRequest remoteRequest;

    for (const std::string &uri : uris) {
        remoteRequest.add_uris(uri);
    }

    for (const std::pair<std::string, std::string> &qualifier : qualifiers) {
        Qualifier *const ptr = remoteRequest.add_qualifiers();
        ptr->set_name(qualifier.first);
        ptr->set_value(qualifier.second);
    }

    remoteRequest.set_instance_name(d_grpcClient->instanceName());

    FetchBlobResponse response;
    auto retrier = d_grpcClient->makeRetrier(
        [&](grpc::ClientContext &context) {
            return d_fetchClient->FetchBlob(&context, remoteRequest,
                                            &response);
        },
        __func__);

    // Treat NOT_FOUND as OK (for cache misses)
    retrier.addOkStatusCode(grpc::StatusCode::NOT_FOUND);

    retrier.issueRequest();

    if (requestStats != nullptr) {
        requestStats->d_grpcRetryCount += retrier.retryAttempts();
    }

    // Return nullopt for NOT_FOUND
    if (retrier.status().error_code() == grpc::StatusCode::NOT_FOUND) {
        return std::nullopt;
    }

    if (!retrier.status().ok()) {
        GrpcError::throwGrpcError(retrier.status());
    }

    std::vector<std::pair<std::string, std::string>> convertedQualifiers;
    convertedQualifiers.reserve(response.qualifiers().size());
    for (const Qualifier &qualifier : response.qualifiers()) {
        convertedQualifiers.emplace_back(qualifier.name(), qualifier.value());
    }

    return FetchResult{
        response.uri(), std::move(convertedQualifiers),
        std::chrono::time_point<std::chrono::steady_clock>(
            std::chrono::duration_cast<std::chrono::steady_clock::duration>(
                std::chrono::seconds(response.expires_at().seconds()) +
                std::chrono::nanoseconds(response.expires_at().nanos()))),
        response.blob_digest()};
}

PushBlobResponse AssetClient::pushBlob(const PushBlobRequest &request,
                                       GrpcClient::RequestStats *requestStats)
{
    PushBlobResponse response;
    d_grpcClient->issueRequest(
        [&](grpc::ClientContext &context) {
            return d_pushClient->PushBlob(&context, request, &response);
        },
        __func__, requestStats);

    return response;
}

void AssetClient::pushBlob(
    const std::vector<std::string> &uris,
    const std::vector<std::pair<std::string, std::string>> &qualifiers,
    const Digest &digest, const std::vector<Digest> &referencedBlobs,
    const std::vector<Digest> &referencedDirectories,
    GrpcClient::RequestStats *requestStats)
{
    PushBlobRequest remoteRequest;
    for (const std::string &uri : uris) {
        remoteRequest.add_uris(uri);
    }

    for (const std::pair<std::string, std::string> &qualifier : qualifiers) {
        Qualifier *const ptr = remoteRequest.add_qualifiers();
        ptr->set_name(qualifier.first);
        ptr->set_value(qualifier.second);
    }

    *remoteRequest.mutable_blob_digest() = digest;

    // Add referenced blobs
    for (const Digest &refBlob : referencedBlobs) {
        *remoteRequest.add_references_blobs() = refBlob;
    }

    // Add referenced directories
    for (const Digest &refDir : referencedDirectories) {
        *remoteRequest.add_references_directories() = refDir;
    }

    remoteRequest.set_instance_name(d_grpcClient->instanceName());

    pushBlob(remoteRequest, requestStats);
}

PushDirectoryResponse
AssetClient::pushDirectory(const PushDirectoryRequest &request,
                           GrpcClient::RequestStats *requestStats)
{
    PushDirectoryResponse response;
    d_grpcClient->issueRequest(
        [&](grpc::ClientContext &context) {
            return d_pushClient->PushDirectory(&context, request, &response);
        },
        __func__, requestStats);
    return response;
}

void AssetClient::pushDirectory(
    const std::vector<std::string> &uris,
    const std::vector<std::pair<std::string, std::string>> &qualifiers,
    const Digest &rootDirDigest, const std::vector<Digest> &referencedBlobs,
    const std::vector<Digest> &referencedDirectories,
    GrpcClient::RequestStats *requestStats)
{
    PushDirectoryRequest remoteRequest;
    for (const std::string &uri : uris) {
        remoteRequest.add_uris(uri);
    }

    for (const std::pair<std::string, std::string> &qualifier : qualifiers) {
        Qualifier *const ptr = remoteRequest.add_qualifiers();
        ptr->set_name(qualifier.first);
        ptr->set_value(qualifier.second);
    }

    *remoteRequest.mutable_root_directory_digest() = rootDirDigest;

    // Add referenced blobs
    for (const Digest &refBlob : referencedBlobs) {
        *remoteRequest.add_references_blobs() = refBlob;
    }

    // Add referenced directories
    for (const Digest &refDir : referencedDirectories) {
        *remoteRequest.add_references_directories() = refDir;
    }

    remoteRequest.set_instance_name(d_grpcClient->instanceName());

    pushDirectory(remoteRequest, requestStats);
}
} // namespace buildboxcommon
