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

#ifndef INCLUDED_BUILDBOXCASD_TESTSERVER_H
#define INCLUDED_BUILDBOXCASD_TESTSERVER_H

#include <buildboxcasd_server.h>
#include <gmock/gmock.h>
#include <google/protobuf/duration.pb.h>
#include <gtest/gtest.h>
#include <ThreadPool.h>

#include <buildboxcasd_fslocalactionstorage.h>
#include <buildboxcommon_casclient.h>
#include <buildboxcommon_fslocalcas.h>
#include <buildboxcommon_remoteexecutionclient.h>
#include <buildboxcommon_temporarydirectory.h>
#include <grpcpp/create_channel.h>
#include <grpcpp/impl/codegen/call.h>
#include <grpcpp/support/sync_stream.h>

#include <memory>
#include <tuple>

using namespace buildboxcasd;
using namespace buildboxcommon;

const std::vector<std::string> TEST_INSTANCE_NAMES{"", "test"};

enum CasProxyModeParam {
    STATIC_INSTANCE,
    DYNAMIC_INSTANCE,
};

enum CasProxyFindMissingBlobsCacheParam {
    FINDMISSINGBLOBS_CACHE_ON,
    FINDMISSINGBLOBS_CACHE_OFF
};

enum AcProxyValidationCacheParam {
    AC_VALIDATION_CACHE_ON,
    AC_VALIDATION_CACHE_OFF
};

enum CasProxyThreadedParam {
    NON_THREADED,
    THREADED,
};

class CasProxyModeFixture
    : public ::testing::TestWithParam<std::tuple<
          CasProxyModeParam, CasProxyFindMissingBlobsCacheParam, std::string,
          CasProxyThreadedParam, AcProxyValidationCacheParam>> {
    /**
     * This fixture can be used to test the functionality of buildbox-casd in
     * proxy mode.
     *
     * It sets up two instances: one running as a proxy and another as a
     * regular server ("remote"). Both instances use sockets for receiving
     * connections.
     *
     *                  (`cas_client`)
     * [ Proxy Server ] --------------> [ Remote Server ]
     *        |                                 |
     * (`local_storage`)                (`remote_storage`)
     *
     *
     * The gtest parameters allow to test static instances (a proxy connected
     * to a remote on instantiation time) and dynamic ones (a connection opened
     * on request). They also allow to enable or disable the FindMissingBlobs()
     * cache in the proxy.
     */
  protected:
    CasProxyModeFixture()
        : TEST_PROXY_SERVER_ADDRESS("unix://" +
                                    std::string(sockets_directory.name()) +
                                    "/proxy.sock"),
          TEST_REMOTE_SERVER_ADDRESS("unix://" +
                                     std::string(sockets_directory.name()) +
                                     "/remote.sock"),
          local_storage(std::make_shared<FsLocalCas>(
              local_storage_root_directory.name(), true)),
          // (allow moving files to local_storage)
          local_asset_storage(std::make_shared<FsLocalAssetStorage>(
              local_storage_root_directory.name())),
          local_action_storage(std::make_shared<FsLocalActionStorage>(
              local_storage_root_directory.name())),
          grpc_client(std::make_shared<buildboxcommon::GrpcClient>()),
          cas_client(std::make_shared<buildboxcommon::CASClient>(grpc_client)),
          remote_storage(std::make_shared<FsLocalCas>(
              remote_storage_root_directory.name())),
          remote_asset_storage(std::make_shared<FsLocalAssetStorage>(
              remote_storage_root_directory.name())),
          remote_action_storage(std::make_shared<FsLocalActionStorage>(
              remote_storage_root_directory.name())),
          remote_server(std::make_shared<Server>(
              remote_storage, remote_asset_storage, remote_action_storage)),
          proxy_cas_channel(grpc::CreateChannel(
              TEST_PROXY_SERVER_ADDRESS, grpc::InsecureChannelCredentials())),
          cas_stub(ContentAddressableStorage::NewStub(proxy_cas_channel)),
          cas_bytestream_stub(ByteStream::NewStub(proxy_cas_channel)),
          localcas_stub(
              LocalContentAddressableStorage::NewStub(proxy_cas_channel))
    {
        const bool dynamic_instance =
            (std::get<0>(GetParam()) == DYNAMIC_INSTANCE);

        const int proxy_findmissingblobs_cache_ttl_seconds =
            (std::get<1>(GetParam()) == FINDMISSINGBLOBS_CACHE_ON)
                ? 5 * 60 // Let entries live throughout a unit test.
                : 0;
        const int proxyAcValidationCheckTtlSeconds =
            (std::get<4>(GetParam()) == AC_VALIDATION_CACHE_ON) ? 5 * 60 : 0;
        instance_name = std::get<2>(GetParam());
        if (!instance_name.empty()) {
            resource_name_prefix = instance_name + "/";
        }

        if (std::get<3>(GetParam()) == THREADED) {
            threadPool = std::make_shared<ThreadPool>(2, "");
        }

        // Building and starting the remote server:
        for (const auto &instance : TEST_INSTANCE_NAMES) {
            remote_server->addLocalServerInstance(instance);
        }

        remote_server->addListeningPort(TEST_REMOTE_SERVER_ADDRESS);
        remote_server->start();

        // Setting up ConnectionOptions used by the proxy to connect with the
        // remote server:
        buildboxcommon::ConnectionOptions connection_options;
        std::string url = TEST_REMOTE_SERVER_ADDRESS;
        connection_options.setUrl(url);
        connection_options.setInstanceName("");
        connection_options.setRetryDelay("1");
        connection_options.setRetryLimit("1");
        connection_options.setRequestTimeout("10");

        if (!dynamic_instance) {
            std::unordered_map<std::string, std::string> instanceMappings;
            local_server = std::make_shared<Server>(
                local_storage, local_asset_storage, local_action_storage,
                nullptr, nullptr, nullptr, connection_options,
                std::optional<ConnectionOptions>(connection_options),
                std::optional<ConnectionOptions>(connection_options),
                std::optional<ConnectionOptions>(connection_options),
                TEST_INSTANCE_NAMES, instanceMappings, false,
                proxy_findmissingblobs_cache_ttl_seconds,
                proxyAcValidationCheckTtlSeconds);
        }
        else {
            local_server = std::make_shared<Server>(
                local_storage, local_asset_storage, local_action_storage);
        }

        cas_servicer = std::dynamic_pointer_cast<CasRemoteExecutionServicer>(
            local_server->remoteExecutionCasServicer());

        ac_servicer = std::dynamic_pointer_cast<ActionCacheServicer>(
            local_server->localAcServicer());

        remote_ac_servicer = std::dynamic_pointer_cast<ActionCacheServicer>(
            remote_server->localAcServicer());

        // Building and starting the proxy server:
        local_server->addListeningPort(TEST_PROXY_SERVER_ADDRESS);
        local_server->start();

        if (dynamic_instance) {
            grpc::ClientContext client_context;
            GetInstanceNameForRemotesRequest request;
            GetInstanceNameForRemotesResponse response;

            request.mutable_content_addressable_storage()->set_url(url);
            request.mutable_content_addressable_storage()->set_retry_limit(1);
            google::protobuf::Duration retry_delay;
            retry_delay.set_seconds(0);
            // 1 ms
            retry_delay.set_nanos(1000 * 1000);
            request.mutable_content_addressable_storage()
                ->mutable_retry_delay()
                ->CopyFrom(retry_delay);
            request.mutable_remote_asset()->set_url(url);
            request.mutable_remote_asset()->set_retry_limit(1);
            request.mutable_remote_asset()->mutable_retry_delay()->CopyFrom(
                retry_delay);
            localcas_stub->GetInstanceNameForRemotes(&client_context, request,
                                                     &response);

            instance_name = response.instance_name();
            resource_name_prefix = instance_name + "/";
        }

        // Initialize the CAS Client for direct access to the remote server
        grpc_client->init(connection_options);
        cas_client->init();
    }

    ~CasProxyModeFixture()
    {
        stopRemoteServer(); // make sure non-default timeout is used
    }

    inline std::string readBlobFromLocalStorage(const Digest &digest) const
    {
        const auto blob_ptr = local_storage->readBlob(digest);
        if (blob_ptr == nullptr) {
            throw std::invalid_argument("Digest " + toString(digest) +
                                        " is not present in local storage.");
        }

        return *blob_ptr;
    }

    inline std::string readBlobFromRemoteStorage(const Digest &digest) const
    {
        const auto blob_ptr = remote_storage->readBlob(digest);
        if (blob_ptr == nullptr) {
            throw std::invalid_argument("Digest " + toString(digest) +
                                        " is not present in remote storage.");
        }

        return *blob_ptr;
    }

    void stopRemoteServer()
    {
        remote_server->shutdown(gpr_time_0(GPR_TIMESPAN));
        remote_server->wait();
    }

    // Directory to contain the sockets where the services will listen:
    buildboxcommon::TemporaryDirectory sockets_directory;
    const std::string TEST_PROXY_SERVER_ADDRESS, TEST_REMOTE_SERVER_ADDRESS;

    // Local proxy:
    buildboxcommon::TemporaryDirectory local_storage_root_directory;
    std::shared_ptr<FsLocalCas> local_storage;
    std::shared_ptr<FsLocalAssetStorage> local_asset_storage;
    std::shared_ptr<FsLocalActionStorage> local_action_storage;
    std::shared_ptr<buildboxcommon::GrpcClient> grpc_client;
    std::shared_ptr<buildboxcommon::CASClient> cas_client;
    std::shared_ptr<Server> local_server;
    std::shared_ptr<CasRemoteExecutionServicer> cas_servicer;
    std::shared_ptr<ActionCacheServicer> ac_servicer;
    std::string instance_name;
    std::string resource_name_prefix;

    // Remote CAS server:
    buildboxcommon::TemporaryDirectory remote_storage_root_directory;
    std::shared_ptr<FsLocalCas> remote_storage;
    std::shared_ptr<FsLocalAssetStorage> remote_asset_storage;
    std::shared_ptr<FsLocalActionStorage> remote_action_storage;
    std::shared_ptr<Server> remote_server;
    std::shared_ptr<ActionCacheServicer> remote_ac_servicer;

    // Thread pools
    std::shared_ptr<ThreadPool> threadPool;

    // Contexts:
    grpc::ClientContext client_context;
    grpc::ServerContext server_context;

    // gRPC stubs (allow to invoke gRPC calls):
    std::shared_ptr<grpc::Channel> proxy_cas_channel;
    std::unique_ptr<ContentAddressableStorage::Stub> cas_stub;
    std::unique_ptr<ByteStream::Stub> cas_bytestream_stub;
    std::unique_ptr<LocalContentAddressableStorage::Stub> localcas_stub;
};

#endif
