/*
 * 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_localraproxyinstance.h>
#include <buildboxcasd_metricnames.h>
#include <buildboxcasd_requestcontextmanager.h>
#include <buildboxcommon_grpcclient.h>
#include <buildboxcommon_grpcerror.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_stringutils.h>
#include <buildboxcommonmetrics_countingmetricutil.h>
#include <buildboxcommonmetrics_durationmetrictimer.h>
#include <buildboxcommonmetrics_metricguard.h>

using namespace buildboxcasd;
using namespace buildboxcommon;

LocalRaProxyInstance::LocalRaProxyInstance(
    const std::shared_ptr<FsLocalAssetStorage> &asset_storage,
    const buildboxcommon::ConnectionOptions &asset_endpoint)
    : RaInstance(), d_assetStorage(asset_storage),
      d_remoteInstanceName(asset_endpoint.d_instanceName)
{
    std::shared_ptr<buildboxcommon::GrpcClient> grpcClient =
        std::make_shared<buildboxcommon::GrpcClient>();
    grpcClient->init(asset_endpoint);
    RequestContextManager::configureGrpcClient(grpcClient.get());
    d_assetClient =
        std::make_unique<buildboxcommon::AssetClient>(std::move(grpcClient));
    d_assetClient->init();
}

grpc::Status LocalRaProxyInstance::FetchBlob(const FetchBlobRequest &request,
                                             FetchBlobResponse *response)
{
    buildboxcommon::buildboxcommonmetrics::MetricGuard<
        buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
        mt(MetricNames::TIMER_NAME_RA_FETCH_BLOB);

    const std::string urisString = StringUtils::join(
        std::vector<std::string>(request.uris().begin(), request.uris().end()),
        ", ");
    BUILDBOX_LOG_INFO("Fetching blob for URIs: " + urisString);

    if (d_assetStorage) {
        // First, try to fetch from local asset storage
        if (d_assetStorage->lookup(request, response)) {
            BUILDBOX_LOG_DEBUG("Blob found in local cache for URIs: " +
                               urisString);
            buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
                MetricNames::COUNTER_NAME_RA_FETCH_BLOB_HITS, 1);
            return grpc::Status::OK;
        }

        BUILDBOX_LOG_DEBUG(
            "Blob not found locally, fetching from remote for URIs: " +
            urisString);
    }

    // Not found locally, try remote
    FetchBlobRequest newRequest = request;
    newRequest.set_instance_name(d_remoteInstanceName);

    try {
        auto remoteResponse = d_assetClient->fetchBlob(newRequest);
        *response = remoteResponse;
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NAME_RA_FETCH_BLOB_HITS, 1);
    }
    catch (const GrpcError &e) {
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NAME_RA_FETCH_BLOB_MISSES, 1);
        throw;
    }

    if (response->status().code() == grpc::StatusCode::OK) {
        BUILDBOX_LOG_INFO("Blob fetched from remote, digest: " +
                          toString(response->blob_digest()) +
                          ", URI: " + response->uri() +
                          (d_assetStorage ? ", caching locally" : ""));

        if (d_assetStorage) {
            // Push to local cache
            PushBlobRequest pushRequest;
            pushRequest.add_uris(response->uri());
            pushRequest.mutable_qualifiers()->CopyFrom(request.qualifiers());
            pushRequest.mutable_expire_at()->CopyFrom(response->expires_at());
            pushRequest.mutable_blob_digest()->CopyFrom(
                response->blob_digest());

            d_assetStorage->insert(pushRequest);
        }

        return grpc::Status::OK;
    }
    else {
        BUILDBOX_LOG_WARNING(
            "Remote fetch failed with status: " << response->status().code());
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NAME_RA_FETCH_BLOB_MISSES, 1);
        return grpc::Status(
            static_cast<grpc::StatusCode>(response->status().code()),
            response->status().message());
    }
}

grpc::Status
LocalRaProxyInstance::FetchDirectory(const FetchDirectoryRequest &request,
                                     FetchDirectoryResponse *response)
{
    buildboxcommon::buildboxcommonmetrics::MetricGuard<
        buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
        mt(MetricNames::TIMER_NAME_RA_FETCH_DIRECTORY);

    const std::string urisString = StringUtils::join(
        std::vector<std::string>(request.uris().begin(), request.uris().end()),
        ", ");
    BUILDBOX_LOG_INFO("Fetching directory for URIs: " + urisString);

    if (d_assetStorage) {
        // First, try to fetch from local asset storage
        if (d_assetStorage->lookup(request, response)) {
            BUILDBOX_LOG_DEBUG("Directory found in local cache for URIs: " +
                               urisString);
            buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
                MetricNames::COUNTER_NAME_RA_FETCH_DIRECTORY_HITS, 1);
            return grpc::Status::OK;
        }

        BUILDBOX_LOG_DEBUG(
            "Directory not found locally, fetching from remote for URIs: " +
            urisString);
    }

    // Not found locally, try remote
    FetchDirectoryRequest newRequest = request;
    newRequest.set_instance_name(d_remoteInstanceName);

    try {
        auto remoteResponse = d_assetClient->fetchDirectory(newRequest);
        *response = remoteResponse;
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NAME_RA_FETCH_DIRECTORY_HITS, 1);
    }
    catch (const GrpcError &e) {
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NAME_RA_FETCH_DIRECTORY_MISSES, 1);
        throw;
    }

    if (response->status().code() == grpc::StatusCode::OK) {
        BUILDBOX_LOG_INFO("Directory fetched from remote, digest: " +
                          toString(response->root_directory_digest()) +
                          ", URI: " + response->uri() +
                          (d_assetStorage ? ", caching locally" : ""));

        if (d_assetStorage) {
            // Push to local cache
            PushDirectoryRequest pushRequest;
            pushRequest.add_uris(response->uri());
            pushRequest.mutable_qualifiers()->CopyFrom(request.qualifiers());
            pushRequest.mutable_expire_at()->CopyFrom(response->expires_at());
            pushRequest.mutable_root_directory_digest()->CopyFrom(
                response->root_directory_digest());

            d_assetStorage->insert(pushRequest);
        }

        return grpc::Status::OK;
    }
    else {
        BUILDBOX_LOG_WARNING("Remote directory fetch failed with status: "
                             << response->status().code());
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NAME_RA_FETCH_DIRECTORY_MISSES, 1);
        return grpc::Status(
            static_cast<grpc::StatusCode>(response->status().code()),
            response->status().message());
    }
}

grpc::Status LocalRaProxyInstance::PushBlob(const PushBlobRequest &request,
                                            PushBlobResponse *response)
{
    buildboxcommon::buildboxcommonmetrics::MetricGuard<
        buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
        mt(MetricNames::TIMER_NAME_RA_PUSH_BLOB);

    const std::string urisString = StringUtils::join(
        std::vector<std::string>(request.uris().begin(), request.uris().end()),
        ", ");
    BUILDBOX_LOG_INFO("Pushing blob for URIs: " + urisString +
                      ", digest: " + toString(request.blob_digest()));

    // Push to remote storage first
    PushBlobRequest newRequest = request;
    newRequest.set_instance_name(d_remoteInstanceName);

    BUILDBOX_LOG_DEBUG("Pushing blob to remote for digest: " +
                       toString(request.blob_digest()));
    PushBlobResponse pushResponse = d_assetClient->pushBlob(newRequest);
    if (response != nullptr) {
        *response = pushResponse;
    }

    if (d_assetStorage) {
        // Push to local storage
        BUILDBOX_LOG_DEBUG("Pushing blob to local storage for digest: " +
                           toString(request.blob_digest()));
        d_assetStorage->insert(request);
    }

    return grpc::Status::OK;
}

grpc::Status
LocalRaProxyInstance::PushDirectory(const PushDirectoryRequest &request,
                                    PushDirectoryResponse *response)
{
    buildboxcommon::buildboxcommonmetrics::MetricGuard<
        buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
        mt(MetricNames::TIMER_NAME_RA_PUSH_DIRECTORY);

    const std::string urisString = StringUtils::join(
        std::vector<std::string>(request.uris().begin(), request.uris().end()),
        ", ");
    BUILDBOX_LOG_INFO(
        "Pushing directory for URIs: " + urisString +
        ", digest: " + toString(request.root_directory_digest()));

    // Push to remote storage first
    PushDirectoryRequest newRequest = request;
    newRequest.set_instance_name(d_remoteInstanceName);

    BUILDBOX_LOG_DEBUG("Pushing directory to remote for digest: " +
                       toString(request.root_directory_digest()));
    PushDirectoryResponse pushResponse =
        d_assetClient->pushDirectory(newRequest);
    if (response != nullptr) {
        *response = pushResponse;
    }

    if (d_assetStorage) {
        // Push to local storage
        BUILDBOX_LOG_DEBUG("Pushing directory to local storage for digest: " +
                           toString(request.root_directory_digest()));
        d_assetStorage->insert(request);
    }

    return grpc::Status::OK;
}
