/*
 * Copyright 2023 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_executionproxyinstance.h>
#include <buildboxcasd_fslocalactionstorage.h>
#include <buildboxcasd_fslocalassetstorage.h>
#include <buildboxcasd_server.h>
#include <buildboxcommon_fslocalcas.h>

#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommon_grpctestserver.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_remoteexecutionclient.h>

#include <build/bazel/remote/execution/v2/remote_execution_mock.grpc.pb.h>
#include <google/longrunning/operations_mock.grpc.pb.h>

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

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

class MockWriter final : public grpc::ServerWriterInterface<Operation> {
  public:
    MockWriter() {};

    bool Write(const Operation &msg, grpc::WriteOptions options) override
    {
        write_calls += 1;
        operations.push_back(msg);
        return true;
    };

    void SendInitialMetadata() override {};

    int write_calls = 0;
    std::vector<Operation> operations;
};

/***
 * Fixture to support testing buildbox-casd as a proxy for Execution
 *
 * This fixture sets up a test gRPC server whose behaviour is defined via
 * a handler thread in each test implementation.
 *
 * It also provides a `setUpProxy` method which configures a buildbox-casd
 * server to function as a proxy for Execution and Operations service
 * requests. The remote server is set as the test gRPC server, so tests
 * can verify that the request was proxied through as intended.
 */
class ExecutionServiceTestFixture : public ::testing::Test {
  protected:
    GrpcTestServer testServer;

    buildboxcommon::TemporaryDirectory socket_directory;
    const std::string TEST_PROXY_SERVER_ADDRESS;

    buildboxcommon::TemporaryDirectory proxy_storage_root_directory;
    std::shared_ptr<FsLocalCas> proxy_cas_storage;
    std::shared_ptr<FsLocalAssetStorage> proxy_asset_storage;
    std::shared_ptr<FsLocalActionStorage> proxy_action_storage;
    std::shared_ptr<Server> proxy_server;

    std::shared_ptr<grpc::Channel> channel;
    std::shared_ptr<Execution::StubInterface> execution_stub;
    std::shared_ptr<MockActionCacheStub> action_cache_stub;
    std::shared_ptr<google::longrunning::Operations::StubInterface>
        operations_stub;
    std::shared_ptr<GrpcClient> grpc_client;
    std::shared_ptr<RemoteExecutionClient> client;

    const std::string instance_name = "test";
    const std::vector<std::string> instance_names = {instance_name};
    const std::unordered_map<std::string, std::string> instanceMappings;
    const std::string remote_instance_name = "exec-dev";

    GetCapabilitiesRequest expectedGetCapabilitiesRequest;
    ServerCapabilities capabilities;
    Digest actionDigest;
    ExecuteRequest expectedExecuteRequest;
    ExecuteRequest executeRequest;
    ExecuteResponse executeResponse;
    WaitExecutionRequest expectedWaitExecutionRequest;
    WaitExecutionRequest waitExecutionRequest;
    GetOperationRequest expectedGetOperationRequest;
    GetOperationRequest getOperationRequest;
    CancelOperationRequest expectedCancelOperationRequest;
    CancelOperationRequest cancelOperationRequest;
    ListOperationsRequest expectedListOperationsRequest;
    ListOperationsRequest listOperationsRequest;

    Digest stdErrDigest;

    Operation operation;
    Operation expected_operation;
    ListOperationsResponse expectedListOperationsResponse;
    ListOperationsResponse listOperationsResponse;

    ExecutionServiceTestFixture()
        : TEST_PROXY_SERVER_ADDRESS("unix://" +
                                    std::string(socket_directory.name()) +
                                    "/proxy.sock"),
          grpc_client(std::make_shared<buildboxcommon::GrpcClient>()),
          client(std::make_shared<RemoteExecutionClient>(grpc_client, nullptr))
    {

        proxy_cas_storage =
            std::make_shared<FsLocalCas>(proxy_storage_root_directory.name());
        proxy_asset_storage = std::make_shared<FsLocalAssetStorage>(
            proxy_storage_root_directory.name());
        proxy_action_storage = std::make_shared<FsLocalActionStorage>(
            proxy_storage_root_directory.name());

        ConnectionOptions exec_endpoint;
        exec_endpoint.setUrl(TEST_PROXY_SERVER_ADDRESS);
        exec_endpoint.setInstanceName(instance_name);

        // Construct the expected requests we expect the proxy to send.
        actionDigest.set_hash("Action digest hash here");
        *expectedExecuteRequest.mutable_action_digest() = actionDigest;
        expectedExecuteRequest.set_instance_name(remote_instance_name);
        executeRequest.CopyFrom(expectedExecuteRequest);
        executeRequest.set_instance_name(instance_name);

        // Populate the example Operation message we'll respond with
        operation.set_name(remote_instance_name + "/operation-name");
        operation.set_done(true);
        operation.mutable_response()->PackFrom(executeResponse);

        expected_operation.CopyFrom(operation);
        expected_operation.set_name(instance_name + "#" + operation.name());

        expectedGetOperationRequest.set_name(operation.name());
        getOperationRequest.set_name(instance_name + "#" + operation.name());

        expectedCancelOperationRequest.set_name(operation.name());
        cancelOperationRequest.set_name(instance_name + "#" +
                                        operation.name());

        expectedListOperationsRequest.set_name(remote_instance_name);
        listOperationsRequest.set_name(instance_name + "#" +
                                       expectedListOperationsRequest.name());

        Operation *ptr = listOperationsResponse.add_operations();
        *ptr = operation;
        Operation *expected_ptr =
            expectedListOperationsResponse.add_operations();
        *expected_ptr = expected_operation;

        *expectedWaitExecutionRequest.mutable_name() = operation.name();
        *waitExecutionRequest.mutable_name() = expected_operation.name();

        expectedGetCapabilitiesRequest.set_instance_name(remote_instance_name);
        CacheCapabilities *const cache_capabilities =
            capabilities.mutable_cache_capabilities();

        cache_capabilities->add_digest_functions(
            DigestGenerator::digestFunction());

        cache_capabilities->set_symlink_absolute_path_strategy(
            SymlinkAbsolutePathStrategy_Value_ALLOWED);

        cache_capabilities->set_max_batch_total_size_bytes(0);
    }

    // NOTE: This needs to be separate to the test setup, since the
    // GrpcTestServer needs a handler thread running to deal with the
    // GetCapabilities request made on Server construction.
    void setUpProxy()
    {
        ConnectionOptions connection_options;
        connection_options.setUrl(testServer.url());
        connection_options.setInstanceName(remote_instance_name);
        connection_options.setRetryDelay("10");
        connection_options.setRetryLimit("1");
        connection_options.setRequestTimeout("10");

        proxy_server = std::make_shared<Server>(
            proxy_cas_storage, proxy_asset_storage, proxy_action_storage,
            nullptr, nullptr, nullptr, connection_options,
            std::optional(connection_options),
            std::optional(connection_options),
            std::optional(connection_options), instance_names,
            instanceMappings);

        proxy_server->addListeningPort(TEST_PROXY_SERVER_ADDRESS);
        proxy_server->start();

        channel = grpc::CreateChannel(TEST_PROXY_SERVER_ADDRESS,
                                      grpc::InsecureChannelCredentials());
        execution_stub = Execution::NewStub(channel);
        operations_stub = Operations::NewStub(channel);
        grpc_client->setInstanceName(instance_name + "#" +
                                     remote_instance_name);
        client->init(execution_stub, nullptr, operations_stub);
    }

    void handleCapabilities()
    {
        GrpcTestServerContext capabilities_ctx(
            &testServer,
            "/build.bazel.remote.execution.v2.Capabilities/GetCapabilities");
        capabilities_ctx.read(expectedGetCapabilitiesRequest);
        capabilities_ctx.writeAndFinish(capabilities);
    }
};

TEST_F(ExecutionServiceTestFixture, ExecuteRequest)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/build.bazel.remote.execution.v2.Execution/Execute");
        ctx.read(expectedExecuteRequest);
        ctx.writeAndFinish(operation);
    });

    setUpProxy();

    // Send an Execute request and ensure that the operation is sent to the
    // writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status = client->proxyExecuteRequest(
            executeRequest, [&] { return false; }, &writer);

        EXPECT_TRUE(writer.write_calls == 1);
        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            writer.operations[0], expected_operation));
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, ExecuteRequestFailedPrecondition)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/build.bazel.remote.execution.v2.Execution/Execute");
        grpc::Status s = grpc::Status(grpc::StatusCode::FAILED_PRECONDITION,
                                      "Command not in storage");
        ctx.finish(s);
    });

    setUpProxy();

    // Send an Execute request and ensure that the operation is sent to the
    // writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status = client->proxyExecuteRequest(
            executeRequest, [&] { return false; }, &writer);
    }
    catch (GrpcError &e) {
        EXPECT_TRUE(e.status.error_code() ==
                    grpc::StatusCode::FAILED_PRECONDITION);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, WaitExecutionRequest)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer,
            "/build.bazel.remote.execution.v2.Execution/WaitExecution");
        ctx.read(expectedWaitExecutionRequest);
        ctx.writeAndFinish(operation);
    });

    setUpProxy();

    // Send a WaitExecution request and ensure that the operation is sent to
    // the writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status = client->proxyWaitExecutionRequest(
            waitExecutionRequest, [&] { return false; }, &writer);

        EXPECT_TRUE(writer.write_calls == 1);
        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            writer.operations[0], expected_operation));
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, WaitExecutionRequestInvalidArgument)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer,
            "/build.bazel.remote.execution.v2.Execution/WaitExecution");
        grpc::Status s = grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                      "Operation name does not exist");
        ctx.finish(s);
    });

    setUpProxy();

    // Send an WaitExecution request and ensure that the operation is sent to
    // the writer and the status is returned.
    try {
        grpc::ServerContext context;
        MockWriter writer = MockWriter();
        grpc::Status status = client->proxyWaitExecutionRequest(
            waitExecutionRequest, [&] { return false; }, &writer);
    }
    catch (GrpcError &e) {
        EXPECT_TRUE(e.status.error_code() ==
                    grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, GetOperationRequest)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/GetOperation");
        ctx.read(expectedGetOperationRequest);
        ctx.writeAndFinish(operation);
    });

    setUpProxy();

    // Send a GetOperation request and ensure that the operation correctly
    // received in response.
    try {
        Operation actual;
        grpc::Status status =
            client->proxyGetOperationRequest(getOperationRequest, &actual);

        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            actual, expected_operation));
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, GetOperationRequestMissingOperation)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/GetOperation");
        ctx.read(expectedGetOperationRequest);
        grpc::Status s = grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                      "Operation not found");
        ctx.finish(s);
    });

    setUpProxy();

    // Send a GetOperation request and ensure that a non-OK status is
    // returned correctly.
    try {
        Operation actual;
        grpc::Status status =
            client->proxyGetOperationRequest(getOperationRequest, &actual);
    }
    catch (GrpcError &e) {
        EXPECT_TRUE(e.status.error_code() ==
                    grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, CancelOperationRequest)
{
    bool cancelled = false;
    std::thread serverHandler([this, &cancelled]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/CancelOperation");
        ctx.read(expectedCancelOperationRequest);
        cancelled = true;
        ctx.writeAndFinish(google::protobuf::Empty());
    });

    setUpProxy();

    // Send a CancelOperation request and ensure that the server handler was
    // called as expected.
    try {
        google::protobuf::Empty empty;
        grpc::Status status = client->proxyCancelOperationRequest(
            cancelOperationRequest, &empty);

        EXPECT_TRUE(status.ok());
        EXPECT_TRUE(cancelled);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture,
       CancelOperationRequestProxyMissingOperation)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/CancelOperation");
        ctx.read(expectedCancelOperationRequest);
        grpc::Status s = grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
                                      "Operation not found");
        ctx.finish(s);
    });

    setUpProxy();

    // Send a CancelOperation request and ensure that the server handler was
    // called as expected.
    try {
        google::protobuf::Empty empty;
        grpc::Status status = client->proxyCancelOperationRequest(
            cancelOperationRequest, &empty);
    }
    catch (GrpcError &e) {
        EXPECT_TRUE(e.status.error_code() ==
                    grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, ListOperationsRequestProxy)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/ListOperations");
        ctx.read(expectedListOperationsRequest);
        ctx.writeAndFinish(listOperationsResponse);
    });

    setUpProxy();

    // Send a ListOperations request and ensure that the server handler was
    // called as expected.
    try {
        ListOperationsResponse response;
        grpc::Status status = client->proxyListOperationsRequest(
            listOperationsRequest, &response);

        EXPECT_TRUE(google::protobuf::util::MessageDifferencer::Equals(
            response, expectedListOperationsResponse));
        EXPECT_TRUE(status.ok());
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}

TEST_F(ExecutionServiceTestFixture, ListOperationsRequestProxyBadQuery)
{
    std::thread serverHandler([this]() {
        handleCapabilities();

        GrpcTestServerContext ctx(
            &testServer, "/google.longrunning.Operations/ListOperations");
        ctx.read(expectedListOperationsRequest);
        grpc::Status s =
            grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "Bad query");
        ctx.finish(s);
    });

    setUpProxy();

    // Send a ListOperations request and ensure that the non-OK status code
    // is propagated to the client correctly
    try {
        ListOperationsResponse response;
        grpc::Status status = client->proxyListOperationsRequest(
            listOperationsRequest, &response);
    }
    catch (GrpcError &e) {
        EXPECT_TRUE(e.status.error_code() ==
                    grpc::StatusCode::INVALID_ARGUMENT);
    }
    catch (...) {
        serverHandler.join();
        throw;
    }
    serverHandler.join();
}
