#include <buildboxcasd_metricnames.h>
#include <buildboxcasd_proxynoupdatefixture.h>
#include <buildboxcasd_server.h>

#include <fcntl.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommonmetrics_countingmetricutil.h>
#include <buildboxcommonmetrics_countingmetricvalue.h>
#include <buildboxcommonmetrics_distributionmetric.h>
#include <buildboxcommonmetrics_metriccollectorfactory.h>
#include <buildboxcommonmetrics_testingutils.h>
#include <grpcpp/client_context.h>
#include <grpcpp/create_channel.h>
#include <grpcpp/impl/codegen/call.h>
#include <grpcpp/server_context.h>
#include <grpcpp/support/sync_stream.h>

using namespace buildboxcasd;
using namespace buildboxcommon;
using namespace buildboxcommon::buildboxcommonmetrics;

namespace {
template <typename ValueType>
std::unordered_map<std::string, ValueType> getMetricSnapshot()
{
    buildboxcommonmetrics::MetricCollector<ValueType> *collector =
        buildboxcommonmetrics::MetricCollectorFactory::getCollector<
            ValueType>();
    return collector->getSnapshot();
}
template <typename ValueType>
using MetricVector = std::vector<std::pair<std::string, ValueType>>;

template <typename ValueType>
void validateMetrics(
    const std::vector<std::pair<std::string, ValueType>> &expected_metrics)
{
    const auto actual_metrics = getMetricSnapshot<ValueType>();
    for (const auto &expected : expected_metrics) {
        const auto actual = actual_metrics.find(expected.first);
        if (actual == actual_metrics.end()) {
            FAIL() << "Expected metric " << expected.first
                   << " was not collected!";
        }
        else if (actual->second != expected.second) {
            FAIL() << "Metric " << expected.first
                   << " expected=" << expected.second.value()
                   << " actual=" << actual->second.value();
        }
    }
}

const auto digestFunctionInitializer = []() {
    buildboxcommon::DigestGenerator::init();
    return 0;
}();

} // namespace

class ProxyParamFixture : public ProxyNoUpdateFixture {
  protected:
    ProxyParamFixture()
    {
        // Pre-loading some content in the local and remote storages:
        local_storage->writeBlob(DigestGenerator::hash("data1"), "data1");
        local_storage->writeBlob(DigestGenerator::hash("data2"), "data2");

        remote_storage->writeBlob(DigestGenerator::hash("remotedata3"),
                                  "remotedata3");
        remote_storage->writeBlob(DigestGenerator::hash("remotedata4"),
                                  "remotedata4");

        // Clear out any metrics collected from the priming writeBlob calls
        buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    }
};

TEST_P(ProxyParamFixture, FindMissingBlobs)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    FindMissingBlobsRequest request;
    request.set_instance_name(instance_names[0]);

    const auto local_digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(local_digest));

    const auto remote_digest = DigestGenerator::hash("remotedata4");
    ASSERT_TRUE(remote_storage->hasBlob(remote_digest));

    const auto missing_digest = DigestGenerator::hash("non-existent");
    ASSERT_FALSE(local_storage->hasBlob(missing_digest));
    ASSERT_FALSE(remote_storage->hasBlob(missing_digest));

    const auto r1 = request.add_blob_digests();
    r1->CopyFrom(local_digest);
    const auto r2 = request.add_blob_digests();
    r2->CopyFrom(remote_digest);
    const auto r3 = request.add_blob_digests();
    r3->CopyFrom(missing_digest);

    FindMissingBlobsResponse response;
    cas_servicer->FindMissingBlobs(&server_context, &request, &response);

    ASSERT_EQ(response.missing_blob_digests().size(), 1);
    ASSERT_EQ(response.missing_blob_digests(0).size_bytes(),
              missing_digest.size_bytes());
    ASSERT_EQ(response.missing_blob_digests(0).hash(), missing_digest.hash());

    // The blob missing in the remote but available locally was NOT
    // uploaded:
    ASSERT_FALSE(remote_storage->hasBlob(local_digest));

    const std::vector<std::string> missingMetrics = {
        MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
        MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL};

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {};

    ASSERT_TRUE(
        buildboxcommonmetrics::allCollectedByNameWithValuesAndAllMissingByName<
            buildboxcommonmetrics::CountingMetricValue>(expected,
                                                        missingMetrics));

    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::allCollectedByNameWithValues<
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue>({
            {MetricNames::DISTRIBUTION_NAME_CAS_FIND_MISSING_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 local_digest.size_bytes())},
            {MetricNames::DISTRIBUTION_NAME_CAS_FIND_MISSING_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 remote_digest.size_bytes())},
            {MetricNames::DISTRIBUTION_NAME_CAS_FIND_MISSING_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 missing_digest.size_bytes())},
        }));
}

TEST_P(ProxyParamFixture, BatchUpdateBlobs)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    BatchUpdateBlobsRequest request;
    request.set_instance_name(instance_names[0]);

    const auto data4 = "data4";
    const auto data5 = "data5";

    const Digest d4 = DigestGenerator::hash(data4);
    const Digest d5 = DigestGenerator::hash(data5);

    auto entry1 = request.add_requests();
    entry1->mutable_digest()->CopyFrom(d4);
    entry1->set_data(data4);

    auto entry2 = request.add_requests();
    entry2->mutable_digest()->CopyFrom(d5);
    entry2->set_data(data5);

    BatchUpdateBlobsResponse response;
    cas_servicer->BatchUpdateBlobs(&server_context, &request, &response);

    // Response is valid:
    ASSERT_EQ(response.responses().size(), 2);

    std::set<std::string> digests;
    for (const auto &r : response.responses()) {
        ASSERT_EQ(r.status().code(), grpc::StatusCode::OK);
        digests.insert(r.digest().hash());
    }

    ASSERT_EQ(digests.count(d4.hash()), 1);
    ASSERT_EQ(digests.count(d5.hash()), 1);

    // And data was stored locally but not remotely
    ASSERT_TRUE(local_storage->hasBlob(d4));
    ASSERT_TRUE(local_storage->hasBlob(d5));
    ASSERT_FALSE(remote_storage->hasBlob(d4));
    ASSERT_FALSE(remote_storage->hasBlob(d5));

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_BATCH_UPDATE,
         buildboxcommonmetrics::CountingMetricValue(2)},
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(2)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(d4.size_bytes() +
                                                    d5.size_bytes())}};

    validateMetrics(expected);
}

TEST_P(ProxyParamFixture, BytestreamWrite)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    const std::string data = "data10";
    const Digest digest = DigestGenerator::hash(data);

    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    WriteRequest request;
    request.set_resource_name(resource_name_prefix +
                              "uploads/uuid-goes-here/blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));
    request.set_data(data);
    request.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request));
    ASSERT_TRUE(writer->WritesDone());
    ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);

    // Should have triggered no local read metrics,
    // but a local/remote write
    const std::vector<std::string> missingMetrics = {
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE,
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL,
        MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
        MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE};

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(1)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())}};

    ASSERT_TRUE(
        buildboxcommonmetrics::allCollectedByNameWithValuesAndAllMissingByName<
            buildboxcommonmetrics::CountingMetricValue>(expected,
                                                        missingMetrics));

    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue>(
            MetricNames::DISTRIBUTION_NAME_CAS_BYTESTREAM_WRITE_BLOB_SIZES,
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                digest.size_bytes())));

    ASSERT_EQ(response.committed_size(),
              readBlobFromLocalStorage(digest).size());

    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    ASSERT_EQ(readBlobFromLocalStorage(digest), data);
}

TEST_P(ProxyParamFixture, UploadMissingBlobs)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    UploadMissingBlobsRequest request;
    request.set_instance_name(instance_names[0]);

    const auto missing_digest = DigestGenerator::hash("notRemote");
    local_storage->writeBlob(missing_digest, "notRemote");
    ASSERT_TRUE(local_storage->hasBlob(missing_digest));
    ASSERT_FALSE(remote_storage->hasBlob(missing_digest));

    const auto r1 = request.add_blob_digests();
    r1->CopyFrom(missing_digest);

    UploadMissingBlobsResponse response;
    grpc::Status code = local_cas_servicer->UploadMissingBlobs(
        &server_context, &request, &response);
    ASSERT_FALSE(code.ok());
    ASSERT_TRUE(code.error_code() == grpc::StatusCode::UNIMPLEMENTED);

    // Make sure it was NOT uploaded to remote
    ASSERT_TRUE(local_storage->hasBlob(missing_digest));
    ASSERT_FALSE(remote_storage->hasBlob(missing_digest));

    const std::vector<std::string> missingMetrics = {
        MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
        MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE};

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {};

    ASSERT_TRUE(
        buildboxcommonmetrics::allCollectedByNameWithValuesAndAllMissingByName<
            buildboxcommonmetrics::CountingMetricValue>(expected,
                                                        missingMetrics));
}

INSTANTIATE_TEST_SUITE_P(CasProxyModeParam, ProxyParamFixture,
                         testing::Values(ENABLED, DISABLED));
