/*
 * 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_HYBRIDFALLBACKEXECUTIONSERVICEFIXTURE_H
#define INCLUDED_BUILDBOXCASD_HYBRIDFALLBACKEXECUTIONSERVICEFIXTURE_H

#include <buildboxcasd_fslocalactionstorage.h>
#include <buildboxcasd_fslocalassetstorage.h>
#include <buildboxcasd_localexecutioninstance.h>
#include <buildboxcasd_localexecutionscheduler.h>
#include <buildboxcasd_server.h>
#include <buildboxcommon_fslocalcas.h>
#include <buildboxcommon_grpctestserver.h>
#include <buildboxcommon_remoteexecutionclient.h>

#include <google/rpc/code.pb.h>
#include <gtest/gtest.h>
#include <unistd.h>

using namespace buildboxcasd;
using namespace buildboxcommon;
using namespace google::longrunning;

/***
 * Fixture to support testing buildbox-casd as a hybrid + fallback Execution
 * service
 */
class HybridFallbackExecutionServiceTestFixture : public ::testing::Test {
  protected:
    GrpcTestServer testRemoteServer;

    buildboxcommon::TemporaryDirectory socketDirectory;

    const std::string TEST_REMOTE_SERVER_ADDRESS;
    const std::string TEST_PROXY_SERVER_ADDRESS;

    buildboxcommon::TemporaryDirectory remoteStorageRootDirectory;
    std::shared_ptr<FsLocalCas> remoteCasStorage;
    std::shared_ptr<FsLocalAssetStorage> remoteAssetStorage;
    std::shared_ptr<FsLocalActionStorage> remoteActionStorage;
    std::shared_ptr<Server> remoteServer;

    buildboxcommon::TemporaryDirectory proxyStorageRootDirectory;
    std::shared_ptr<FsLocalCas> proxyCasStorage;
    std::shared_ptr<FsLocalAssetStorage> proxyAssetStorage;
    std::shared_ptr<FsLocalActionStorage> proxyActionStorage;
    std::shared_ptr<Server> proxyServer;

    std::shared_ptr<grpc::Channel> channel;
    std::shared_ptr<ByteStream::StubInterface> byteStreamStub;
    std::shared_ptr<ContentAddressableStorage::StubInterface> casStub;
    std::shared_ptr<Execution::StubInterface> executionStub;
    std::shared_ptr<google::longrunning::Operations::StubInterface>
        operationsStub;
    std::shared_ptr<GrpcClient> grpcClient;
    std::shared_ptr<CASClient> casClient;
    std::shared_ptr<RemoteExecutionClient> execClient;

    const std::string remoteInstanceName = "testremote";
    const std::string proxyInstanceName = "testproxy";

    HybridFallbackExecutionServiceTestFixture()
        : TEST_REMOTE_SERVER_ADDRESS("unix://" +
                                     std::string(socketDirectory.name()) +
                                     "/remote.sock"),
          TEST_PROXY_SERVER_ADDRESS(
              "unix://" + std::string(socketDirectory.name()) + "/proxy.sock"),
          grpcClient(std::make_shared<buildboxcommon::GrpcClient>()),
          casClient(std::make_shared<CASClient>(grpcClient)),
          execClient(
              std::make_shared<RemoteExecutionClient>(grpcClient, nullptr))
    {
        const char *runnerCommand = getenv("BUILDBOX_RUN");
        EXPECT_NE(runnerCommand, nullptr);
        const std::vector<std::string> extraRunArgs = {};

        // Create a server instance to act as remote server but not for
        // execution, for execution we will use testRemoteServer
        remoteCasStorage =
            std::make_shared<FsLocalCas>(remoteStorageRootDirectory.name());
        remoteAssetStorage = std::make_shared<FsLocalAssetStorage>(
            remoteStorageRootDirectory.name());
        remoteActionStorage = std::make_shared<FsLocalActionStorage>(
            remoteStorageRootDirectory.name());

        remoteServer = std::make_shared<Server>(
            remoteCasStorage, remoteAssetStorage, remoteActionStorage);
        remoteServer->addLocalServerInstance(remoteInstanceName);
        remoteServer->addListeningPort(TEST_REMOTE_SERVER_ADDRESS);
        remoteServer->start();

        // Create a server instance to act as proxy to the remote server,
        // also supporting local execution.
        ConnectionOptions remoteServerEndpoint;
        remoteServerEndpoint.setInstanceName(remoteInstanceName);
        remoteServerEndpoint.setUrl(TEST_REMOTE_SERVER_ADDRESS);
        remoteServerEndpoint.setRetryDelay("10");
        remoteServerEndpoint.setRetryLimit("1");

        // This is used only for execution and points to testRemoteServer
        ConnectionOptions fakeRemoteServerEndpoint;
        fakeRemoteServerEndpoint.setInstanceName(remoteInstanceName);
        fakeRemoteServerEndpoint.setUrl(testRemoteServer.url());
        remoteServerEndpoint.setRetryDelay("10");
        fakeRemoteServerEndpoint.setRetryLimit("0");

        proxyCasStorage =
            std::make_shared<FsLocalCas>(proxyStorageRootDirectory.name());
        proxyAssetStorage = std::make_shared<FsLocalAssetStorage>(
            proxyStorageRootDirectory.name());
        proxyActionStorage = std::make_shared<FsLocalActionStorage>(
            proxyStorageRootDirectory.name());

        auto localScheduler = std::make_shared<LocalExecutionScheduler>(
            TEST_PROXY_SERVER_ADDRESS, runnerCommand, extraRunArgs,
            /* maxJobs */ 2);

        proxyServer =
            std::make_shared<Server>(proxyCasStorage, proxyAssetStorage,
                                     proxyActionStorage, localScheduler);
        proxyServer->addProxyInstance(
            proxyInstanceName, /* casEndpoint */ remoteServerEndpoint,
            /* raEndpoint */ remoteServerEndpoint,
            /* acEndpoint */ remoteServerEndpoint,
            /* executionEndpoint */ fakeRemoteServerEndpoint,
            /* readOnlyRemote */ false, /* proxyFindmissingblobsCacheTtl */ 0,
            /* proxyAcValidationCacheTtl */ 0,
            /*execHybridQueueLimit*/ 0,
            /*enableFallbackToLocalExecution */ true);
        proxyServer->addListeningPort(TEST_PROXY_SERVER_ADDRESS);
        proxyServer->start();

        // Initialize a gRPC client to the proxy
        channel = grpc::CreateChannel(TEST_PROXY_SERVER_ADDRESS,
                                      grpc::InsecureChannelCredentials());
        byteStreamStub = ByteStream::NewStub(channel);
        casStub = ContentAddressableStorage::NewStub(channel);
        executionStub = Execution::NewStub(channel);
        operationsStub = Operations::NewStub(channel);
        grpcClient->setInstanceName(proxyInstanceName);
        casClient->init(byteStreamStub, casStub, nullptr, nullptr);
        execClient->init(executionStub, nullptr, operationsStub);
    }

    ~HybridFallbackExecutionServiceTestFixture() {}

    Digest generateSimpleShellAction(const std::string &commandString)
    {
        Action action;
        // Empty input root
        action.mutable_input_root_digest()->CopyFrom(
            DigestGenerator::hash(""));
        Command command;
        command.add_arguments("/bin/sh");
        command.add_arguments("-c");
        command.add_arguments(commandString);

        const auto commandDigest = casClient->uploadMessage(command);
        action.mutable_command_digest()->CopyFrom(commandDigest);
        const auto actionDigest = casClient->uploadMessage(action);
        return actionDigest;
    }

    ExecuteRequest getExpectedExecuteRequest(const Digest &actionDigest)
    {
        ExecuteRequest executeRequest;
        executeRequest.set_instance_name(remoteInstanceName);
        *executeRequest.mutable_action_digest() = actionDigest;
        return executeRequest;
    }

    WaitExecutionRequest
    getExpectedWaitExecutionRequest(const Operation &operation)
    {
        WaitExecutionRequest waitExecutionRequest;
        waitExecutionRequest.set_name(operation.name());
        return waitExecutionRequest;
    }

    Operation getExpectedOperation()
    {
        Operation operation;
        operation.set_name(remoteInstanceName + "#operation-name");
        return operation;
    }

    void completeOperation(Operation *operation)
    {
        operation->set_done(true);
        ExecuteResponse executeResponse;
        ActionResult actionResult;
        actionResult.set_exit_code(0);
        executeResponse.mutable_result()->CopyFrom(actionResult);
        executeResponse.mutable_status()->set_code(google::rpc::Code::OK);
        operation->mutable_response()->PackFrom(executeResponse);
    }

    void internalErrorOperationExecuteResponse(Operation *operation)
    {
        operation->set_done(true);
        ExecuteResponse executeResponse;
        executeResponse.mutable_status()->set_code(
            google::rpc::Code::INTERNAL);
        operation->mutable_response()->PackFrom(executeResponse);
    }

    ExecutedActionMetadata expectOperationCancelled(const Operation &operation)
    {
        EXPECT_TRUE(operation.done());

        ExecuteResponse executeResponse;
        EXPECT_TRUE(operation.response().UnpackTo(&executeResponse));

        EXPECT_EQ(executeResponse.status().code(),
                  google::rpc::Code::CANCELLED);

        const ActionResult actionResult = executeResponse.result();
        return actionResult.execution_metadata();
    }

    ExecutedActionMetadata
    expectOperationSuccessful(const Operation &operation)
    {
        EXPECT_TRUE(operation.done());

        ExecuteResponse executeResponse;
        EXPECT_TRUE(operation.response().UnpackTo(&executeResponse));

        EXPECT_EQ(executeResponse.status().code(), google::rpc::Code::OK);

        const ActionResult actionResult = executeResponse.result();
        EXPECT_EQ(actionResult.exit_code(), 0);

        return actionResult.execution_metadata();
    }

    void expectOperationExitCode(const Operation &operation,
                                 google::rpc::Code code)
    {
        EXPECT_TRUE(operation.done());

        ExecuteResponse executeResponse;
        ASSERT_TRUE(operation.response().UnpackTo(&executeResponse));

        EXPECT_EQ(executeResponse.status().code(), code);
    }

    bool isRemoteOperation(const Operation &operation)
    {
        const std::string remotePrefix =
            proxyInstanceName + "#" + remoteInstanceName + "#";
        return operation.name().substr(0, remotePrefix.size()) == remotePrefix;
    }
};

#endif // INCLUDED_BUILDBOXCASD_HYBRIDFALLBACKEXECUTIONSERVICEFIXTURE_H
